145 lines
5.0 KiB
Python
145 lines
5.0 KiB
Python
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()
|