81 lines
4.1 KiB
Python
81 lines
4.1 KiB
Python
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.predict_v1":
|
|
args = conda_python(conda_env, YOLO_DIR / "yolo_predict_V1.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 legacy YOLO V1 script", 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")
|
|
|
|
if job_type == "yolo.layer_tester":
|
|
return CommandSpec(conda_python(conda_env, YOLO_DIR / "Yolo可视化测试" / "yolo_layer_tester.py"), YOLO_DIR, "test YOLO heatmap target layers")
|
|
|
|
return None
|