60 lines
3.0 KiB
Python
60 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
from ...commands import CommandSpec, append_flag, conda_python
|
|
from ...config import settings
|
|
from ..dataset.tasks import _append_label_rebuild_flags
|
|
|
|
|
|
VISUAL_DIR = settings.source_root / "Tool-可视化"
|
|
VISUAL_LABEL_DIR = VISUAL_DIR / "0_图片Labels生成"
|
|
|
|
|
|
def build_visual_task(job_type: str, params: dict, conda_env: str) -> CommandSpec | None:
|
|
if job_type == "visual.train":
|
|
return CommandSpec(conda_python(conda_env, VISUAL_DIR / "train.py"), VISUAL_DIR, "run standalone YOLO visual training script")
|
|
|
|
if job_type == "visual.inference":
|
|
return CommandSpec(conda_python(conda_env, VISUAL_DIR / "inference.py"), VISUAL_DIR, "run standalone YOLO visual inference script")
|
|
|
|
if job_type == "visual.fps":
|
|
args = conda_python(conda_env, VISUAL_DIR / "get_FPS.py")
|
|
append_flag(args, "--weights", params.get("weights", "yolov8n.pt"))
|
|
append_flag(args, "--batch", params.get("batch", 1))
|
|
imgs = params.get("imgs", [640, 640])
|
|
if isinstance(imgs, str):
|
|
imgs = [part for part in imgs.replace(",", " ").split() if part]
|
|
if imgs:
|
|
args.append("--imgs")
|
|
args.extend(str(item) for item in imgs)
|
|
append_flag(args, "--device", params.get("device", ""))
|
|
append_flag(args, "--warmup", params.get("warmup"))
|
|
append_flag(args, "--testtime", params.get("testtime"))
|
|
if params.get("half"):
|
|
args.append("--half")
|
|
return CommandSpec(args, VISUAL_DIR, "measure YOLO visual model FPS")
|
|
|
|
if job_type == "visual.yolo11_heatmap_v1":
|
|
return CommandSpec(conda_python(conda_env, VISUAL_DIR / "yolov11_heatmap_V1.py"), VISUAL_DIR, "run standalone YOLOv11 heatmap V1")
|
|
|
|
if job_type == "visual.yolo11_heatmap_v2":
|
|
return CommandSpec(conda_python(conda_env, VISUAL_DIR / "yolov11_heatmap_V2.py"), VISUAL_DIR, "run standalone YOLOv11 heatmap V2")
|
|
|
|
if job_type == "visual.label_ori":
|
|
return CommandSpec(conda_python(conda_env, VISUAL_DIR / "Tool_Check_and_Gen_Txt_Label_ori_label.py"), VISUAL_DIR, "generate YOLO txt labels from original label ids")
|
|
|
|
if job_type == "visual.label_sort":
|
|
return CommandSpec(conda_python(conda_env, VISUAL_DIR / "Tool_Check_and_Gen_Txt_Label_sort_label.py"), VISUAL_DIR, "generate YOLO txt labels from sorted label ids")
|
|
|
|
if job_type == "visual.gen_8bit_png":
|
|
return CommandSpec(conda_python(conda_env, VISUAL_DIR / "Tool_Gen_8_Bit_PNG[没用，不认].py"), VISUAL_DIR, "generate 8-bit PNG labels with legacy visual tool")
|
|
|
|
if job_type == "visual.deal_labels":
|
|
args = conda_python(conda_env, VISUAL_LABEL_DIR / "4_deal_labels.py")
|
|
_append_label_rebuild_flags(args, params)
|
|
return CommandSpec(args, VISUAL_LABEL_DIR, "run visual label rebuild/conversion script")
|
|
|
|
if job_type == "visual.tool_deal_labels_demo":
|
|
return CommandSpec(conda_python(conda_env, VISUAL_LABEL_DIR / "Tool_deal_labels.py"), VISUAL_LABEL_DIR, "run visual label helper demo script")
|
|
|
|
return None
|