Add custom YOLO prediction and heatmap workflow

backend/app/modules/yolo/custom_predict.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+import argparse
+import json
+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()
+
+
+def resolve_project_path(value: str | Path) -> Path:
+    path = Path(value).expanduser()
+    if path.is_absolute():
+        return path.resolve()
+    return (PROJECT_ROOT / path).resolve()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Predict segmentation masks with a supplied YOLO checkpoint.")
+    parser.add_argument("--weights", required=True, help="Path to best.pt/last.pt or another YOLO segmentation checkpoint.")
+    parser.add_argument("--source", required=True, help="Image file or image directory to predict.")
+    parser.add_argument("--project", default=str(PROJECT_ROOT / "var" / "custom_yolo_runs"))
+    parser.add_argument("--name", default="custom_predict")
+    parser.add_argument("--imgsz", type=int, default=640)
+    parser.add_argument("--conf", type=float, default=0.25)
+    parser.add_argument("--device", default="cpu")
+    parser.add_argument("--save-txt", action="store_true")
+    parser.add_argument("--save-conf", action="store_true")
+    parser.add_argument("--exist-ok", action="store_true")
+    args = parser.parse_args()
+
+    weights = resolve_project_path(args.weights)
+    source = resolve_project_path(args.source)
+    project = resolve_project_path(args.project)
+    if not weights.exists():
+        raise FileNotFoundError(f"weights not found: {weights}")
+    if not source.exists():
+        raise FileNotFoundError(f"source not found: {source}")
+
+    model = YOLO(str(weights))
+    results = model.predict(
+        source=str(source),
+        imgsz=args.imgsz,
+        conf=args.conf,
+        device=args.device,
+        project=str(project),
+        name=args.name,
+        save=True,
+        save_txt=args.save_txt,
+        save_conf=args.save_conf,
+        exist_ok=args.exist_ok,
+        verbose=True,
+    )
+    save_dir = Path(getattr(model.predictor, "save_dir", project / args.name))
+    metadata = {
+        "weights": str(weights),
+        "source": str(source),
+        "save_dir": str(save_dir),
+        "count": len(results),
+    }
+    save_dir.mkdir(parents=True, exist_ok=True)
+    (save_dir / "predict_manifest.json").write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8")
+    print(json.dumps(metadata, ensure_ascii=False))
+
+
+if __name__ == "__main__":
+    main()