admin/Seg_Data_Server_Net
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add custom YOLO prediction and heatmap workflow

Browse Source
This commit is contained in:
admin
2026-06-30 15:11:47 +08:00
parent 4d0c26be05
commit 777f168a75
12 changed files with 393 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
backend/app/modules/yolo/tasks.py
View File

@@ -7,6 +7,8 @@ from ...config import settings
YOLO_DIR = settings.source_root / "Seg_All_In_One_YoloModel"
VIDEO_YOLO_DIR = settings.source_root / "Seg_Predict_YoloModel"
CUSTOM_TRAIN = settings.project_root / "backend" / "app" / "modules" / "yolo" / "custom_train.py"
CUSTOM_PREDICT = settings.project_root / "backend" / "app" / "modules" / "yolo" / "custom_predict.py"
CUSTOM_HEATMAP = settings.project_root / "backend" / "app" / "modules" / "yolo" / "custom_heatmap.py"
def build_yolo_task(job_type: str, params: dict, conda_env: str) -> CommandSpec | None:
@@ -32,6 +34,36 @@ def build_yolo_task(job_type: str, params: dict, conda_env: str) -> CommandSpec
args.append("--exist-ok")
return CommandSpec(args, YOLO_DIR, "train YOLO segmentation on a supplied dataset.yaml")
if job_type == "yolo.predict_custom":
args = conda_python(conda_env, CUSTOM_PREDICT)
append_flag(args, "--weights", required(params, "weights"))
append_flag(args, "--source", required(params, "source"))
append_flag(args, "--project", params.get("project", settings.project_root / "var" / "custom_yolo_runs"))
append_flag(args, "--name", params.get("name", "custom_predict"))
append_flag(args, "--imgsz", params.get("imgsz", 640))
append_flag(args, "--conf", params.get("conf", 0.25))
append_flag(args, "--device", params.get("device", "cpu"))
if params.get("save_txt", False):
args.append("--save-txt")
if params.get("save_conf", False):
args.append("--save-conf")
if params.get("exist_ok", True):
args.append("--exist-ok")
return CommandSpec(args, YOLO_DIR, "predict uploaded/custom images with a YOLO checkpoint")
if job_type == "yolo.heatmap_custom":
args = conda_python(conda_env, CUSTOM_HEATMAP)
append_flag(args, "--weights", required(params, "weights"))
append_flag(args, "--source", required(params, "source"))
append_flag(args, "--project", params.get("project", settings.project_root / "var" / "custom_yolo_runs"))
append_flag(args, "--name", params.get("name", "custom_heatmap"))
append_flag(args, "--model-key", params.get("model_key", "YOLO11n-seg"))
append_flag(args, "--pt-name", params.get("pt_name", "best.pt"))
append_flag(args, "--cam-method", params.get("cam_method", "GradCAM"))
append_flag(args, "--target-layers", params.get("target_layers", "model.model.model[9]"))
append_flag(args, "--limit", params.get("limit", 3))
return CommandSpec(args, YOLO_DIR, "generate YOLO heatmaps for uploaded/custom images")
if job_type == "yolo.batch_train":
return CommandSpec(bash(YOLO_DIR / "yolo_train.sh"), YOLO_DIR, "run legacy YOLO batch training", env=env)