refactor: 按实际的五步工作流裁掉无关代码

实际工作流是:simclr 训练 -> 聚类 -> 人工归类 -> 分类训练 -> GradCAM/SHAP 生成
热力图与 box。据此移除不参与其中的检测/定位链路,共 31 个文件、约 2600 行:

- 模型:defect_net、postprocess、necks/、heads/center_head
- 数据:film_defect(YOLO 读取端)、target_encoder
- 引擎:trainer、evaluator、losses
- 入口:train.py、eval.py
- 脚本:convert_annotations、split_dataset、export_onnx、profile_inference
- 配置:configs/{train,eval,model,dataset}
- utils/labels(多类/单类标签模式,仅检测分类头使用)

共享模块按实际用量瘦身:engine/metrics 只留 AverageMeter(检测指标随之移除);
datasets/yolo 只留写出侧(IMAGE_SUFFIXES 与 write_data_yaml)——第 5 步产出的
YOLO 标签是给外部检测训练用的交换格式,读取端不再需要。

裁剪过程中修掉一个真 bug:generate_boxes 的 --stage 默认 6(layer3)对 tiny_conv
这类非 nn.Sequential 骨架直接抛异常。现在自动退回骨架最终输出并警告,
另加 --stage -1 显式选择最终输出。

五步工作流已端到端实跑验证(simclr 1 轮 -> 聚类 -> 导出 -> 分类 1 轮 ->
gradcam 与 shap 两种方法各生成一遍)。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
2026-09-11 10:27:20 +08:00
co-authored by Claude Opus 5
parent 57446c2f10
commit 86c86b105c
45 changed files with 186 additions and 3218 deletions
+2 -2
View File
@@ -1,4 +1,4 @@
from models.heads.center_head import CenterHead
from models.heads.classification import ClassificationHead
from models.heads.projection import ProjectionHead
__all__ = ["CenterHead", "ClassificationHead"]
__all__ = ["ClassificationHead", "ProjectionHead"]
-51
View File
@@ -1,51 +0,0 @@
"""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),
}