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>
84 lines
3.1 KiB
Python
84 lines
3.1 KiB
Python
"""缺陷分类 + 定位主模型。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from collections.abc import Mapping
|
||
from typing import Any
|
||
|
||
import torch
|
||
from torch import nn
|
||
|
||
from models.heads.center_head import CenterHead
|
||
from models.heads.classification import ClassificationHead
|
||
from models.registry import BACKBONES, MODELS, NECKS
|
||
from utils.labels import LABEL_MODES, MULTILABEL, check_label_mode
|
||
|
||
__all__ = ["DefectNet", "classification_outputs", "LABEL_MODES"]
|
||
|
||
|
||
def classification_outputs(num_classes: int, label_mode: str) -> int:
|
||
"""分类头的输出维度；multiclass 需要额外一维表示 "无缺陷"。"""
|
||
check_label_mode(label_mode)
|
||
return num_classes if label_mode == MULTILABEL else num_classes + 1
|
||
|
||
|
||
@MODELS.register("defect_net")
|
||
class DefectNet(nn.Module):
|
||
"""backbone -> neck -> {分类头, 中心点定位头}。
|
||
|
||
Args:
|
||
num_classes: 缺陷类别数（不含背景）。
|
||
backbone: backbone 注册表的配置，如 `{"name": "tiny_conv", ...}`。
|
||
neck: neck 配置，`None` 表示直接在骨架特征上出图。
|
||
head: CenterHead 的额外参数。
|
||
with_classification: 是否输出图像级分类 logits。
|
||
label_mode: 图像级标签形式，由 yaml 配置决定。
|
||
"multilabel" —— 多类：一张图可同时含多种缺陷，逐类 sigmoid + BCE；
|
||
"multiclass" —— 单类：一张图只归一个类别（含 "无缺陷"），softmax + CE。
|
||
仅影响分类头与分类损失；定位头始终按 num_classes 个类别输出 heatmap。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
num_classes: int,
|
||
backbone: Mapping[str, Any] | None = None,
|
||
neck: Mapping[str, Any] | None = None,
|
||
head: Mapping[str, Any] | None = None,
|
||
with_classification: bool = True,
|
||
classification: Mapping[str, Any] | None = None,
|
||
label_mode: str = "multilabel",
|
||
) -> None:
|
||
super().__init__()
|
||
self.num_classes = num_classes
|
||
self.label_mode = check_label_mode(label_mode)
|
||
|
||
self.backbone = BACKBONES.build(dict(backbone or {"name": "tiny_conv"}))
|
||
stride = self.backbone.out_stride
|
||
channels = self.backbone.out_channels
|
||
|
||
if neck:
|
||
self.neck: nn.Module = NECKS.build(dict(neck), in_channels=channels)
|
||
channels = self.neck.out_channels
|
||
stride = stride // self.neck.scale_factor
|
||
else:
|
||
self.neck = nn.Identity()
|
||
|
||
self.out_stride = stride
|
||
self.center_head = CenterHead(channels, num_classes, **dict(head or {}))
|
||
self.cls_head = (
|
||
ClassificationHead(
|
||
channels,
|
||
classification_outputs(num_classes, label_mode),
|
||
**dict(classification or {}),
|
||
)
|
||
if with_classification
|
||
else None
|
||
)
|
||
|
||
def forward(self, images: torch.Tensor) -> dict[str, torch.Tensor]:
|
||
feat = self.neck(self.backbone(images))
|
||
outputs = self.center_head(feat)
|
||
if self.cls_head is not None:
|
||
outputs["logits"] = self.cls_head(feat)
|
||
return outputs
|