#!/usr/bin/env python """用 SHAP Partition explainer 生成框标签，并与 Grad-CAM 并排对比。 uv run python scripts/shap_boxes.py --checkpoint runs/classify/best.pth --per-class 10 与 Grad-CAM 的取舍： | | Grad-CAM | SHAP Partition | | --- | --- | --- | | 分辨率 | 受限于特征图（layer3 @128² 只有 8×8） | 与输入同分辨率（128×128） | | 单张耗时 | 一次 backward，全量 11k 张约 2 分钟 | 几百次前向，单张约 1 秒 | 所以 SHAP 默认**按类别抽样**跑（`--per-class`），用来判断值不值得为它付出时间； 确实要全量请加 `--all`，注意 11k 张需要数小时。 产出与 gradcam_boxes.py 同构：labels/ 是 YOLO 标签，viz/ 是对比图 （原图 | Grad-CAM 热力+框 | SHAP 热力+框），一眼看出谁定得更准。 """ from __future__ import annotations import argparse import json import random import sys import time from collections import Counter from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parents[1])) import numpy as np # noqa: E402 import torch # noqa: E402 from PIL import Image # noqa: E402 from datasets import FolderClassificationDataset # noqa: E402 from datasets.yolo import write_data_yaml # noqa: E402 from engine import load_checkpoint # noqa: E402 from engine.gradcam import cam_to_boxes, grad_cam, to_yolo_line # noqa: E402 from engine.shap_explain import shap_attribution, smooth # noqa: E402 from models import build_model # noqa: E402 from utils import load_config, setup_logging # noqa: E402 from utils.device import DeviceNotAvailableError, describe_device, resolve_device # noqa: E402 logger = None def save_comparison(path: Path, image_path: Path, panels: list[tuple[str, np.ndarray, list]]) -> None: """一行画出原图与各方法的热力图+框，便于直接比较定位质量。""" import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from matplotlib.patches import Rectangle with Image.open(image_path) as raw: image = np.asarray(raw.convert("RGB")) height, width = image.shape[:2] fig, axes = plt.subplots(1, len(panels) + 1, figsize=(4 * (len(panels) + 1), 4)) axes[0].imshow(image) axes[0].set_title("original") axes[0].axis("off") for ax, (title, heat, boxes) in zip(axes[1:], panels, strict=True): resized = np.asarray( Image.fromarray((heat * 255).astype(np.uint8)).resize((width, height), Image.BILINEAR) ) / 255.0 ax.imshow(image) ax.imshow(resized, cmap="jet", alpha=0.45, vmin=0, vmax=1) for x1, y1, x2, y2 in boxes: ax.add_patch( Rectangle((x1 * width, y1 * height), (x2 - x1) * width, (y2 - y1) * height, fill=False, edgecolor="lime", linewidth=2) ) ax.set_title(f"{title} ({heat.shape[0]}x{heat.shape[1]}, {len(boxes)} box)") ax.axis("off") fig.tight_layout() path.parent.mkdir(parents=True, exist_ok=True) fig.savefig(path, dpi=110) plt.close(fig) def main() -> int: global logger parser = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter ) parser.add_argument("--checkpoint", default="runs/classify/best.pth") parser.add_argument("--config", default="configs/classify/film.yaml") parser.add_argument("--out", default="runs/shap_boxes") parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE") parser.add_argument("--device", default=None) parser.add_argument("--per-class", type=int, default=10, help="每个类别抽多少张") parser.add_argument("--all", action="store_true", help="跑全量（11k 张需数小时）") parser.add_argument("--seed", type=int, default=0, help="抽样随机种子") parser.add_argument("--max-evals", type=int, default=500, help="SHAP 每张图的前向次数") parser.add_argument("--smooth-radius", type=int, default=2, help="归因图方框平滑半径，0=不平滑") parser.add_argument("--threshold", type=float, default=0.4) parser.add_argument("--padding", type=float, default=0.1) parser.add_argument( "--max-boxes", type=int, default=1, help="每张图最多保留几个框（按归因总量排序取前 N 个）。" "默认 1：一张裁片通常只对应一处缺陷，多出来的框多半是噪声", ) parser.add_argument( "--min-area-ratio", type=float, default=0.002, help="连通域面积下限（占整图比例）。默认值远小于 gradcam_boxes.py 的 0.02——" "SHAP 的归因图是 128x128 而不是 8x8，同样一个小黑点在这里只占百来像素，" "沿用 0.02（=328 像素）会把它整个滤掉", ) parser.add_argument("--peak-ratio", type=float, default=0.75) parser.add_argument("--stage", type=int, default=6, help="对比用的 Grad-CAM 目标层") parser.add_argument("--no-compare", action="store_true", help="不算 Grad-CAM，只出 SHAP") args = parser.parse_args() out_dir = Path(args.out) out_dir.mkdir(parents=True, exist_ok=True) logger = setup_logging(out_dir / "shap.log") cfg = load_config(args.config, args.set) try: device = resolve_device(args.device or cfg.get("device", "auto")) except DeviceNotAvailableError as exc: logger.error("%s", exc) return 1 data_cfg = cfg["data"] dataset = FolderClassificationDataset( root=data_cfg["root"], split="all", classes=data_cfg.get("classes"), transforms=data_cfg["val_transforms"], exclude=data_cfg.get("exclude") or [], ) classes = dataset.classes if not args.all: by_class: dict[int, list] = {} for sample in dataset.samples: by_class.setdefault(sample[1], []).append(sample) rng = random.Random(args.seed) dataset.samples = [ sample for samples in by_class.values() for sample in rng.sample(samples, min(args.per_class, len(samples))) ] model_cfg = dict(cfg["model"]) model_cfg.setdefault("num_classes", len(classes)) model_cfg["init_from"] = None model = build_model(model_cfg) load_checkpoint(args.checkpoint, model, map_location=device) model.to(device).eval() for param in model.parameters(): param.requires_grad_(True) logger.info("设备: %s | 权重: %s", describe_device(device), args.checkpoint) logger.info( "样本 %d 张（%s）| max_evals=%d | 平滑半径 %d", len(dataset), "全量" if args.all else f"每类 {args.per_class} 张", args.max_evals, args.smooth_radius, ) box_kwargs = { "threshold": args.threshold, "min_area_ratio": args.min_area_ratio, "max_boxes": args.max_boxes, "padding": args.padding, "peak_ratio": args.peak_ratio, } labels_dir = out_dir / "labels" viz_dir = out_dir / "viz" shap_counts: Counter[int] = Counter() cam_counts: Counter[int] = Counter() started = time.time() for index in range(len(dataset)): sample = dataset[index] target = sample["label"] name = classes[target] source = Path(sample["meta"]["path"]) image_chw = sample["image"] attribution = shap_attribution( model, image_chw.permute(1, 2, 0).numpy().astype(np.float32), target, device, max_evals=args.max_evals, ) attribution = smooth(attribution, args.smooth_radius) shap_boxes = cam_to_boxes(attribution, **box_kwargs) shap_counts[len(shap_boxes)] += 1 panels = [] if not args.no_compare: cam = grad_cam( model.backbone, model.head, image_chw[None].to(device), torch.tensor([target], device=device), stage=args.stage, )[0].cpu().numpy() cam_boxes = cam_to_boxes(cam, **box_kwargs) cam_counts[len(cam_boxes)] += 1 panels.append(("grad-cam", cam, cam_boxes)) panels.append(("shap", attribution, shap_boxes)) label_path = labels_dir / name / f"{source.stem}.txt" label_path.parent.mkdir(parents=True, exist_ok=True) label_path.write_text( "\n".join(to_yolo_line(box, target) for box in shap_boxes) + ("\n" if shap_boxes else ""), encoding="utf-8", ) save_comparison(viz_dir / name / f"{source.stem}.png", source, panels) if index % 20 == 0: elapsed = time.time() - started logger.info( "[%d/%d] 已用 %.0fs，平均 %.2fs/张", index, len(dataset), elapsed, elapsed / max(index, 1), ) total = sum(shap_counts.values()) write_data_yaml(out_dir / "data.yaml", classes, root=out_dir) summary = { "checkpoint": args.checkpoint, "num_images": total, "seconds_per_image": (time.time() - started) / max(total, 1), "max_evals": args.max_evals, "smooth_radius": args.smooth_radius, "box_params": box_kwargs, "shap_boxes_per_image": {str(k): v for k, v in sorted(shap_counts.items())}, "gradcam_boxes_per_image": {str(k): v for k, v in sorted(cam_counts.items())}, } (out_dir / "summary.json").write_text( json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8" ) logger.info("共 %d 张，平均 %.2fs/张", total, summary["seconds_per_image"]) logger.info("SHAP 出框分布: %s", dict(sorted(shap_counts.items()))) if cam_counts: logger.info("Grad-CAM 出框分布: %s", dict(sorted(cam_counts.items()))) logger.info("对比图写入 %s", viz_dir) return 0 if __name__ == "__main__": raise SystemExit(main())