Add custom YOLO prediction and heatmap workflow

2026-06-30 15:11:47 +08:00
parent 4d0c26be05
commit 777f168a75
12 changed files with 393 additions and 17 deletions
--- 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`
--- 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 {
--- 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})
--- 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",
--- 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"},
--- 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():
--- a/backend/app/modules/yolo/custom_heatmap.py
+++ 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()
--- a/backend/app/modules/yolo/custom_predict.py
+++ 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()
--- 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,
--- 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)

--- 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",
--- a/frontend/src/main.tsx
+++ b/frontend/src/main.tsx
@@ -276,7 +276,9 @@ const defaultParams: Record<string, Record<string, unknown>> = {
  "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<Job>("/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<Job>("/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<Job>("/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() {
                <button className="iconButton" disabled={busy || !selectedValidation?.ready.yolo} onClick={generateSelectedYoloYaml} title="生成 YOLO dataset.yaml">
                  <FileSearch size={18} />
                </button>
+                <button className="iconButton" disabled={busy || !selectedValidation?.ready.yolo} onClick={startSelectedYoloTrain} title="启动自定义 YOLO 训练">
+                  <Play size={18} />
+                </button>
+                <button className="iconButton" disabled={busy || !selectedDataset?.absolute_layout} onClick={startSelectedYoloPredict} title="使用自定义 best.pt 预测">
+                  <FileImage size={18} />
+                </button>
+                <button className="iconButton" disabled={busy || !selectedDataset?.absolute_layout} onClick={startSelectedYoloHeatmap} title="使用自定义 best.pt 生成热度图">
+                  <Zap size={18} />
+                </button>
                <FileImage size={22} />
              </div>
            </div>