Add dataset QA and custom YOLO training flow

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)