Add dataset QA and custom YOLO training flow
This commit is contained in:
@@ -13,6 +13,7 @@ from ...config import settings
|
||||
|
||||
DATASET_KINDS = ("images", "labels", "masks")
|
||||
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"}
|
||||
LABEL_EXTS = {".txt", ".json", ".yaml", ".yml"}
|
||||
|
||||
|
||||
def uploads_root() -> Path:
|
||||
@@ -76,6 +77,62 @@ def _iter_files(root: Path) -> Iterable[Path]:
|
||||
return sorted(path for path in root.rglob("*") if path.is_file())
|
||||
|
||||
|
||||
def _stem_map(paths: Iterable[Path], exts: set[str] | None = None) -> dict[str, Path]:
|
||||
result: dict[str, Path] = {}
|
||||
for path in paths:
|
||||
if exts and path.suffix.lower() not in exts:
|
||||
continue
|
||||
result.setdefault(path.stem, path)
|
||||
return result
|
||||
|
||||
|
||||
def _image_shape(path: Path) -> dict | None:
|
||||
try:
|
||||
import cv2
|
||||
|
||||
image = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
|
||||
if image is None:
|
||||
return None
|
||||
height, width = image.shape[:2]
|
||||
channels = 1 if image.ndim == 2 else image.shape[2]
|
||||
return {"width": int(width), "height": int(height), "channels": int(channels)}
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _validate_yolo_txt(path: Path) -> dict:
|
||||
errors = []
|
||||
classes: set[int] = set()
|
||||
lines = path.read_text(encoding="utf-8", errors="replace").splitlines()
|
||||
annotation_count = 0
|
||||
for line_number, line in enumerate(lines, 1):
|
||||
raw = line.strip()
|
||||
if not raw:
|
||||
continue
|
||||
parts = raw.split()
|
||||
annotation_count += 1
|
||||
if len(parts) < 5:
|
||||
errors.append(f"{path.name}:{line_number} has fewer than 5 tokens")
|
||||
continue
|
||||
try:
|
||||
class_id = int(float(parts[0]))
|
||||
coords = [float(item) for item in parts[1:]]
|
||||
except ValueError:
|
||||
errors.append(f"{path.name}:{line_number} contains non-numeric values")
|
||||
continue
|
||||
if class_id < 0:
|
||||
errors.append(f"{path.name}:{line_number} has negative class id")
|
||||
classes.add(class_id)
|
||||
if len(coords) not in {4} and len(coords) < 6:
|
||||
errors.append(f"{path.name}:{line_number} has too few segmentation coordinates")
|
||||
if len(coords) != 4 and len(coords) % 2 != 0:
|
||||
errors.append(f"{path.name}:{line_number} has an odd number of polygon coordinates")
|
||||
out_of_range = [value for value in coords if value < 0 or value > 1]
|
||||
if out_of_range:
|
||||
errors.append(f"{path.name}:{line_number} has coordinates outside 0..1")
|
||||
return {"annotations": annotation_count, "classes": sorted(classes), "errors": errors}
|
||||
|
||||
|
||||
def describe_dataset(name: str) -> dict:
|
||||
safe_name = slugify(name)
|
||||
root = dataset_dir(safe_name)
|
||||
@@ -99,6 +156,113 @@ def describe_dataset(name: str) -> dict:
|
||||
return {**meta, "absolute_layout": absolute_layout, "counts": counts, "samples": samples}
|
||||
|
||||
|
||||
def validate_dataset(name: str) -> dict:
|
||||
safe_name = slugify(name)
|
||||
root = dataset_dir(safe_name)
|
||||
image_files = list(_iter_files(root / "images"))
|
||||
label_files = list(_iter_files(root / "labels"))
|
||||
mask_files = list(_iter_files(root / "masks"))
|
||||
images = _stem_map(image_files, IMAGE_EXTS)
|
||||
labels = _stem_map(label_files, LABEL_EXTS)
|
||||
masks = _stem_map(mask_files, IMAGE_EXTS)
|
||||
|
||||
image_stems = set(images)
|
||||
label_stems = set(labels)
|
||||
mask_stems = set(masks)
|
||||
paired_label_stems = sorted(image_stems & label_stems)
|
||||
paired_mask_stems = sorted(image_stems & mask_stems)
|
||||
|
||||
yolo_errors = []
|
||||
class_ids: set[int] = set()
|
||||
annotation_count = 0
|
||||
for label in labels.values():
|
||||
if label.suffix.lower() != ".txt":
|
||||
continue
|
||||
detail = _validate_yolo_txt(label)
|
||||
yolo_errors.extend(detail["errors"])
|
||||
class_ids.update(detail["classes"])
|
||||
annotation_count += detail["annotations"]
|
||||
|
||||
shape_mismatches = []
|
||||
sample_shapes = []
|
||||
for stem in paired_mask_stems[:80]:
|
||||
image_shape = _image_shape(images[stem])
|
||||
mask_shape = _image_shape(masks[stem])
|
||||
if image_shape and len(sample_shapes) < 8:
|
||||
sample_shapes.append({"name": images[stem].name, **image_shape})
|
||||
if image_shape and mask_shape and (image_shape["width"], image_shape["height"]) != (mask_shape["width"], mask_shape["height"]):
|
||||
shape_mismatches.append({"stem": stem, "image": image_shape, "mask": mask_shape})
|
||||
if not sample_shapes:
|
||||
for image in list(images.values())[:8]:
|
||||
image_shape = _image_shape(image)
|
||||
if image_shape:
|
||||
sample_shapes.append({"name": image.name, **image_shape})
|
||||
|
||||
checks = [
|
||||
{"name": "has_images", "passed": len(images) > 0, "count": len(images)},
|
||||
{"name": "has_labels_or_masks", "passed": len(labels) > 0 or len(masks) > 0, "labels": len(labels), "masks": len(masks)},
|
||||
{"name": "image_label_pairs", "passed": len(label_stems) == 0 or len(paired_label_stems) > 0, "count": len(paired_label_stems)},
|
||||
{"name": "image_mask_pairs", "passed": len(mask_stems) == 0 or len(paired_mask_stems) > 0, "count": len(paired_mask_stems)},
|
||||
{"name": "yolo_txt_valid", "passed": not yolo_errors, "errors": yolo_errors[:20]},
|
||||
{"name": "mask_shapes_match", "passed": not shape_mismatches, "errors": shape_mismatches[:20]},
|
||||
]
|
||||
yolo_ready = len(images) > 0 and len(paired_label_stems) > 0 and not yolo_errors
|
||||
mask_ready = len(images) > 0 and len(paired_mask_stems) > 0 and not shape_mismatches
|
||||
return {
|
||||
"dataset": safe_name,
|
||||
"root": str(root.resolve()),
|
||||
"counts": {"images": len(images), "labels": len(labels), "masks": len(masks), "annotations": annotation_count},
|
||||
"pairs": {
|
||||
"image_label": len(paired_label_stems),
|
||||
"image_mask": len(paired_mask_stems),
|
||||
"images_without_labels": sorted(image_stems - label_stems)[:50],
|
||||
"labels_without_images": sorted(label_stems - image_stems)[:50],
|
||||
"images_without_masks": sorted(image_stems - mask_stems)[:50],
|
||||
"masks_without_images": sorted(mask_stems - image_stems)[:50],
|
||||
},
|
||||
"classes": sorted(class_ids),
|
||||
"sample_shapes": sample_shapes,
|
||||
"checks": checks,
|
||||
"ready": {"yolo": yolo_ready, "mask": mask_ready, "any": yolo_ready or mask_ready},
|
||||
}
|
||||
|
||||
|
||||
def generate_yolo_dataset_yaml(name: str, class_names: list[str] | None = None) -> dict:
|
||||
validation = validate_dataset(name)
|
||||
if not validation["ready"]["yolo"]:
|
||||
raise ValueError("dataset is not YOLO-ready; upload matching images and valid txt labels first")
|
||||
safe_name = slugify(name)
|
||||
root = dataset_dir(safe_name)
|
||||
classes = validation["classes"] or [0]
|
||||
class_count = max(classes) + 1
|
||||
names = list(class_names or [])
|
||||
if len(names) < class_count:
|
||||
names.extend(f"class_{index}" for index in range(len(names), class_count))
|
||||
yaml_path = root / "dataset.yaml"
|
||||
names_block = "\n".join(f" {index}: {label}" for index, label in enumerate(names))
|
||||
yaml_text = "\n".join(
|
||||
[
|
||||
f"path: {root.resolve()}",
|
||||
"train: images",
|
||||
"val: images",
|
||||
f"nc: {len(names)}",
|
||||
"names:",
|
||||
names_block,
|
||||
"",
|
||||
]
|
||||
)
|
||||
yaml_path.write_text(yaml_text, encoding="utf-8")
|
||||
return {
|
||||
"dataset": safe_name,
|
||||
"path": str(yaml_path.resolve()),
|
||||
"relative_path": str(yaml_path.resolve().relative_to(settings.project_root)),
|
||||
"classes": classes,
|
||||
"names": names,
|
||||
"content": yaml_text,
|
||||
"validation": validation,
|
||||
}
|
||||
|
||||
|
||||
def list_uploaded_datasets() -> list[dict]:
|
||||
root = uploads_root()
|
||||
datasets = []
|
||||
|
||||
Reference in New Issue
Block a user