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

144
backend/app/modules/yolo/custom_heatmap.py Normal file
View File

@@ -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()

82
backend/app/modules/yolo/custom_predict.py Normal file
View File

@@ -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()

26
backend/app/modules/yolo/custom_train.py
View File

@@ -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,

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)