Add custom YOLO prediction and heatmap workflow

backend/app/modules/yolo/custom_train.py

@@ -1,10 +1,25 @@
 from __future__ import annotations
 
 import argparse
+import os
 from pathlib import Path
 
 from ultralytics import YOLO
 
+DEFAULT_PROJECT_ROOT = Path(__file__).resolve().parents[4]
+
+
+def project_root() -> Path:
+    raw = os.getenv("SEG_DATA_SERVER_ROOT")
+    if not raw:
+        return DEFAULT_PROJECT_ROOT
+    path = Path(raw).expanduser()
+    if path.is_absolute():
+        return path.resolve()
+    return (DEFAULT_PROJECT_ROOT / path).resolve()
+
+
+PROJECT_ROOT = project_root()
 
 MODEL_ALIASES = {
     "YOLOv8n-seg": "yolov8n-seg.pt",
@@ -22,6 +37,13 @@ MODEL_ALIASES = {
 }
 
 
+def resolve_project_path(value: str) -> Path:
+    path = Path(value).expanduser()
+    if path.is_absolute():
+        return path.resolve()
+    return (PROJECT_ROOT / path).resolve()
+
+
 def resolve_model(value: str) -> str:
     candidate = MODEL_ALIASES.get(value, value)
     path = Path(candidate).expanduser()
@@ -46,13 +68,13 @@ def main() -> None:
 
     model = YOLO(resolve_model(args.model))
     result = model.train(
-        data=str(Path(args.data).expanduser().resolve()),
+        data=str(resolve_project_path(args.data)),
         epochs=args.epochs,
         imgsz=args.imgsz,
         batch=args.batch,
         workers=args.workers,
         device=args.device,
-        project=str(Path(args.project).expanduser().resolve()),
+        project=str(resolve_project_path(args.project)),
         name=args.name,
         exist_ok=args.exist_ok,
         verbose=True,