Shaun ZhouandClaude Opus 5
9aa87d31a6
新增基于 SHAP Partition explainer 的定位，与 Grad-CAM 形成互补对照： - engine/shap_explain.py：对目标类别求 SHAP 归因（只保留正贡献），以及方框平滑。 遮挡用数值 masker 而非教程里的 inpaint_telea——后者依赖 opencv，而我们的图已做过 逐图标准化，用 0（即该图均值）遮挡语义更干净 - scripts/shap_boxes.py：按类别抽样生成标签，并输出「原图 | Grad-CAM | SHAP」并排图 - shap 放入可选依赖组：uv sync --group explain 实测（225 张，每类 15）：SHAP 在小目标类别上明显更准，框面积占比 blackdot 0.019 vs 0.269（14 倍）、mosquito 0.030 vs 0.173、whitedot 0.076 vs 0.253。 原因是 Grad-CAM 受限于 8x8 特征图（每格 16x16 像素），几像素的黑点根本不可分辨。 代价是单张 0.55 秒，全量 11k 张约 1.7 小时，而 Grad-CAM 全量只要 2 分钟。 调参时发现的坑：SHAP 第一版全部出 0 框，因为沿用了给 8x8 调的 min_area_ratio=0.02（在 128x128 上等于 328 像素）。参数不能跨分辨率照搬， SHAP 默认改为 0.002。 --max-boxes 两个脚本都默认改为 1（原为 2）：一张裁片通常只对应一处缺陷。 库函数 cam_to_boxes 的默认值同步收紧，相关测试改为显式传参而不依赖默认值。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
240 lines
9.6 KiB
Python
240 lines
9.6 KiB
Python
#!/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())
|