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>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
33fdf396ee
commit
fc7136fddf
@@ -0,0 +1,159 @@
|
||||
"""Grad-CAM 与由热力图导出检测框。
|
||||
|
||||
分类模型只学过"这是什么缺陷"，没学过"缺陷在哪"。Grad-CAM 把分类分数对某层特征图的
|
||||
梯度当作该层各通道的重要性，加权求和后得到一张空间热力图——分类模型隐含的定位线索。
|
||||
再对热力图阈值化、取连通域外接矩形，就得到弱监督的框标签。
|
||||
|
||||
分辨率是主要限制：resnet18 在 128×128 输入下 layer4 只有 4×4、layer3 有 8×8。
|
||||
层越浅空间越精细，但类别判别性越弱，需要按数据实测取舍。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
def grad_cam(
|
||||
backbone: nn.Module,
|
||||
head: nn.Module,
|
||||
images: torch.Tensor,
|
||||
targets: torch.Tensor,
|
||||
stage: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""对一个 batch 求 Grad-CAM，返回 (B, h, w) 的热力图（每张已归一化到 [0, 1]）。
|
||||
|
||||
Args:
|
||||
stage: 取 `backbone.features` 的第几层输出做 CAM；None 表示用骨架最终输出
|
||||
（即最后一个卷积块）。层越浅分辨率越高、类别判别性越弱。
|
||||
|
||||
对每个样本只回传它自己那一类的分数，所以整批可以一次 backward——
|
||||
样本之间的梯度互不干扰。
|
||||
"""
|
||||
layers = getattr(backbone, "features", None)
|
||||
if stage is not None and not isinstance(layers, nn.Sequential):
|
||||
raise ValueError("该骨架不是 nn.Sequential，无法按 stage 取中间层")
|
||||
|
||||
with torch.enable_grad():
|
||||
if stage is None:
|
||||
activations = backbone(images)
|
||||
features = activations
|
||||
else:
|
||||
features = images
|
||||
for index, layer in enumerate(layers):
|
||||
features = layer(features)
|
||||
if index == stage:
|
||||
break
|
||||
else:
|
||||
raise ValueError(f"stage {stage} 超出骨架层数 {len(layers)}")
|
||||
activations = features
|
||||
for layer in list(layers)[stage + 1 :]:
|
||||
features = layer(features)
|
||||
|
||||
activations.retain_grad()
|
||||
logits = head(features)
|
||||
# 只取每个样本对应目标类别的分数；求和后一次 backward 即可
|
||||
score = logits.gather(1, targets.view(-1, 1)).sum()
|
||||
backbone.zero_grad(set_to_none=True)
|
||||
head.zero_grad(set_to_none=True)
|
||||
score.backward()
|
||||
|
||||
gradients = activations.grad
|
||||
if gradients is None:
|
||||
raise RuntimeError("没有拿到梯度，检查目标层是否参与了前向计算")
|
||||
|
||||
weights = gradients.mean(dim=(2, 3), keepdim=True) # 各通道的重要性
|
||||
cam = F.relu((weights * activations).sum(dim=1)) # (B, h, w)
|
||||
|
||||
flat = cam.flatten(1)
|
||||
lo = flat.min(dim=1, keepdim=True).values
|
||||
hi = flat.max(dim=1, keepdim=True).values
|
||||
cam = ((flat - lo) / (hi - lo).clamp(min=1e-8)).view_as(cam)
|
||||
return cam.detach()
|
||||
|
||||
|
||||
def connected_components(mask: np.ndarray) -> list[np.ndarray]:
|
||||
"""4-邻接连通域标记，返回每个连通域的布尔掩码。
|
||||
|
||||
自己写而不引 scipy/opencv：项目只依赖 numpy/torch/PIL，
|
||||
为这一处再加一个重依赖不划算。
|
||||
"""
|
||||
height, width = mask.shape
|
||||
seen = np.zeros_like(mask, dtype=bool)
|
||||
components = []
|
||||
for start_y in range(height):
|
||||
for start_x in range(width):
|
||||
if not mask[start_y, start_x] or seen[start_y, start_x]:
|
||||
continue
|
||||
stack = [(start_y, start_x)]
|
||||
seen[start_y, start_x] = True
|
||||
pixels = []
|
||||
while stack:
|
||||
y, x = stack.pop()
|
||||
pixels.append((y, x))
|
||||
for ny, nx in ((y - 1, x), (y + 1, x), (y, x - 1), (y, x + 1)):
|
||||
if 0 <= ny < height and 0 <= nx < width and mask[ny, nx] and not seen[ny, nx]:
|
||||
seen[ny, nx] = True
|
||||
stack.append((ny, nx))
|
||||
component = np.zeros_like(mask, dtype=bool)
|
||||
ys, xs = zip(*pixels, strict=True)
|
||||
component[list(ys), list(xs)] = True
|
||||
components.append(component)
|
||||
return components
|
||||
|
||||
|
||||
def cam_to_boxes(
|
||||
cam: np.ndarray,
|
||||
threshold: float = 0.4,
|
||||
min_area_ratio: float = 0.01,
|
||||
max_boxes: int = 3,
|
||||
padding: float = 0.0,
|
||||
peak_ratio: float = 0.75,
|
||||
) -> list[tuple[float, float, float, float]]:
|
||||
"""把一张热力图变成若干归一化 xyxy 框。
|
||||
|
||||
Args:
|
||||
threshold: 相对最大值的阈值；CAM 已归一化到 [0,1]，所以直接比较。
|
||||
min_area_ratio: 连通域面积小于整图这个比例就丢掉，滤掉零星的高响应点。
|
||||
max_boxes: 最多保留几个框（按连通域内热力总和排序）。
|
||||
padding: 框的四周外扩比例——CAM 倾向只覆盖最具判别性的部位（比如缺陷中心），
|
||||
往往比真实缺陷范围小，适当外扩更贴合。
|
||||
peak_ratio: 连通域的峰值至少要达到全局最大值的这个比例才保留。
|
||||
CAM 归一化后最强的区域峰值恒为 1，次强区域若只有 0.45 多半是噪声——
|
||||
实测不加这条会在没有缺陷的角落刷出假框。
|
||||
|
||||
返回的坐标已归一化到 [0, 1]，与图像实际尺寸无关，可直接写 YOLO 标签。
|
||||
"""
|
||||
height, width = cam.shape
|
||||
mask = cam >= threshold
|
||||
if not mask.any():
|
||||
return []
|
||||
|
||||
boxes = []
|
||||
for component in connected_components(mask):
|
||||
if component.sum() < max(1, min_area_ratio * height * width):
|
||||
continue
|
||||
if cam[component].max() < peak_ratio: # 弱响应区多半是噪声，不是第二个缺陷
|
||||
continue
|
||||
ys, xs = np.where(component)
|
||||
# +1 让框覆盖住整个格子，而不是格子中心之间的距离
|
||||
x1, x2 = xs.min() / width, (xs.max() + 1) / width
|
||||
y1, y2 = ys.min() / height, (ys.max() + 1) / height
|
||||
if padding:
|
||||
dx, dy = (x2 - x1) * padding, (y2 - y1) * padding
|
||||
x1, x2 = max(0.0, x1 - dx), min(1.0, x2 + dx)
|
||||
y1, y2 = max(0.0, y1 - dy), min(1.0, y2 + dy)
|
||||
boxes.append((float(cam[component].sum()), (x1, y1, x2, y2)))
|
||||
|
||||
boxes.sort(key=lambda item: -item[0])
|
||||
return [box for _, box in boxes[:max_boxes]]
|
||||
|
||||
|
||||
def to_yolo_line(box: tuple[float, float, float, float], class_id: int) -> str:
|
||||
"""归一化 xyxy -> YOLO 的 `class_id cx cy w h`。"""
|
||||
x1, y1, x2, y2 = box
|
||||
return (
|
||||
f"{class_id} {(x1 + x2) / 2:.6f} {(y1 + y2) / 2:.6f} {x2 - x1:.6f} {y2 - y1:.6f}"
|
||||
)
|
||||
Reference in New Issue
Block a user