Shaun ZhouandClaude Opus 5
3a36aad6c3
面向实时设备部署，全程只用 convolution，不含 attention / transformer。 - backbones: 自研 tiny_conv（深度可分离，默认）与 torchvision 封装（resnet/mobilenet/shufflenet） - necks: nearest 上采样 + 3x3 卷积，把特征恢复到 stride 4，兼顾小缺陷定位与速度 - heads: anchor-free 中心点定位头（heatmap + wh + offset）与图像级分类头； 分类头输出维度随 label_mode 切换（多类多热 / 单类互斥含背景位） - postprocess: 3x3 max-pool 峰值抑制替代 NMS 循环，无循环、可导出 TensorRT / OpenVINO Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
"""推理后处理：从 heatmap 解码出缺陷框。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import torch
|
||
from torch.nn import functional as F
|
||
|
||
|
||
def heatmap_nms(heatmap: torch.Tensor, kernel: int = 3) -> torch.Tensor:
|
||
"""用 max-pool 做峰值抑制，替代 NMS（无循环，可导出）。"""
|
||
pooled = F.max_pool2d(heatmap, kernel, stride=1, padding=kernel // 2)
|
||
return heatmap * (pooled == heatmap).float()
|
||
|
||
|
||
@torch.no_grad()
|
||
def decode_heatmap(
|
||
heatmap: torch.Tensor,
|
||
wh: torch.Tensor,
|
||
offset: torch.Tensor,
|
||
stride: int,
|
||
topk: int = 50,
|
||
score_threshold: float = 0.3,
|
||
) -> list[dict[str, torch.Tensor]]:
|
||
"""解码为原图坐标下的检测结果。
|
||
|
||
Args:
|
||
heatmap: (B, C, H, W) logits。
|
||
wh / offset: (B, 2, H, W)，单位为特征图像素。
|
||
stride: 特征图相对原图的下采样倍数。
|
||
|
||
Returns:
|
||
每张图一个 dict：`boxes` (N, 4) xyxy、`scores` (N,)、`labels` (N,)。
|
||
"""
|
||
scores_map = heatmap_nms(heatmap.sigmoid())
|
||
batch, _num_classes, height, width = scores_map.shape
|
||
|
||
flat = scores_map.view(batch, -1)
|
||
topk = min(topk, flat.shape[1])
|
||
topk_scores, topk_idx = flat.topk(topk, dim=1)
|
||
|
||
labels = torch.div(topk_idx, height * width, rounding_mode="floor")
|
||
pixel_idx = topk_idx % (height * width)
|
||
ys = torch.div(pixel_idx, width, rounding_mode="floor").float()
|
||
xs = (pixel_idx % width).float()
|
||
|
||
wh_flat = wh.view(batch, 2, -1)
|
||
offset_flat = offset.view(batch, 2, -1)
|
||
gather_idx = pixel_idx.unsqueeze(1).expand(-1, 2, -1)
|
||
pred_wh = wh_flat.gather(2, gather_idx)
|
||
pred_offset = offset_flat.gather(2, gather_idx)
|
||
|
||
cx = (xs + pred_offset[:, 0]) * stride
|
||
cy = (ys + pred_offset[:, 1]) * stride
|
||
half_w = pred_wh[:, 0].clamp(min=0) * stride / 2
|
||
half_h = pred_wh[:, 1].clamp(min=0) * stride / 2
|
||
boxes = torch.stack([cx - half_w, cy - half_h, cx + half_w, cy + half_h], dim=-1)
|
||
|
||
results = []
|
||
for i in range(batch):
|
||
keep = topk_scores[i] >= score_threshold
|
||
results.append(
|
||
{
|
||
"boxes": boxes[i][keep],
|
||
"scores": topk_scores[i][keep],
|
||
"labels": labels[i][keep],
|
||
}
|
||
)
|
||
return results
|