Add dataset QA and custom YOLO training flow

backend/app/modules/dataset/service.py

@@ -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 = []

backend/app/modules/yolo/custom_train.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from ultralytics import YOLO
+
+
+MODEL_ALIASES = {
+    "YOLOv8n-seg": "yolov8n-seg.pt",
+    "YOLOv8s-seg": "yolov8s-seg.pt",
+    "YOLOv8m-seg": "yolov8m-seg.pt",
+    "YOLOv8l-seg": "yolov8l-seg.pt",
+    "YOLOv8x-seg": "yolov8x-seg.pt",
+    "YOLOv9c-seg": "yolov9c-seg.pt",
+    "YOLOv9e-seg": "yolov9e-seg.pt",
+    "YOLO11n-seg": "yolo11n-seg.pt",
+    "YOLO11s-seg": "yolo11s-seg.pt",
+    "YOLO11m-seg": "yolo11m-seg.pt",
+    "YOLO11l-seg": "yolo11l-seg.pt",
+    "YOLO11x-seg": "yolo11x-seg.pt",
+}
+
+
+def resolve_model(value: str) -> str:
+    candidate = MODEL_ALIASES.get(value, value)
+    path = Path(candidate).expanduser()
+    if path.is_absolute() or path.exists():
+        return str(path.resolve())
+    return candidate
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Train a YOLO segmentation model from a supplied dataset.yaml.")
+    parser.add_argument("--data", required=True, help="Path to YOLO dataset.yaml.")
+    parser.add_argument("--model", default="YOLO11n-seg", help="Model alias, weight path, or Ultralytics model name.")
+    parser.add_argument("--epochs", type=int, default=10)
+    parser.add_argument("--imgsz", type=int, default=640)
+    parser.add_argument("--batch", type=int, default=1)
+    parser.add_argument("--workers", type=int, default=0)
+    parser.add_argument("--device", default="cpu")
+    parser.add_argument("--project", required=True)
+    parser.add_argument("--name", default="custom_upload")
+    parser.add_argument("--exist-ok", action="store_true")
+    args = parser.parse_args()
+
+    model = YOLO(resolve_model(args.model))
+    result = model.train(
+        data=str(Path(args.data).expanduser().resolve()),
+        epochs=args.epochs,
+        imgsz=args.imgsz,
+        batch=args.batch,
+        workers=args.workers,
+        device=args.device,
+        project=str(Path(args.project).expanduser().resolve()),
+        name=args.name,
+        exist_ok=args.exist_ok,
+        verbose=True,
+    )
+    save_dir = getattr(result, "save_dir", None) or getattr(model.trainer, "save_dir", "")
+    print(f"save_dir={save_dir}")
+
+
+if __name__ == "__main__":
+    main()

backend/app/modules/yolo/tasks.py

@@ -6,6 +6,7 @@ from ...config import settings
 
 YOLO_DIR = settings.source_root / "Seg_All_In_One_YoloModel"
 VIDEO_YOLO_DIR = settings.source_root / "Seg_Predict_YoloModel"
+CUSTOM_TRAIN = settings.project_root / "backend" / "app" / "modules" / "yolo" / "custom_train.py"
 
 
 def build_yolo_task(job_type: str, params: dict, conda_env: str) -> CommandSpec | None:
@@ -16,6 +17,21 @@ def build_yolo_task(job_type: str, params: dict, conda_env: str) -> CommandSpec
         append_flag(args, "--model", required(params, "model"))
         return CommandSpec(args, YOLO_DIR, "train one Ultralytics YOLO segmentation model")
 
+    if job_type == "yolo.train_custom":
+        args = conda_python(conda_env, CUSTOM_TRAIN)
+        append_flag(args, "--data", required(params, "data"))
+        append_flag(args, "--model", params.get("model", "YOLO11n-seg"))
+        append_flag(args, "--epochs", params.get("epochs", 10))
+        append_flag(args, "--imgsz", params.get("imgsz", 640))
+        append_flag(args, "--batch", params.get("batch", 1))
+        append_flag(args, "--workers", params.get("workers", 0))
+        append_flag(args, "--device", params.get("device", "cpu"))
+        append_flag(args, "--project", params.get("project", settings.project_root / "var" / "custom_yolo_runs"))
+        append_flag(args, "--name", params.get("name", "custom_upload"))
+        if params.get("exist_ok", True):
+            args.append("--exist-ok")
+        return CommandSpec(args, YOLO_DIR, "train YOLO segmentation on a supplied dataset.yaml")
+
     if job_type == "yolo.batch_train":
         return CommandSpec(bash(YOLO_DIR / "yolo_train.sh"), YOLO_DIR, "run legacy YOLO batch training", env=env)