"""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 # 弱响应区 kwargs = {"threshold": 0.4, "min_area_ratio": 0.0, "max_boxes": 5} # 显式放开数量上限 assert len(cam_to_boxes(cam, peak_ratio=0.0, **kwargs)) == 2 assert len(cam_to_boxes(cam, peak_ratio=0.75, **kwargs)) == 1 def test_max_boxes_defaults_to_one(): """默认只留最强的一个框——一张裁片通常只对应一处缺陷。""" cam = np.zeros((12, 12), dtype=np.float32) cam[1:4, 1:4] = 1.0 cam[7:10, 7:10] = 0.95 # 峰值够高，不会被 peak_ratio 滤掉 assert len(cam_to_boxes(cam, threshold=0.4, min_area_ratio=0.0)) == 1 assert len(cam_to_boxes(cam, threshold=0.4, min_area_ratio=0.0, max_boxes=2)) == 2 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, max_boxes=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)