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

4
backend/app/catalog.py
View File

@@ -84,9 +84,11 @@ TASK_TYPES = [
"yolo.train_custom",
"yolo.batch_train",
"yolo.predict",
"yolo.predict_custom",
"yolo.predict_v1",
"yolo.batch_predict",
"yolo.heatmap",
"yolo.heatmap_custom",
"yolo.compare",
"yolo.raw_mask_check",
"yolo.copy_best",
@@ -185,9 +187,11 @@ TASK_DEFAULTS: dict[str, dict[str, Any]] = {
"yolo.train_custom": {"data": "var/uploads/datasets/example/dataset.yaml", "model": "YOLO11n-seg", "epochs": 10, "imgsz": 640, "batch": 1, "workers": 0, "device": "cpu", "exist_ok": True},
"yolo.batch_train": {},
"yolo.predict": {"model": "YOLOv8n-seg", "pt_name": "best.pt", "conf": 0.2, "run_choice": 1},
"yolo.predict_custom": {"weights": "var/custom_yolo_runs/example/weights/best.pt", "source": "var/uploads/datasets/example/images", "imgsz": 640, "conf": 0.25, "device": "cpu", "name": "example_predict", "exist_ok": True},
"yolo.predict_v1": {"model": "YOLOv8n-seg", "pt_name": "best.pt", "conf": 0.2, "run_choice": 1},
"yolo.batch_predict": {"pt_name": "best.pt", "conf": 0.2},
"yolo.heatmap": {"model": "YOLOv8n-seg", "cam_method": "All", "pt_name": "best.pt", "run_choice": 1},
"yolo.heatmap_custom": {"weights": "var/custom_yolo_runs/example/weights/best.pt", "source": "var/uploads/datasets/example/images", "model_key": "YOLO11n-seg", "cam_method": "GradCAM", "target_layers": "model.model.model[9]", "limit": 3, "name": "example_heatmap"},
"yolo.compare": {"pt_name": "all"},
"yolo.raw_mask_check": {"pt_name": "best.pt"},
"yolo.copy_best": {"pt_name": "best.pt"},