Shaun ZhouandClaude Opus 5
fc7136fddf
分类模型只学过"是什么缺陷"，Grad-CAM 把分类分数对卷积特征的梯度当作通道重要性， 加权求和得到空间热力图，阈值化 + 连通域外接矩形即得到框标签，输出 YOLO 格式， 可直接接入既有的检测训练链路。 - engine/gradcam.py：整批一次 backward 的 Grad-CAM（各样本只回传自己那一类的分数）、 numpy 实现的 4-邻接连通域标记（不为这一处引入 scipy/opencv）、CAM 转归一化框 - scripts/gradcam_boxes.py：遍历数据集生成标签，并输出「原图 | 热力图 | 叠加+框」 三联图供人工核对；框的类别取自图片所在目录而非模型预测 两处按实图定下的默认值： - --stage 默认 6（layer3, 8×8）。layer4 只有 4×4 太粗，layer2 响应太散导致 40% 出不了框 - --peak-ratio 0.75：次要连通域峰值需达全局最大值的 75%，否则会在无缺陷的角落刷假框； 实测把两框比例从 53% 降到 20% 全量 11,212 张已生成：61.8% 出 1 框、35.2% 出 2 框、3.0% 无框，每类留 25 张可视化。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
248 lines
10 KiB
Python
248 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=2, help="每张图最多几个框")
|
||
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())
|