Initial Seg Data Server Net platform

backend/app/modules/yolo/tasks.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from ...commands import CommandSpec, append_flag, bash, conda_python, required
+from ...config import settings
+
+
+YOLO_DIR = settings.source_root / "Seg_All_In_One_YoloModel"
+VIDEO_YOLO_DIR = settings.source_root / "Seg_Predict_YoloModel"
+
+
+def build_yolo_task(job_type: str, params: dict, conda_env: str) -> CommandSpec | None:
+    env = {"SEG_CONDA_ENV": conda_env}
+
+    if job_type == "yolo.train":
+        args = conda_python(conda_env, YOLO_DIR / "yolo_train.py")
+        append_flag(args, "--model", required(params, "model"))
+        return CommandSpec(args, YOLO_DIR, "train one Ultralytics YOLO segmentation model")
+
+    if job_type == "yolo.batch_train":
+        return CommandSpec(bash(YOLO_DIR / "yolo_train.sh"), YOLO_DIR, "run legacy YOLO batch training", env=env)
+
+    if job_type == "yolo.predict":
+        args = conda_python(conda_env, YOLO_DIR / "yolo_predict_V2.py")
+        append_flag(args, "--model", required(params, "model"))
+        append_flag(args, "--source", params.get("source"))
+        append_flag(args, "--pt_name", params.get("pt_name", "best.pt"))
+        append_flag(args, "--conf", params.get("conf", 0.2))
+        choice = str(params.get("run_choice", 1))
+        return CommandSpec(args, YOLO_DIR, "predict with one YOLO model", stdin_text=f"{choice}\n")
+
+    if job_type == "yolo.batch_predict":
+        args = bash(YOLO_DIR / "yolo_predict.sh")
+        append_flag(args, "--pt_name", params.get("pt_name", "best.pt"))
+        append_flag(args, "--conf", params.get("conf", 0.2))
+        append_flag(args, "--heatmap_method", params.get("heatmap_method"))
+        return CommandSpec(args, YOLO_DIR, "run legacy YOLO batch prediction", env=env)
+
+    if job_type == "yolo.heatmap":
+        args = conda_python(conda_env, YOLO_DIR / "yolo_predict_visualize_nn.py")
+        append_flag(args, "--model", required(params, "model"))
+        append_flag(args, "--target_layers", params.get("target_layers", "default"))
+        append_flag(args, "--cam_method", params.get("cam_method", "All"))
+        append_flag(args, "--pt_name", params.get("pt_name", "best.pt"))
+        choice = str(params.get("run_choice", 1))
+        return CommandSpec(args, YOLO_DIR, "generate YOLO heatmaps", stdin_text=f"{choice}\n")
+
+    if job_type == "yolo.compare":
+        args = conda_python(conda_env, YOLO_DIR / "yolo_predict_V2_compare_all.py")
+        append_flag(args, "--pt_name", params.get("pt_name", "all"))
+        return CommandSpec(args, YOLO_DIR, "compare all YOLO prediction outputs")
+
+    if job_type == "yolo.raw_mask_check":
+        args = conda_python(conda_env, YOLO_DIR / "yolo_predict_raw_masks_check.py")
+        append_flag(args, "--pt_name", params.get("pt_name", "best.pt"))
+        return CommandSpec(args, YOLO_DIR, "check YOLO raw mask completeness")
+
+    if job_type == "yolo.copy_best":
+        args = bash(YOLO_DIR / "Tool_Yolo_Copy_Best_Model.sh")
+        append_flag(args, "--pt_name", params.get("pt_name", "best.pt"))
+        return CommandSpec(args, YOLO_DIR, "copy YOLO best weights into prediction area")
+
+    if job_type == "yolo.video_visible":
+        return CommandSpec(conda_python(conda_env, VIDEO_YOLO_DIR / "yolo_Seg_Video-V1-Visible.py"), VIDEO_YOLO_DIR, "render visible YOLO video prediction")
+
+    if job_type == "yolo.video_unvisible":
+        return CommandSpec(conda_python(conda_env, VIDEO_YOLO_DIR / "yolo_Seg_Video-V2-UnVisible.py"), VIDEO_YOLO_DIR, "render invisible/headless YOLO video prediction")
+
+    return None
+