From 777f168a75888ac943da411fce1d0fc247c39f4e Mon Sep 17 00:00:00 2001 From: admin <572701190@qq.com> Date: Tue, 30 Jun 2026 15:11:47 +0800 Subject: [PATCH] Add custom YOLO prediction and heatmap workflow --- README.md | 12 +- backend/app/agents/evaluation_agent.py | 6 +- backend/app/agents/validation_agent.py | 9 +- backend/app/capabilities.py | 2 + backend/app/catalog.py | 4 + backend/app/modules/results/service.py | 1 + backend/app/modules/yolo/custom_heatmap.py | 144 +++++++++++++++++++++ backend/app/modules/yolo/custom_predict.py | 82 ++++++++++++ backend/app/modules/yolo/custom_train.py | 26 +++- backend/app/modules/yolo/tasks.py | 32 +++++ backend/tests/test_catalog.py | 2 + frontend/src/main.tsx | 90 ++++++++++++- 12 files changed, 393 insertions(+), 17 deletions(-) create mode 100644 backend/app/modules/yolo/custom_heatmap.py create mode 100644 backend/app/modules/yolo/custom_predict.py diff --git a/README.md b/README.md index aa53ea6..2d86a93 100644 --- a/README.md +++ b/README.md @@ -51,9 +51,12 @@ resize, pair-check, label rebuild, transparent overlay, stitch, and video-frame jobs. Selecting an uploaded dataset fills task JSON with its images, labels, and masks directories. The dataset panel validates image/label/mask pairing, checks YOLO txt labels and mask dimensions, and can generate a `dataset.yaml` -for the `yolo.train_custom` task. Segmentation previews, YOLO heatmaps, and -loss/metric artifacts are grouped on the results dashboard, and YOLO-style -`results.csv` files are parsed into lightweight training curves. +for the `yolo.train_custom` task. The selected upload dataset also exposes +direct YOLO custom train, predict, and heatmap actions; custom outputs are +written under `var/custom_yolo_runs` and are scanned by the results dashboard. +Segmentation previews, YOLO heatmaps, and loss/metric artifacts are grouped on +the results dashboard, and YOLO-style `results.csv` files are parsed into +lightweight training curves. Job APIs and the SSE log stream also expose structured progress parsed from YOLO, MMSeg/MMEngine, SegModel-style epoch logs, and generic tqdm percentages, so the queue and live log panel can show stage, epoch/iteration, and percent @@ -139,7 +142,8 @@ The backend exposes all current Seg capabilities as job types. Examples: `segmodel.batch_predict`, `segmodel.flops`, `segmodel.params_flops`, `segmodel.benchmark`, `segmodel.raw_mask_check` - `yolo.train`, `yolo.batch_train`, `yolo.predict`, `yolo.batch_predict`, - `yolo.train_custom`, `yolo.heatmap`, `yolo.compare`, + `yolo.train_custom`, `yolo.predict_custom`, `yolo.heatmap`, + `yolo.heatmap_custom`, `yolo.compare`, `yolo.raw_mask_check`, `yolo.video_visible` - `mmseg.generate_data`, `mmseg.generate_alg`, `mmseg.train`, `mmseg.metrics`, `mmseg.flops_fps`, `mmseg.draw`, `mmseg.extract_loss_miou` diff --git a/backend/app/agents/evaluation_agent.py b/backend/app/agents/evaluation_agent.py index f442743..97b7dcc 100644 --- a/backend/app/agents/evaluation_agent.py +++ b/backend/app/agents/evaluation_agent.py @@ -14,6 +14,8 @@ REQUIRED_TASKS = { "segmodel.train": "job", "segmodel.predict": "job", "yolo.heatmap": "job", + "yolo.predict_custom": "job", + "yolo.heatmap_custom": "job", "yolo.video_visible": "job", "mmseg.flops_fps": "job", "analysis.all": "job", @@ -41,6 +43,7 @@ def evaluate_project() -> dict: "left_nav_dataset": "数据集" in frontend_text and "#datasets" in frontend_text, "upload_ui": "uploadDatasetFiles" in frontend_text and "labels" in frontend_text and "masks" in frontend_text, "dataset_quality_ui": "DatasetQuality" in frontend_text and "generateSelectedYoloYaml" in frontend_text, + "uploaded_yolo_workflow_ui": "startSelectedYoloTrain" in frontend_text and "startSelectedYoloPredict" in frontend_text and "startSelectedYoloHeatmap" in frontend_text, "loss_result_ui": "loss" in frontend_text.lower() and "heatmap" in frontend_text.lower() and "CurvePanel" in frontend_text, "job_progress_ui": "JobProgressBar" in frontend_text and "progressTrack" in frontend_text, "runtime_readiness_ui": "runtimeReadiness" in frontend_text and "环境就绪" in frontend_text, @@ -61,6 +64,7 @@ def evaluate_project() -> dict: "coverage_api": "/api/coverage" in backend_text and coverage["task_build_passed"], "visual_tools": "visual.yolo11_heatmap_v2" in catalog["task_types"] and "visual.fps" in catalog["task_types"], "yolo_custom_train": "yolo.train_custom" in catalog["task_types"], + "yolo_custom_predict_heatmap": "yolo.predict_custom" in catalog["task_types"] and "yolo.heatmap_custom" in catalog["task_types"], "yolo_dataset_tools": "dataset.yolo_txt_sort" in catalog["task_types"] and "dataset.yolo_resize" in catalog["task_types"], "no_weight_to_gitea": "Do not push" in readme_text and "check_no_weight_git" in readme_text, "all_core_tasks": all(task in catalog["task_types"] for task in REQUIRED_TASKS if REQUIRED_TASKS[task] == "job"), @@ -79,7 +83,7 @@ def evaluate_project() -> dict: if coverage["unmapped_user_scripts"]: suggestions.append(f"Map remaining user-facing scripts: {len(coverage['unmapped_user_scripts'])}") if not suggestions: - suggestions.append("Current platform covers the requested control-plane features and synthetic deep training/heatmap acceptance; next focus is a user-supplied dataset end-to-end run.") + suggestions.append("Current platform covers the requested control-plane features, uploaded YOLO dataset train/predict/heatmap actions, and synthetic deep training/heatmap acceptance; next focus is running a real user-supplied dataset through the full workflow.") score = sum(1 for item in checks if item["passed"]) / max(len(checks), 1) return { diff --git a/backend/app/agents/validation_agent.py b/backend/app/agents/validation_agent.py index d3978dd..a76799c 100644 --- a/backend/app/agents/validation_agent.py +++ b/backend/app/agents/validation_agent.py @@ -51,6 +51,7 @@ def validate_project(run_build: bool = False, run_acceptance: bool | None = None checks.append({"name": "catalog_has_yolo_heatmap", "passed": "yolo.heatmap" in catalog["task_types"]}) checks.append({"name": "catalog_has_yolo_custom_train", "passed": "yolo.train_custom" in catalog["task_types"]}) + checks.append({"name": "catalog_has_yolo_custom_predict_heatmap", "passed": "yolo.predict_custom" in catalog["task_types"] and "yolo.heatmap_custom" in catalog["task_types"]}) checks.append({"name": "catalog_has_mmseg_31_algs", "passed": len(catalog["mmseg_algorithms"]) >= 31}) checks.append({"name": "catalog_has_visual_tools", "passed": "visual.yolo11_heatmap_v2" in catalog["task_types"] and "visual.fps" in catalog["task_types"]}) checks.append({"name": "catalog_has_yolo_dataset_tools", "passed": "dataset.yolo_txt_sort" in catalog["task_types"] and "dataset.yolo_convert_png" in catalog["task_types"]}) @@ -111,10 +112,10 @@ def validate_project(run_build: bool = False, run_acceptance: bool | None = None health = _fetch(f"{backend_url}/api/health") datasets = _fetch(f"{backend_url}/api/datasets") live_jobs = _fetch(f"{backend_url}/api/jobs") - live_readiness = _fetch(f"{backend_url}/api/system/readiness") - live_capabilities = _fetch(f"{backend_url}/api/capabilities") - live_coverage = _fetch(f"{backend_url}/api/coverage") - live_curves = _fetch(f"{backend_url}/api/results/curves") + live_readiness = _fetch(f"{backend_url}/api/system/readiness", timeout=20) + live_capabilities = _fetch(f"{backend_url}/api/capabilities", timeout=20) + live_coverage = _fetch(f"{backend_url}/api/coverage", timeout=20) + live_curves = _fetch(f"{backend_url}/api/results/curves", timeout=20) frontend = _fetch(frontend_url) checks.append({"name": "live_backend_health", "passed": health["passed"] and '"ok":true' in health.get("body", "").replace(" ", ""), "detail": health}) checks.append({"name": "live_dataset_api", "passed": datasets["passed"] and datasets.get("body", "").lstrip().startswith("["), "detail": datasets}) diff --git a/backend/app/capabilities.py b/backend/app/capabilities.py index 7a0cec5..ecb3148 100644 --- a/backend/app/capabilities.py +++ b/backend/app/capabilities.py @@ -59,8 +59,10 @@ CAPABILITY_GROUPS = [ "yolo.train", "yolo.train_custom", "yolo.predict", + "yolo.predict_custom", "yolo.batch_predict", "yolo.heatmap", + "yolo.heatmap_custom", "yolo.compare", "yolo.raw_mask_check", "yolo.video_visible", diff --git a/backend/app/catalog.py b/backend/app/catalog.py index 9610950..64d0795 100644 --- a/backend/app/catalog.py +++ b/backend/app/catalog.py @@ -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"}, diff --git a/backend/app/modules/results/service.py b/backend/app/modules/results/service.py index 8c6ac1b..f52ffb6 100644 --- a/backend/app/modules/results/service.py +++ b/backend/app/modules/results/service.py @@ -28,6 +28,7 @@ def result_roots() -> list[Path]: source / "Tool-可视化" / "runs", source / "Tool-可视化" / "Data" / "result", source / "Tool-图片堆叠" / "result_0.3透明度", + project / "var" / "custom_yolo_runs", ] upload_root = project / "var" / "uploads" / "datasets" if upload_root.exists(): diff --git a/backend/app/modules/yolo/custom_heatmap.py b/backend/app/modules/yolo/custom_heatmap.py new file mode 100644 index 0000000..f8458ca --- /dev/null +++ b/backend/app/modules/yolo/custom_heatmap.py @@ -0,0 +1,144 @@ +from __future__ import annotations + +import argparse +import importlib.util +import json +import os +import sys +import types +from pathlib import Path + +DEFAULT_PROJECT_ROOT = Path(__file__).resolve().parents[4] + + +def project_root() -> Path: + raw = os.getenv("SEG_DATA_SERVER_ROOT") + if not raw: + return DEFAULT_PROJECT_ROOT + path = Path(raw).expanduser() + if path.is_absolute(): + return path.resolve() + return (DEFAULT_PROJECT_ROOT / path).resolve() + + +PROJECT_ROOT = project_root() + + +def source_root() -> Path: + raw = os.getenv("SEG_SOURCE_ROOT") + if not raw: + return (PROJECT_ROOT.parent / "Seg").resolve() + path = Path(raw).expanduser() + if path.is_absolute(): + return path.resolve() + return (PROJECT_ROOT.parent / path).resolve() + + +SOURCE_ROOT = source_root() +YOLO_DIR = SOURCE_ROOT / "Seg_All_In_One_YoloModel" +LEGACY_HEATMAP = YOLO_DIR / "yolo_predict_visualize_nn.py" +IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"} + + +def resolve_project_path(value: str | Path) -> Path: + path = Path(value).expanduser() + if path.is_absolute(): + return path.resolve() + return (PROJECT_ROOT / path).resolve() + + +def iter_images(source: Path, limit: int) -> list[Path]: + if source.is_file(): + return [source] + images = sorted(path for path in source.rglob("*") if path.is_file() and path.suffix.lower() in IMAGE_EXTS) + return images[:limit] if limit > 0 else images + + +def load_legacy_module(): + fake_config = types.ModuleType("yolo_config") + fake_config.MODEL_CONFIGS = { + "YOLOv8n-seg": {}, + "YOLOv8s-seg": {}, + "YOLOv8m-seg": {}, + "YOLOv8l-seg": {}, + "YOLOv8x-seg": {}, + "YOLOv9c-seg": {}, + "YOLOv9e-seg": {}, + "YOLO11n-seg": {}, + "YOLO11s-seg": {}, + "YOLO11m-seg": {}, + "YOLO11l-seg": {}, + "YOLO11x-seg": {}, + "YOLO12-seg": {}, + } + fake_config.TEST_IMAGE_DIR = Path(".") + fake_config.PREDICT_BEST_MODEL_DIR = Path(".") + fake_config.show_config_summary = lambda: None + sys.modules["yolo_config"] = fake_config + sys.path.insert(0, str(YOLO_DIR)) + spec = importlib.util.spec_from_file_location("seg_yolo_heatmap", LEGACY_HEATMAP) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot import heatmap script: {LEGACY_HEATMAP}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def main() -> None: + parser = argparse.ArgumentParser(description="Generate Grad-CAM style YOLO heatmaps for uploaded/custom datasets.") + parser.add_argument("--weights", required=True, help="Path to best.pt/last.pt or another YOLO segmentation checkpoint.") + parser.add_argument("--source", required=True, help="Image file or image directory.") + parser.add_argument("--project", default=str(PROJECT_ROOT / "var" / "custom_yolo_runs")) + parser.add_argument("--name", default="custom_heatmap") + parser.add_argument("--model-key", default="YOLO11n-seg") + parser.add_argument("--pt-name", default="best.pt") + parser.add_argument("--cam-method", default="GradCAM") + parser.add_argument("--target-layers", default="model.model.model[9]") + parser.add_argument("--limit", type=int, default=3) + args = parser.parse_args() + + weights = resolve_project_path(args.weights) + source = resolve_project_path(args.source) + project = resolve_project_path(args.project) + save_dir = project / args.name + if not weights.exists(): + raise FileNotFoundError(f"weights not found: {weights}") + if not source.exists(): + raise FileNotFoundError(f"source not found: {source}") + + images = iter_images(source, args.limit) + if not images: + raise FileNotFoundError(f"no images found in source: {source}") + + module = load_legacy_module() + methods = list(module.CAM_METHODS.keys()) if args.cam_method == "All" else [args.cam_method] + save_dir.mkdir(parents=True, exist_ok=True) + for method in methods: + for image in images: + module.visualize_nn_comprehensive( + model_path=str(weights), + source_dir=str(image), + base_save_dir=save_dir, + pt_name=args.pt_name, + cam_method_name=method, + target_layer_str=args.target_layers, + model_key=args.model_key, + ) + + outputs = sorted((save_dir / "HeartMap_Visual").rglob("*.jpg")) if (save_dir / "HeartMap_Visual").exists() else [] + metadata = { + "weights": str(weights), + "source": str(source), + "save_dir": str(save_dir), + "images": [str(image) for image in images], + "methods": methods, + "outputs": len(outputs), + } + (save_dir / "heatmap_manifest.json").write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8") + if not outputs: + raise RuntimeError("heatmap generation completed without image outputs") + print(json.dumps(metadata, ensure_ascii=False)) + + +if __name__ == "__main__": + main() diff --git a/backend/app/modules/yolo/custom_predict.py b/backend/app/modules/yolo/custom_predict.py new file mode 100644 index 0000000..65ad7d5 --- /dev/null +++ b/backend/app/modules/yolo/custom_predict.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path + +from ultralytics import YOLO + +DEFAULT_PROJECT_ROOT = Path(__file__).resolve().parents[4] + + +def project_root() -> Path: + raw = os.getenv("SEG_DATA_SERVER_ROOT") + if not raw: + return DEFAULT_PROJECT_ROOT + path = Path(raw).expanduser() + if path.is_absolute(): + return path.resolve() + return (DEFAULT_PROJECT_ROOT / path).resolve() + + +PROJECT_ROOT = project_root() + + +def resolve_project_path(value: str | Path) -> Path: + path = Path(value).expanduser() + if path.is_absolute(): + return path.resolve() + return (PROJECT_ROOT / path).resolve() + + +def main() -> None: + parser = argparse.ArgumentParser(description="Predict segmentation masks with a supplied YOLO checkpoint.") + parser.add_argument("--weights", required=True, help="Path to best.pt/last.pt or another YOLO segmentation checkpoint.") + parser.add_argument("--source", required=True, help="Image file or image directory to predict.") + parser.add_argument("--project", default=str(PROJECT_ROOT / "var" / "custom_yolo_runs")) + parser.add_argument("--name", default="custom_predict") + parser.add_argument("--imgsz", type=int, default=640) + parser.add_argument("--conf", type=float, default=0.25) + parser.add_argument("--device", default="cpu") + parser.add_argument("--save-txt", action="store_true") + parser.add_argument("--save-conf", action="store_true") + parser.add_argument("--exist-ok", action="store_true") + args = parser.parse_args() + + weights = resolve_project_path(args.weights) + source = resolve_project_path(args.source) + project = resolve_project_path(args.project) + if not weights.exists(): + raise FileNotFoundError(f"weights not found: {weights}") + if not source.exists(): + raise FileNotFoundError(f"source not found: {source}") + + model = YOLO(str(weights)) + results = model.predict( + source=str(source), + imgsz=args.imgsz, + conf=args.conf, + device=args.device, + project=str(project), + name=args.name, + save=True, + save_txt=args.save_txt, + save_conf=args.save_conf, + exist_ok=args.exist_ok, + verbose=True, + ) + save_dir = Path(getattr(model.predictor, "save_dir", project / args.name)) + metadata = { + "weights": str(weights), + "source": str(source), + "save_dir": str(save_dir), + "count": len(results), + } + save_dir.mkdir(parents=True, exist_ok=True) + (save_dir / "predict_manifest.json").write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8") + print(json.dumps(metadata, ensure_ascii=False)) + + +if __name__ == "__main__": + main() diff --git a/backend/app/modules/yolo/custom_train.py b/backend/app/modules/yolo/custom_train.py index f6b304e..a2243dd 100644 --- a/backend/app/modules/yolo/custom_train.py +++ b/backend/app/modules/yolo/custom_train.py @@ -1,10 +1,25 @@ from __future__ import annotations import argparse +import os from pathlib import Path from ultralytics import YOLO +DEFAULT_PROJECT_ROOT = Path(__file__).resolve().parents[4] + + +def project_root() -> Path: + raw = os.getenv("SEG_DATA_SERVER_ROOT") + if not raw: + return DEFAULT_PROJECT_ROOT + path = Path(raw).expanduser() + if path.is_absolute(): + return path.resolve() + return (DEFAULT_PROJECT_ROOT / path).resolve() + + +PROJECT_ROOT = project_root() MODEL_ALIASES = { "YOLOv8n-seg": "yolov8n-seg.pt", @@ -22,6 +37,13 @@ MODEL_ALIASES = { } +def resolve_project_path(value: str) -> Path: + path = Path(value).expanduser() + if path.is_absolute(): + return path.resolve() + return (PROJECT_ROOT / path).resolve() + + def resolve_model(value: str) -> str: candidate = MODEL_ALIASES.get(value, value) path = Path(candidate).expanduser() @@ -46,13 +68,13 @@ def main() -> None: model = YOLO(resolve_model(args.model)) result = model.train( - data=str(Path(args.data).expanduser().resolve()), + data=str(resolve_project_path(args.data)), epochs=args.epochs, imgsz=args.imgsz, batch=args.batch, workers=args.workers, device=args.device, - project=str(Path(args.project).expanduser().resolve()), + project=str(resolve_project_path(args.project)), name=args.name, exist_ok=args.exist_ok, verbose=True, diff --git a/backend/app/modules/yolo/tasks.py b/backend/app/modules/yolo/tasks.py index f84d4ad..c56ae04 100644 --- a/backend/app/modules/yolo/tasks.py +++ b/backend/app/modules/yolo/tasks.py @@ -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) diff --git a/backend/tests/test_catalog.py b/backend/tests/test_catalog.py index 4e2afd3..6d3b489 100644 --- a/backend/tests/test_catalog.py +++ b/backend/tests/test_catalog.py @@ -8,6 +8,8 @@ def test_catalog_contains_required_capabilities(): "dataset.video_frames", "segmodel.train", "yolo.train_custom", + "yolo.predict_custom", + "yolo.heatmap_custom", "yolo.predict", "mmseg.flops_fps", "analysis.all", diff --git a/frontend/src/main.tsx b/frontend/src/main.tsx index 76cabe1..952bfad 100644 --- a/frontend/src/main.tsx +++ b/frontend/src/main.tsx @@ -276,7 +276,9 @@ const defaultParams: Record> = { "yolo.train": { model: "YOLOv8n-seg" }, "yolo.train_custom": { model: "YOLO11n-seg", data: "var/uploads/datasets/example/dataset.yaml", epochs: 10, imgsz: 640, batch: 1, workers: 0, device: "cpu", exist_ok: true }, "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.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" }, "mmseg.generate_alg": { dataset_choice: 1, gpu_count: 1, gpu_ids: [0], schedule_mode: 2, max_epochs: 300, algorithm_choice: 1 }, "mmseg.train": { config: "configs/example.py", work_dir: "../DataSet_Public_outputs/example" }, "mmseg.metrics": { input_dir: "../Hardisk", output_dir: "../BestMode_Predict_Results_DataSet_Public", dataset_choice: 1, algorithm_choice: 0 }, @@ -558,17 +560,84 @@ function App() { } } + function customYoloWeightPath(dataset: UploadedDataset) { + const expected = `var/custom_yolo_runs/${dataset.name}/weights/best.pt`; + return results.find((item) => item.relative_path === expected || item.relative_path.endsWith(`/custom_yolo_runs/${dataset.name}/weights/best.pt`))?.relative_path ?? expected; + } + + async function createSelectedYoloYaml() { + if (!selectedDataset) return; + const classNames = selectedValidation?.classes.map((classId) => `class_${classId}`) ?? undefined; + return api<{ relative_path: string; path: string }>(`/api/datasets/${encodeURIComponent(selectedDataset.name)}/yolo-yaml`, { + method: "POST", + body: JSON.stringify({ class_names: classNames }) + }); + } + async function generateSelectedYoloYaml() { if (!selectedDataset) return; setBusy(true); try { - const classNames = selectedValidation?.classes.map((classId) => `class_${classId}`) ?? undefined; - const generated = await api<{ relative_path: string; path: string }>(`/api/datasets/${encodeURIComponent(selectedDataset.name)}/yolo-yaml`, { - method: "POST", - body: JSON.stringify({ class_names: classNames }) - }); + const generated = await createSelectedYoloYaml(); + if (!generated) return; setTaskType("yolo.train_custom"); - setParams(JSON.stringify({ model: "YOLO11n-seg", data: generated.path, epochs: 10, imgsz: 640, batch: 1, workers: 0, device: "cpu", exist_ok: true }, null, 2)); + setParams(JSON.stringify({ model: "YOLO11n-seg", data: generated.path, epochs: 10, imgsz: 640, batch: 1, workers: 0, device: "cpu", project: "var/custom_yolo_runs", name: selectedDataset.name, exist_ok: true }, null, 2)); + await refresh(); + } finally { + setBusy(false); + } + } + + async function startSelectedYoloTrain() { + if (!selectedDataset) return; + setBusy(true); + try { + const generated = await createSelectedYoloYaml(); + if (!generated) return; + await api("/api/jobs", { + method: "POST", + body: JSON.stringify({ + type: "yolo.train_custom", + params: { model: "YOLO11n-seg", data: generated.path, epochs: 10, imgsz: 640, batch: 1, workers: 0, device: "cpu", project: "var/custom_yolo_runs", name: selectedDataset.name, exist_ok: true } + }) + }); + window.location.hash = "jobs"; + await refresh(); + } finally { + setBusy(false); + } + } + + async function startSelectedYoloPredict() { + if (!selectedDataset?.absolute_layout) return; + setBusy(true); + try { + await api("/api/jobs", { + method: "POST", + body: JSON.stringify({ + type: "yolo.predict_custom", + params: { weights: customYoloWeightPath(selectedDataset), source: selectedDataset.absolute_layout.images, imgsz: 640, conf: 0.25, device: "cpu", project: "var/custom_yolo_runs", name: `${selectedDataset.name}_predict`, exist_ok: true } + }) + }); + window.location.hash = "jobs"; + await refresh(); + } finally { + setBusy(false); + } + } + + async function startSelectedYoloHeatmap() { + if (!selectedDataset?.absolute_layout) return; + setBusy(true); + try { + await api("/api/jobs", { + method: "POST", + body: JSON.stringify({ + type: "yolo.heatmap_custom", + params: { weights: customYoloWeightPath(selectedDataset), source: selectedDataset.absolute_layout.images, model_key: "YOLO11n-seg", cam_method: "GradCAM", target_layers: "model.model.model[9]", limit: 3, project: "var/custom_yolo_runs", name: `${selectedDataset.name}_heatmap` } + }) + }); + window.location.hash = "jobs"; await refresh(); } finally { setBusy(false); @@ -807,6 +876,15 @@ function App() { + + +