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()