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>
252 lines
10 KiB
Python
252 lines
10 KiB
Python
#!/usr/bin/env python
|
||
"""用训练好的分类模型 + Grad-CAM 为每张图生成框标签。
|
||
|
||
uv run python scripts/gradcam_boxes.py --checkpoint runs/classify/best.pth
|
||
|
||
分类模型只学过"这是什么缺陷"。Grad-CAM 把分类分数对卷积特征的梯度当作通道重要性，
|
||
加权求和得到空间热力图——这是分类模型隐含的定位线索。阈值化 + 取连通域外接矩形，
|
||
就得到弱监督的框标签。
|
||
|
||
产出：
|
||
|
||
```
|
||
runs/gradcam_boxes/
|
||
├── labels/<类别>/<图名>.txt # YOLO 格式：class_id cx cy w h（归一化）
|
||
├── viz/<类别>/<图名>.png # 原图 | 热力图 | 叠加+框，供人工核对匹配程度
|
||
├── data.yaml # 类别表，可直接喂给检测训练
|
||
└── summary.json
|
||
```
|
||
|
||
**框的类别取自图片所在目录**（人工归类的结果），不是模型预测——即使模型分错，
|
||
我们要的也是该图真实类别的热力区域。
|
||
|
||
分辨率是主要限制：resnet18 在 128×128 输入下 layer4 只有 4×4、layer3 有 8×8。
|
||
`--stage` 可选层（4~7 对应 layer1~4），层越浅越精细但类别判别性越弱。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import sys
|
||
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 torch.utils.data import DataLoader # noqa: E402
|
||
|
||
from datasets import FolderClassificationDataset, classification_collate # 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 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_visualization(path: Path, image_path: Path, cam: np.ndarray, boxes, title: str) -> 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]
|
||
heat = np.asarray(
|
||
Image.fromarray((cam * 255).astype(np.uint8)).resize((width, height), Image.BILINEAR)
|
||
) / 255.0
|
||
|
||
fig, axes = plt.subplots(1, 3, figsize=(11, 4))
|
||
axes[0].imshow(image)
|
||
axes[0].set_title("original")
|
||
axes[1].imshow(heat, cmap="jet", vmin=0, vmax=1)
|
||
axes[1].set_title(f"grad-cam {cam.shape[0]}x{cam.shape[1]}")
|
||
axes[2].imshow(image)
|
||
axes[2].imshow(heat, cmap="jet", alpha=0.45, vmin=0, vmax=1)
|
||
for x1, y1, x2, y2 in boxes:
|
||
axes[2].add_patch(
|
||
Rectangle(
|
||
(x1 * width, y1 * height), (x2 - x1) * width, (y2 - y1) * height,
|
||
fill=False, edgecolor="lime", linewidth=2,
|
||
)
|
||
)
|
||
axes[2].set_title(f"overlay + {len(boxes)} box")
|
||
for ax in axes:
|
||
ax.axis("off")
|
||
fig.suptitle(title, fontsize=10)
|
||
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", help="分类模型权重")
|
||
parser.add_argument("--config", default="configs/classify/film.yaml", help="训练时的配置")
|
||
parser.add_argument("--out", default="runs/gradcam_boxes", help="输出目录")
|
||
parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE")
|
||
parser.add_argument("--device", default=None)
|
||
parser.add_argument(
|
||
"--stage", type=int, default=6,
|
||
help="取骨架 features 的第几层做 CAM（4~7 = layer1~4）；默认 6 即 layer3",
|
||
)
|
||
parser.add_argument("--threshold", type=float, default=0.4, help="热力图阈值（相对最大值）")
|
||
parser.add_argument("--padding", type=float, default=0.1, help="框四周外扩比例")
|
||
parser.add_argument(
|
||
"--max-boxes", type=int, default=1,
|
||
help="每张图最多保留几个框（按连通域内热力总和排序取前 N 个）。"
|
||
"默认 1：一张裁片通常只对应一处缺陷，多出来的框多半是噪声",
|
||
)
|
||
parser.add_argument(
|
||
"--min-area-ratio", type=float, default=0.02, help="连通域面积下限（占整图比例）"
|
||
)
|
||
parser.add_argument(
|
||
"--peak-ratio", type=float, default=0.75,
|
||
help="次要连通域的峰值至少要达到全局最大值的这个比例才保留，用于滤掉弱响应假框",
|
||
)
|
||
parser.add_argument(
|
||
"--viz-per-class", type=int, default=20, help="每个类别存多少张可视化，0 = 不存"
|
||
)
|
||
parser.add_argument("--viz-all", action="store_true", help="每张图都存可视化（很慢很占盘）")
|
||
parser.add_argument("--limit", type=int, default=0, help="只处理前 N 张，用于快速试参数")
|
||
args = parser.parse_args()
|
||
|
||
out_dir = Path(args.out)
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
logger = setup_logging(out_dir / "gradcam.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 [],
|
||
)
|
||
if args.limit:
|
||
dataset.samples = dataset.samples[: args.limit]
|
||
classes = dataset.classes
|
||
|
||
model_cfg = dict(cfg["model"])
|
||
model_cfg.setdefault("num_classes", len(classes))
|
||
model_cfg["init_from"] = None # 权重从 checkpoint 来，不要再读自监督骨架
|
||
model = build_model(model_cfg)
|
||
load_checkpoint(args.checkpoint, model, map_location=device)
|
||
model.to(device).eval()
|
||
for param in model.parameters(): # Grad-CAM 需要梯度，但不更新权重
|
||
param.requires_grad_(True)
|
||
|
||
logger.info("设备: %s | 权重: %s", describe_device(device), args.checkpoint)
|
||
logger.info("数据集 %s：%d 张，%d 类", data_cfg["root"], len(dataset), len(classes))
|
||
logger.info(
|
||
"CAM 目标层 features[%d] | 阈值 %.2f | 外扩 %.2f | 每图最多 %d 框",
|
||
args.stage, args.threshold, args.padding, args.max_boxes,
|
||
)
|
||
|
||
loader = DataLoader(
|
||
dataset,
|
||
batch_size=int(data_cfg.get("loader", {}).get("batch_size", 64)),
|
||
shuffle=False,
|
||
num_workers=int(data_cfg.get("loader", {}).get("num_workers", 4)),
|
||
collate_fn=classification_collate,
|
||
)
|
||
|
||
labels_dir = out_dir / "labels"
|
||
viz_dir = out_dir / "viz"
|
||
viz_done: Counter[str] = Counter()
|
||
box_counts: Counter[int] = Counter()
|
||
per_class_boxes: Counter[str] = Counter()
|
||
correct = 0
|
||
cam_shape = None
|
||
|
||
for step, batch in enumerate(loader):
|
||
images = batch["images"].to(device)
|
||
targets = batch["labels"].to(device)
|
||
cams = grad_cam(model.backbone, model.head, images, targets, stage=args.stage)
|
||
with torch.no_grad():
|
||
correct += int((model(images).argmax(1) == targets).sum())
|
||
cam_shape = tuple(cams.shape[-2:])
|
||
|
||
for cam, target, meta in zip(cams.cpu().numpy(), targets.cpu().tolist(), batch["meta"],
|
||
strict=True):
|
||
name = classes[target]
|
||
boxes = cam_to_boxes(
|
||
cam,
|
||
threshold=args.threshold,
|
||
min_area_ratio=args.min_area_ratio,
|
||
max_boxes=args.max_boxes,
|
||
padding=args.padding,
|
||
peak_ratio=args.peak_ratio,
|
||
)
|
||
box_counts[len(boxes)] += 1
|
||
per_class_boxes[name] += len(boxes)
|
||
|
||
source = Path(meta["path"])
|
||
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 boxes) + ("\n" if boxes else ""),
|
||
encoding="utf-8",
|
||
)
|
||
|
||
if args.viz_all or viz_done[name] < args.viz_per_class:
|
||
viz_done[name] += 1
|
||
save_visualization(
|
||
viz_dir / name / f"{source.stem}.png", source, cam, boxes,
|
||
f"{name} | {len(boxes)} box | thr={args.threshold}",
|
||
)
|
||
|
||
if step % 20 == 0:
|
||
logger.info("处理 [%d/%d]", step, len(loader))
|
||
|
||
total = sum(box_counts.values())
|
||
write_data_yaml(out_dir / "data.yaml", classes, root=out_dir)
|
||
summary = {
|
||
"checkpoint": args.checkpoint,
|
||
"dataset": str(data_cfg["root"]),
|
||
"num_images": total,
|
||
"cam_shape": list(cam_shape) if cam_shape else None,
|
||
"stage": args.stage,
|
||
"threshold": args.threshold,
|
||
"padding": args.padding,
|
||
"max_boxes": args.max_boxes,
|
||
"min_area_ratio": args.min_area_ratio,
|
||
"peak_ratio": args.peak_ratio,
|
||
"classification_accuracy_on_all": correct / total if total else 0.0,
|
||
"boxes_per_image": {str(k): v for k, v in sorted(box_counts.items())},
|
||
"boxes_per_class": dict(sorted(per_class_boxes.items())),
|
||
}
|
||
(out_dir / "summary.json").write_text(
|
||
json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8"
|
||
)
|
||
|
||
logger.info("CAM 分辨率 %s | 全量分类准确率 %.4f", cam_shape, summary["classification_accuracy_on_all"])
|
||
for count, number in sorted(box_counts.items()):
|
||
logger.info(" 生成 %d 个框的图片: %d 张 (%.1f%%)", count, number, 100 * number / total)
|
||
logger.info("标签写入 %s，可视化 %d 张写入 %s", labels_dir, sum(viz_done.values()), viz_dir)
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|