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>
52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
"""Anchor-free 中心点定位头（CenterNet 风格，纯卷积）。
|
|
|
|
输出:
|
|
heatmap: (B, num_classes, H, W) 各类别缺陷中心的置信度（logits）
|
|
wh: (B, 2, H, W) 中心点处的框宽高（特征图尺度）
|
|
offset: (B, 2, H, W) 下采样量化误差补偿
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
from models.layers import ConvBNAct
|
|
from models.registry import HEADS
|
|
|
|
|
|
def _branch(in_channels: int, hidden: int, out_channels: int) -> nn.Sequential:
|
|
return nn.Sequential(
|
|
ConvBNAct(in_channels, hidden, kernel_size=3),
|
|
nn.Conv2d(hidden, out_channels, kernel_size=1),
|
|
)
|
|
|
|
|
|
@HEADS.register("center")
|
|
class CenterHead(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
num_classes: int,
|
|
hidden_channels: int = 64,
|
|
prior_prob: float = 0.01,
|
|
) -> None:
|
|
super().__init__()
|
|
self.num_classes = num_classes
|
|
self.heatmap = _branch(in_channels, hidden_channels, num_classes)
|
|
self.wh = _branch(in_channels, hidden_channels, 2)
|
|
self.offset = _branch(in_channels, hidden_channels, 2)
|
|
|
|
# heatmap 正样本极稀疏，用先验概率初始化最后一层 bias，稳定早期训练。
|
|
bias = -math.log((1 - prior_prob) / prior_prob)
|
|
nn.init.constant_(self.heatmap[-1].bias, bias)
|
|
|
|
def forward(self, feat: torch.Tensor) -> dict[str, torch.Tensor]:
|
|
return {
|
|
"heatmap": self.heatmap(feat),
|
|
"wh": self.wh(feat),
|
|
"offset": self.offset(feat),
|
|
}
|