feat: 用 Grad-CAM 从分类模型生成弱监督框标签

分类模型只学过"是什么缺陷"，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>
2026-09-10 15:46:30 +08:00
parent 33fdf396ee
commit fc7136fddf
4 changed files with 631 additions and 0 deletions
@@ -0,0 +1,180 @@
+"""Grad-CAM 与由热力图导出框。"""
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+import numpy as np  # noqa: E402
+from torch import nn  # noqa: E402
+
+from engine.gradcam import (  # noqa: E402
+    cam_to_boxes,
+    connected_components,
+    grad_cam,
+    to_yolo_line,
+)
+from models import build_model  # noqa: E402
+
+
+def test_connected_components_separates_blobs():
+    mask = np.array([[1, 1, 0, 0], [1, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 1]], dtype=bool)
+    components = connected_components(mask)
+    assert len(components) == 2
+    assert sorted(int(c.sum()) for c in components) == [3, 4]
+    assert not (components[0] & components[1]).any()   # 互不重叠
+    assert (components[0] | components[1] == mask).all()
+
+
+def test_connected_components_is_four_connected():
+    """对角相邻不算连通——否则细长划痕会和邻近噪点粘成一个框。"""
+    mask = np.array([[1, 0], [0, 1]], dtype=bool)
+    assert len(connected_components(mask)) == 2
+
+
+def test_cam_to_boxes_finds_the_hot_region():
+    cam = np.zeros((8, 8), dtype=np.float32)
+    cam[2:5, 3:7] = 1.0
+    (box,) = cam_to_boxes(cam, threshold=0.5, min_area_ratio=0.0, peak_ratio=0.75)
+
+    x1, y1, x2, y2 = box
+    assert (x1, x2) == pytest.approx((3 / 8, 7 / 8))   # 框覆盖整格，不是格心距离
+    assert (y1, y2) == pytest.approx((2 / 8, 5 / 8))
+
+
+def test_padding_expands_within_bounds():
+    cam = np.zeros((8, 8), dtype=np.float32)
+    cam[3:5, 3:5] = 1.0
+    plain = cam_to_boxes(cam, threshold=0.5, min_area_ratio=0.0)[0]
+    padded = cam_to_boxes(cam, threshold=0.5, min_area_ratio=0.0, padding=0.5)[0]
+
+    assert padded[0] < plain[0] and padded[2] > plain[2]
+    assert all(0.0 <= v <= 1.0 for v in padded)        # 外扩不会越界
+
+    corner = np.zeros((8, 8), dtype=np.float32)
+    corner[0:2, 0:2] = 1.0
+    assert cam_to_boxes(corner, threshold=0.5, min_area_ratio=0.0, padding=1.0)[0][:2] == (0.0, 0.0)
+
+
+def test_peak_ratio_drops_weak_secondary_blobs():
+    """次强区域峰值只有 0.45 时多半是噪声——实测不过滤会在无缺陷处刷出假框。"""
+    cam = np.zeros((8, 8), dtype=np.float32)
+    cam[1:3, 1:3] = 1.0      # 主区域
+    cam[5:7, 5:7] = 0.45     # 弱响应区
+
+    assert len(cam_to_boxes(cam, threshold=0.4, min_area_ratio=0.0, peak_ratio=0.0)) == 2
+    assert len(cam_to_boxes(cam, threshold=0.4, min_area_ratio=0.0, peak_ratio=0.75)) == 1
+
+
+def test_min_area_and_max_boxes_limits():
+    cam = np.zeros((10, 10), dtype=np.float32)
+    cam[0:1, 0:1] = 1.0      # 1 格，面积很小
+    cam[4:8, 4:8] = 1.0
+    assert len(cam_to_boxes(cam, threshold=0.5, min_area_ratio=0.05, peak_ratio=0.5)) == 1
+
+    many = np.zeros((12, 12), dtype=np.float32)
+    for offset in (0, 4, 8):
+        many[offset:offset + 2, offset:offset + 2] = 1.0
+    assert len(cam_to_boxes(many, threshold=0.5, min_area_ratio=0.0, max_boxes=2)) == 2
+
+
+def test_no_box_when_nothing_activates():
+    assert cam_to_boxes(np.zeros((8, 8), dtype=np.float32), threshold=0.5) == []
+
+
+def test_to_yolo_line_roundtrip():
+    line = to_yolo_line((0.2, 0.4, 0.6, 0.8), class_id=7)
+    parts = line.split()
+    assert parts[0] == "7"
+    cx, cy, w, h = (float(v) for v in parts[1:])
+    assert (cx, cy, w, h) == pytest.approx((0.4, 0.6, 0.4, 0.4))
+
+
+class _TinyNet(nn.Module):
+    """features 是 nn.Sequential，好按 stage 取中间层。"""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.features = nn.Sequential(
+            nn.Conv2d(3, 4, 3, padding=1), nn.ReLU(),
+            nn.Conv2d(4, 8, 3, stride=2, padding=1), nn.ReLU(),
+        )
+        self.out_channels = 8
+
+    def forward(self, x):
+        return self.features(x)
+
+
+def test_grad_cam_shapes_and_normalization():
+    torch.manual_seed(0)
+    backbone = _TinyNet()
+    head = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(8, 3))
+    images = torch.randn(2, 3, 16, 16)
+    targets = torch.tensor([0, 2])
+
+    cam = grad_cam(backbone, head, images, targets)
+    assert cam.shape == (2, 8, 8)                       # 与最终特征图同分辨率
+    assert not cam.requires_grad
+    assert float(cam.min()) >= 0.0 and float(cam.max()) <= 1.0
+    # 每张单独归一化：非退化的那张峰值应当正好是 1
+    peaks = cam.flatten(1).max(dim=1).values
+    assert any(float(p) == pytest.approx(1.0) for p in peaks)
+
+
+def test_grad_cam_all_zero_when_relu_kills_everything():
+    """梯度全为负时 ReLU 会清零整张图，此时归一化结果仍是合法的全 0，
+    对应下游"这张图提取不出框"——全量数据里约 3% 属于这种情况。"""
+    backbone = _TinyNet()
+    head = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(8, 3))
+    with torch.no_grad():   # 让目标类别的权重全为负，梯度随之为负
+        head[2].weight.fill_(-1.0)
+        head[2].bias.zero_()
+    with torch.no_grad():
+        backbone.features[0].weight.fill_(0.1)
+        backbone.features[2].weight.fill_(0.1)
+
+    cam = grad_cam(backbone, head, torch.rand(1, 3, 16, 16), torch.tensor([0]))
+    assert float(cam.max()) == pytest.approx(0.0)
+    assert cam_to_boxes(cam[0].numpy(), threshold=0.4) == []
+
+
+def test_grad_cam_stage_controls_resolution():
+    backbone = _TinyNet()
+    head = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(8, 3))
+    images = torch.randn(1, 3, 16, 16)
+    targets = torch.tensor([1])
+
+    # 浅层分辨率更高：stage 1 在第一次下采样之前
+    assert grad_cam(backbone, head, images, targets, stage=1).shape == (1, 16, 16)
+    assert grad_cam(backbone, head, images, targets, stage=3).shape == (1, 8, 8)
+
+    with pytest.raises(ValueError, match="超出骨架层数"):
+        grad_cam(backbone, head, images, targets, stage=99)
+
+
+def test_grad_cam_responds_to_the_target_class():
+    """不同目标类别应当给出不同的热力图，否则 CAM 就没有类别判别性。"""
+    torch.manual_seed(0)
+    backbone = _TinyNet()
+    head = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(8, 3))
+    images = torch.randn(1, 3, 16, 16)
+
+    a = grad_cam(backbone, head, images, torch.tensor([0]))
+    b = grad_cam(backbone, head, images, torch.tensor([1]))
+    assert not torch.allclose(a, b)
+
+
+def test_grad_cam_works_on_the_real_classifier():
+    model = build_model({
+        "name": "classifier", "num_classes": 5,
+        "backbone": {"name": "torchvision", "arch": "resnet18", "pretrained": False},
+    }).eval()
+    for param in model.parameters():
+        param.requires_grad_(True)
+
+    images = torch.randn(2, 3, 128, 128)
+    targets = torch.tensor([1, 4])
+    cam = grad_cam(model.backbone, model.head, images, targets, stage=6)
+
+    assert cam.shape == (2, 8, 8)        # layer3 在 128 输入下是 8x8
+    boxes = cam_to_boxes(cam[0].cpu().numpy(), threshold=0.4, min_area_ratio=0.0)
+    assert all(0.0 <= v <= 1.0 for box in boxes for v in box)