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

Add dataset QA and custom YOLO training flow

Browse Source
This commit is contained in:
admin
2026-06-30 14:04:11 +08:00
parent 43ed767b4f
commit 93af8bcd3a
14 changed files with 529 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -6,6 +6,7 @@ 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"
def build_yolo_task(job_type: str, params: dict, conda_env: str) -> CommandSpec | None:
@@ -16,6 +17,21 @@ def build_yolo_task(job_type: str, params: dict, conda_env: str) -> CommandSpec
append_flag(args, "--model", required(params, "model"))
return CommandSpec(args, YOLO_DIR, "train one Ultralytics YOLO segmentation model")
if job_type == "yolo.train_custom":
args = conda_python(conda_env, CUSTOM_TRAIN)
append_flag(args, "--data", required(params, "data"))
append_flag(args, "--model", params.get("model", "YOLO11n-seg"))
append_flag(args, "--epochs", params.get("epochs", 10))
append_flag(args, "--imgsz", params.get("imgsz", 640))
append_flag(args, "--batch", params.get("batch", 1))
append_flag(args, "--workers", params.get("workers", 0))
append_flag(args, "--device", params.get("device", "cpu"))
append_flag(args, "--project", params.get("project", settings.project_root / "var" / "custom_yolo_runs"))
append_flag(args, "--name", params.get("name", "custom_upload"))
if params.get("exist_ok", True):
args.append("--exist-ok")
return CommandSpec(args, YOLO_DIR, "train YOLO segmentation on a supplied dataset.yaml")
if job_type == "yolo.batch_train":
return CommandSpec(bash(YOLO_DIR / "yolo_train.sh"), YOLO_DIR, "run legacy YOLO batch training", env=env)