Shaun ZhouandClaude Opus 5 a4d099e41d feat(ml-backend): 把训练好的 YOLO 接进 Label Studio 做预标注
ml_backend/ 是一个 label-studio-ml 后端:Label Studio 打开任务时调 /predict,
后端跑 YOLO 并把框转成 Label Studio 的百分比坐标返回,人工只需修正。

- model.py:FilmDefectDetector。权重懒加载且只加载一次(Flask 多线程,加锁);
  控件名从项目的 labeling config 读取,不写死;model_version 带权重的修改时间,
  换权重后 Label Studio 能区分新旧预测。
- convert.py:坐标转换,刻意不 import label_studio_ml,好在主环境里单测。
- fit():每条标注累加计数,达到 RETRAIN_AFTER 就起子进程执行 RETRAIN_COMMAND
  并立刻返回——训练几十分钟,同步做会卡住 Label Studio。
- _wsgi.py 由官方模板生成,ruff 跳过,便于跟随上游。

必须用独立虚拟环境:label-studio-ml 会把 jinja2 降到 3.0.3,而 label-studio 1.23
需要 3.1,装进主环境会直接弄坏 Label Studio。PyPI 上的 1.0.9 是 2023 年的旧版,
要装 GitHub 上的现行版本。

本地用 COCO 权重端到端验证过:/health、/predict(bus.jpg 出 10 个框、坐标换算正确)、
/webhook(第 3 条标注触发重训命令并清零计数)。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-18 17:36:12 +08:00

ModelTrainEval

薄膜缺陷的无监督归类 → 分类模型 → 弱监督定位流水线。模型最终部署在实时设备上, 因此只使用基于 convolution 的结构(不含 attention / transformer)。

完整工作流五步:

  1. 自监督:在无标注数据上用 SimCLR 训练卷积骨架
  2. 聚类:按指定簇数分组,产出可人工翻看的硬链接目录
  3. 人工归类:在 runs/ 下直接整理簇(可对某一类继续聚类细化),再导出成数据集
  4. 分类训练:在归类好的数据集上训练分类模型
  5. 定位:用 Grad-CAM / SHAP 生成热力图与 YOLO 框标签

目录

目录 作用
configs/ yaml 配置(model/ 模型、dataset/ 数据、train/ 训练、eval/ 评估)
models/ 模型骨架:backbones/ + necks/ + heads/,由 defect_net.py 组装
datasets/ 抛出数据样本的 Dataset 类、数据增强、监督信号编码
data/ 数据集本体或软链接(不入库,见 data/README.md)
scripts/ 格式转换、数据集切分、ONNX 导出等批处理脚本
engine/ 训练 / 评估循环、损失、指标、优化器、检查点、SwanLab 可视化
utils/ 配置、注册表、日志、随机种子、设备解析(cuda/mps/cpu)
tests/ 单元测试
docs/ 算法与数学说明文档
figures/ 文档用的 mermaid 图(.mmd 源 + .md 预览 + .png)

入口:classify.py(分类训练与评估)、scripts/ 下的聚类、导出与框生成脚本。

环境

uv sync                      # 安装依赖(含 swanlab)
uv sync --group export       # 需要导出 ONNX 时再装
uv run python scripts/check_env.py   # 自检(Python / PyTorch / 可用后端)

计算后端

支持 CUDA / MPS(Apple Silicon)/ CPU。配置里的 device: auto 按 cuda > mps > cpu 自动挑最快的可用后端,也可以显式写 cuda / mps / cpu,或用命令行 --device 覆盖:

uv run python train.py --config configs/train/baseline.yaml               # auto
uv run python train.py --config configs/train/baseline.yaml --device mps  # 强制 Mac GPU
uv run python train.py --config configs/train/baseline.yaml --device cpu

显式指定不可用的后端会直接报错退出,不会悄悄退回 CPU——训练慢几十倍却不吭声是最坏的情况。

后端 AMP 默认 显存统计 备注
cuda 开 精确峰值 Linux 上 PyPI 的 torch 即 CUDA 构建;指定版本用 uv pip install --index-url https://download.pytorch.org/whl/cu124 torch torchvision
mps 关 采样下界 Apple Silicon。autocast 可用且数值正常,但实测比 fp32 慢(tiny_conv @128²:fp32 76 ms/步,fp16 90 ms),故默认不开;amp: true 可强制打开
cpu 不支持 无 实测比 MPS 慢约 30 倍,仅适合调试与精度回归

amp 字段接受 auto(随后端)/ true / false。MPS 上 torch.profiler 抓不到 Metal kernel (没有对应的 ProfilerActivity),算子明细只反映 CPU 侧下发耗时。

训练可视化(SwanLab)

训练过程通过 SwanLab 记录:按 step 记录各项损失与学习率,按 epoch 记录 验证指标与显存峰值,实验的超参就是完整的 yaml 配置。在 train.swanlab 里配置:

train:
  swanlab:
    enabled: true
    project: film-defect
    name: null        # 缺省用 work_dir 目录名
    mode: online      # online(云端) | local(本地看板) | offline(仅存盘) | disabled
    tags: [tiny_conv, center-head]
# 本地看板:不需要账号,训练时用 local 模式,另开终端起看板
uv run python train.py --config configs/train/baseline.yaml --set train.swanlab.mode=local
uv run swanlab watch runs/baseline/swanlog          # 默认 http://127.0.0.1:5092

# 云端
uv run swanlab login                               # 首次使用需登录
uv run python train.py --config configs/train/baseline.yaml   # mode 默认 online

# 离线跑完后再补传云端
uv run swanlab sync runs/baseline/swanlog

看板可以在训练进行中打开,指标边训边刷。端口冲突用 -p 5100,要局域网访问加 -h 0.0.0.0。 本地看板依赖 swanlab[dashboard],已包含在 uv sync 里。

未安装 swanlab、未登录或网络不通时会打一条警告并继续训练,不会中断。关掉用 --set train.swanlab.enabled=false。

未分类样本的无监督聚类

调用 ImageNet 预训练模型

data/raw/film 里大量样本还没人工分完。这条流水线把目录下所有图片拉平(不看任何 目录名),做无监督聚类,产出可人工核对的分组:

# ImageNet 预训练卷积特征,不训练,CPU 上 14.5k 张约 100 秒
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml

# 先扫一遍簇数再决定分几类
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml --sweep-k 2-8

# 换成在本地数据上自监督训练的特征(SimCLR,Mac 上自动走 MPS)
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml --set method=simclr

簇数由配置的 num_clusters 决定,也可以 --num-clusters 6 临时覆盖(实测 4 类的换种子稳定性最高,见文档 §10.4)。产出在 work_dir: assignments.csv(每张图的簇号与到质心的距离)、summary.json、 clusters/cluster_XX/(按簇建的软链目录,直接用看图软件翻)、 cluster_XX_preview.png(每簇最靠近质心的代表样本拼图)、 cluster_scatter.png(簇分布图,见下)。

簇分布图

cluster_scatter.png 和 preview 拼图同级,一张图回答"各簇离多远":

  • 左:聚类空间的 PCA 二维投影,散点按簇着色,X 是质心。标题给出两个主成分的 方差占比——占比低说明高维结构被压掉得多,图上的远近只能参考。
  • 右:质心两两距离矩阵(在聚类空间里算的,不是投影后),带数值标注。 这个是准的,看重叠严不严重以它为准。

投影用 PCA 而不是 t-SNE/UMAP:后两者为了把簇拉开会扭曲簇间距离,而这张图的目的 恰恰是看距离。相关配置:

scatter:
  enabled: true
  max_points: 5000      # 点太多会糊成一团且文件很大,超过就随机抽样
  point_size: 6
  alpha: 0.35
  figsize: [14, 6]
  dpi: 150

关掉用 --set scatter.enabled=false。图上的中文需要系统里有 PingFang / Noto Sans CJK 等字体,找不到会自动退回英文标签(不会画成豆腐块)。质心距离矩阵同时写进 summary.json 的 centroid_distances,方便脚本读取。

翻图时注意 clusters/ 里文件名前缀是排序号:0000_ 最靠近质心(最典型), 末尾 20 张是离质心最远的离群样本——最能暴露这个簇是不是混了东西。

clusters/ 里放的是硬链接(link_mode: hard,默认):不占额外磁盘,但表现得和真实 文件一样——看图软件能开、可以直接在 runs/ 下把图挪来挪去、建子目录继续分子类, 整个目录拷到别的机器也不会断链。要求与源图同一文件系统,跨盘会自动退回软链并警告。

link_mode 说明
hard(默认) 硬链接。零额外磁盘,行为同真实文件,可跨机器拷贝
sym 软链接。跨盘可用,但记绝对路径,换机器就断(export_clusters.py 能重定位)
copy 直接复制。跨盘可用且绝不断链,代价是占一份磁盘

hard 与原图是同一份数据:在 clusters/ 里改图会改到原图;只挪动、删除条目 不影响原图。

尺寸不统一的处理:这批图是已裁好的缺陷小图,中位 152×113,宽高比从 0.24 到 39.9。 尺寸和宽高比是上游 ROI 裁剪的产物、不是缺陷属性,所以预处理的目标是把它抹掉: 直接拉伸到 128×128(不用 letterbox——填充比例等价于把宽高比明着喂给网络), 再做逐图 z-score 抹掉整体亮度/对比度。

无标签时用两个指标判断质量:轮廓系数(簇是否紧致)和换随机种子的稳定性 (同一份数据换初始化后划分是否还一样,低说明是在硬切一团连续数据)。

完整的数学推导、算法流程图与全部消融数据见 docs/unsupervised-clustering.md。

按聚类结果导出数据集

聚类产出的 clusters/ 里是软链,方便人工翻看和调整——可以把图在子目录间挪动、 删掉分错的、把 cluster_03 改名成 grease。整理完之后,以这个目录为准把真实文件 复制出来,得到一个可以直接使用的数据集:

uv run python scripts/export_clusters.py --work-dir runs/simclr_r101 --prefix-parent
#                                                              先看计划加 --dry-run
data/simclr_r101/          # 默认取 work_dir 的名字,--out 可完全指定
├── grease/                # 子目录名沿用 clusters/ 里的(人工整理过的类别名)
├── mosquito/
├── normal/
└── manifest.json          # 每张图的原始路径与所属类别,便于追溯

导出的是真实文件(不是软链),文件名用原图的名字、去掉 clusters/ 里的 0125_ 排序前缀。加 --prefix-parent 会再拼上原父目录名(反面纵黑点/abc.png -> 反面纵黑点_abc.png),既保住来源信息,也化解不同目录同名文件撞车。

从服务器拷回来的 clusters/ 软链全都是坏的(指向 /home/deep/...),不影响使用: 脚本读出链接目标,按其中的 data/ 段重新落到本机 --source-root(默认当前目录)上; 万一还找不到,再用文件名在源目录里兜底索引一次。解析不到的条目会列出来并中止, 确认要忽略就加 --skip-missing。

几个保护:同类别内重名(保持原名会互相覆盖)直接报错并提示可用 --prefix-parent; 目标目录非空时必须显式 --overwrite。

自监督预训练(SimCLR)

method: pretrained(默认)用 ImageNet 权重直接抽特征,不训练;method: simclr 则在 本批薄膜图上自监督训练一个项目自己的纯卷积骨架,域适配更好。同一张图的两个随机增强 视图互为正样本,batch 内其余 2B−2 个视图为负样本,用 NT-Xent 损失训练。

训练走 device: auto 选中的后端——Mac 上直接用 MPS 即可,实测 14,579 张、 batch 256 时约 65 秒/epoch。

第一步:训练 + 聚类(一条命令)

method=simclr 会先训练、存权重,然后直接用这份权重抽特征做聚类,不需要分两次调用:

uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
    --work-dir runs/simclr \
    --set method=simclr simclr.epochs=40 loader.num_workers=8

按 65 秒/epoch 算,40 轮约 45 分钟。产出:

runs/simclr/
├── simclr_backbone.pth      # 只含骨架,给下游用(聚类 / 初始化 DefectNet)
├── simclr_last.pth          # 完整训练状态,给断点续训用
├── swanlog/                 # SwanLab 日志
├── cluster.log
├── assignments.csv
├── summary.json
├── clusters/cluster_XX/
└── cluster_XX_preview.png

跑长任务建议丢后台,随时能看进度:

nohup uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
    --work-dir runs/simclr \
    --set method=simclr simclr.epochs=40 loader.num_workers=8 > runs/simclr.log 2>&1 &

tail -f runs/simclr.log

断点续训

训练中断(Ctrl-C、kill、机器重启)后接着练,用 simclr.resume 指向 simclr_last.pth:

uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
    --work-dir runs/simclr \
    --set method=simclr simclr.epochs=40 loader.num_workers=8 \
          simclr.resume=runs/simclr/simclr_last.pth

日志会打印"已载入训练状态 …:从 epoch N 接着练到 40"。恢复的不只是权重,还有投影头、 优化器动量缓冲和 step 计数——少了这些,续训等于把动量清零,损失会先反弹一段。

配置文件是唯一真相源。 存档只提供"进度"(权重、动量、已训到第几轮); epochs、lr、weight_decay、调度器一律以当前 yaml 为准。所以续训时可以顺手调超参:

uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
    --work-dir runs/simclr \
    --set method=simclr simclr.epochs=60 simclr.optimizer.lr=5e-4 \
          simclr.resume=runs/simclr/simclr_last.pth

与存档不一致时日志会逐项列出(配置与存档不一致,按当前配置执行:总轮数 40 -> 60; optimizer.lr 0.001 -> 0.0005),不会出现"改了没生效"。学习率调度按新的总轮数重建 并快进到当前 epoch,不会沿用存档里的 T_max。

续训到调度末段时想调大 lr:cosine 是 epoch 的函数,在 2229/2500 处只剩满值的 3%, 所以新设的 lr=1e-3 实际只有 2.9e-5。日志会警告这一点。要让新 lr 真正生效:

# 按剩余轮数重开一条 cosine 曲线,新 lr 从满值开始
--set simclr.restart_schedule=true

# 或者拉长整条曲线(总轮数变大,当前位置的系数就回升)
--set simclr.epochs=5000

train.py 的 train.restart_schedule 同理。

续训时各项配置的行为:

配置项 续训时
epochs 按 yaml,且是总轮数不是剩余轮数。已训 12 轮、写 40 就再训 28 轮;写成 ≤12 则不再训练并给出提示
optimizer 的超参(lr / weight_decay / betas / momentum …) 按 yaml
scheduler(类型与参数) 按 yaml,重建后快进到当前 epoch
temperature、log_interval、save_interval、amp、batch_size、增强 按 yaml
模型权重、投影头、优化器动量缓冲、已训轮数、global_step 来自存档(这些才是"进度")
optimizer.name(adamw ↔ sgd) 可以改,但动量缓冲无法迁移、会从零积累(有日志提示)
backbone 结构、projection 维度 不能改,与存档权重形状不符会直接报错

resume 与 checkpoint 语义不同,二选一(同时给会报错): resume = 载入完整状态接着训练(要 simclr_last.pth); checkpoint = 有权重就跳过训练直接聚类(要 simclr_backbone.pth)。 拿 simclr_backbone.pth 去 resume 会报错,因为它不含投影头和优化器状态。

训练与聚类分开跑

自监督训练和聚类在代码上本来就是解耦的(engine/simclr.py 与 engine/cluster.py 互不引用,唯一接口是 simclr_backbone.pth),只是默认一条命令串起来跑。 需要分两步时——比如在 GPU 机器上专心训练几十小时,之后在别处反复调 k——用 --train-only 断开:

# 第一步:只训练,存完权重就退出,不抽特征也不聚类
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film_simclr_resnet18.yaml \
    --work-dir runs/simclr_r18 --train-only --set simclr.epochs=40

# 第二步:只聚类,想换几次 k 就跑几次,每次约 1 分钟
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film_simclr_resnet18.yaml \
    --work-dir runs/simclr_r18_k4 \
    --set simclr.checkpoint=runs/simclr_r18/simclr_backbone.pth --num-clusters 4

第一步只产出 simclr_backbone.pth 与 simclr_last.pth,没有 assignments.csv 之类。 --train-only 遇到 method=pretrained(没东西可训)、simclr.checkpoint(会跳过训练) 或 --sweep-k(要靠聚类才能扫)都会直接报错,不会跑完才发现什么都没产出。

第二步:复用权重重新聚类

想换簇数或改聚类参数时不必重训——把权重路径填进 simclr.checkpoint 就会跳过训练 (日志会打印"复用已有的自监督骨架"),这一步只跑特征提取 + 聚类,约 1 分钟:

# 换成 k=2(实测 SimCLR 特征在 k=2 上最干净:轮廓 0.469 / 稳定性 0.991)
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
    --work-dir runs/simclr_k2 \
    --set method=simclr simclr.checkpoint=runs/simclr/simclr_backbone.pth \
    --num-clusters 2

# 先扫一遍簇数再决定(只打指标表,不产出分组)
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
    --set method=simclr simclr.checkpoint=runs/simclr/simclr_backbone.pth \
    --sweep-k 2-8

复用权重时给一个新的 --work-dir,否则新的聚类产出会覆盖同目录下的旧结果。

实测效果(14,579 张,与 pretrained 用完全相同的聚类设置):只训 1 个 epoch 的域内特征 在 k=4 上就追平 ImageNet(轮廓 0.334 vs 0.325),在 k=2 上大幅胜出(轮廓 0.469、 稳定性 0.991)——目视确认这个二分正是"有明显缺陷 / 弱纹理近空白",可直接当作流水线的 第一级过滤。详见 文档 §10.5。

权重存到 <work_dir>/simclr_backbone.pth(只含骨架,不含投影头——投影头被训练成 对增强不变,会丢掉下游有用的信息,所以聚类和保存都只用骨架特征)。

固化到配置

每次敲 --set 麻烦的话,直接改 configs/cluster/film.yaml:

method: simclr

simclr:
  # 自监督用项目自己的纯卷积骨架(与 method=pretrained 的顶层 `backbone` 各自独立,
  # 两者字段完全不同,分开配置以免 _base_ 合并时串味)
  backbone: {name: tiny_conv, in_channels: 3, widths: [32, 64, 128, 256], depths: [1, 2, 2, 2]}
  epochs: 40
  batch_size: 256
  temperature: 0.2          # 越小越关注难负样本
  amp: auto                 # CUDA 上开、MPS 上关
  save_interval: 5          # 每 5 个 epoch 落一次盘
  # 下面两个二选一,语义不同:
  checkpoint: null          # 有权重就跳过训练直接聚类,例 runs/simclr/simclr_backbone.pth
  resume: null              # 载入完整状态接着训练,例 runs/simclr/simclr_last.pth
  projection: {hidden_channels: 256, out_channels: 128}
  optimizer: {name: adamw, lr: 0.001, weight_decay: 0.0001}
  scheduler: {name: cosine}

view_transforms:            # 生成两个视图的随机增强
  - {name: RandomResizedCrop, size: [128, 128], scale: [0.5, 1.0]}
  - {name: RandomHorizontalFlip, p: 0.5}
  - {name: RandomVerticalFlip, p: 0.5}
  - {name: GaussianBlur, p: 0.3}

RandomResizedCrop 的 scale 下限别低于 0.4:裁太狠会让两个视图都只剩背景, "背景 vs 背景"被当成正样本对,学出来的特征会退化。

换一个骨架

仓库自带 resnet18 / resnet50 / resnet101 的现成配置,直接用:

uv run python scripts/cluster_unlabeled.py --config configs/cluster/film_simclr_resnet18.yaml \
    --work-dir runs/simclr_r18 --set simclr.epochs=40

uv run python scripts/cluster_unlabeled.py --config configs/cluster/film_simclr_resnet50.yaml \
    --work-dir runs/simclr_r50 --set simclr.epochs=40

uv run python scripts/cluster_unlabeled.py --config configs/cluster/film_simclr_resnet101.yaml \
    --work-dir runs/simclr_r101 --set simclr.epochs=40

两者在 Apple MPS 上的实测开销(14,579 张,128×128 双视图):

骨架 参数量 layer2+3 维度 每 epoch batch 上限
resnet18 11.2M 384 约 93 秒 1024 可用
resnet50 23.5M 1536 约 155 秒 1024 会 OOM(需 18.7 GiB),配置里用 256
resnet101 42.5M 1536 未实测 配置里用 128(按规模比例保守取值)

resnet50 每 epoch 的耗时与 batch 基本无关(算力受限),所以显存不够时往下调 batch 不花额外时间, 只是 NT-Xent 的负样本变少。

注意默认的 method: pretrained 用的也是 resnet18(ImageNet 权重,不训练,约 51 秒出结果); 上面这个配置是在 data/raw/film 上自监督训练它。两者实测对比(14,579 张,k=4):

方案 轮廓 稳定性 成本
pretrained resnet18(ImageNet,不训练) 0.326 0.964 约 51 秒
SimCLR tiny_conv(1 epoch) 0.334 0.873 65 秒/epoch
SimCLR resnet18(1 epoch) 0.243 0.997 93 秒/epoch,显存 315 MB

resnet18 参数量是 tiny_conv 的 37 倍(11.2M vs 0.3M),1 个 epoch 远不够收敛—— 上表只是下界,要用它就得训几十轮。想快速出结果,默认的 pretrained 或 tiny_conv 更划算。

simclr.backbone 用的是和检测模型同一套注册表,可选:

name arch out_channels 说明
tiny_conv — 256(可调) 自研深度可分离卷积,默认,最轻
torchvision resnet18 / resnet34 512 通用基线
torchvision resnet50 2048 更强也更重,MPS 上注意 batch
torchvision resnet101 2048 最深,首次使用会下载约 171 MB 权重
torchvision mobilenet_v3_small / _large 576 / 960 面向实时部署
torchvision shufflenet_v2_x1_0 1024 同上

全是纯卷积(ViT/Swin 不在支持范围,本项目要部署在实时设备上)。改法是整块替换 simclr.backbone,因为两种骨架的字段完全不同(widths/depths/block vs arch/pretrained):

# configs/cluster/film.yaml
simclr:
  backbone:
    name: torchvision
    arch: resnet18
    pretrained: true      # 先用 ImageNet 权重起步,再在本批数据上自监督微调
    in_channels: 3

如果不想改原文件,用 _base_ 继承时要加 _replace_: true——子字典默认是合并而不是 替换,不加的话 base 里 tiny_conv 的 widths/depths 会残留下来传给新骨架:

_base_: [configs/cluster/film.yaml]
simclr:
  backbone:
    _replace_: true     # 整块换掉,不与 base 合并
    name: torchvision
    arch: resnet18
    pretrained: true

不要用 --set simclr.backbone.name=torchvision:它只改名字,旧骨架的专有字段仍在, 会报 构建 'torchvision' 时收到它不接受的字段 ['block', 'depths', 'widths']。 换组件请改配置文件。

注意 simclr.feature_stages:torchvision 系列的 features 是 nn.Sequential,可以填 [5, 6] 这类中间层下标;tiny_conv 不是,必须保持 null(取最终输出)。

SimCLR 训练同样接了 SwanLab(配置在 simclr.swanlab,字段与 train.swanlab 相同, 默认 mode: local 不需要登录)。训练跑起来后另开一个终端:

uv run swanlab watch runs/simclr/swanlog

记录的指标:按 step 的 simclr/loss、simclr/loss_avg、simclr/lr,按 epoch 的 simclr/epoch_loss、simclr/epoch_seconds、simclr/gpu_mem_peak_mb,超参是完整的 yaml 配置。 不想要就 --set simclr.swanlab.enabled=false;纯看文本日志用 grep "完成 loss" runs/simclr/cluster.log。

训练出的骨架和检测模型用的是同一套卷积骨架,可以直接拿去初始化 DefectNet 的 backbone—— 标注量有限时,通常比 ImageNet 权重更贴合薄膜的成像特性。

缺陷分类模型

在聚类 + 人工归类得到的数据集上训练分类模型(data/<数据集>/<类别名>/*.png):

uv run python classify.py --config configs/classify/film.yaml

# 只评估
uv run python classify.py --config configs/classify/film.yaml --eval-only \
    --checkpoint runs/classify/best.pth

产物在 work_dir:best.pth / last.pth / config.yaml / classify.log / eval_report.json(逐类 P/R/F1 与完整混淆矩阵)。

三个关键取舍

主指标是 macro-F1,不是准确率。 数据从 normal 2679 张到 others 65 张差 41 倍, 全预测成 normal 也能拿 24% 准确率。配置里同时给出 balanced_accuracy(各类召回的均值), 两者一起看才知道小类有没有学会。

损失按类别频率加权(class_weight: balanced,逆频率归一化到均值 1), 否则占 24% 的 normal 会主导梯度。不需要时设 none。

逐图标准化抹掉整体亮度。 15 个类别里有几类几乎只来自单一原始数据集 (ripple 几乎全是 pearlescent,grease/zigzag 全是 film),不抹掉亮度的话模型 可能靠"认出这是哪个数据集"就把类别猜对——验证集上分数很高,换批数据就崩。

划分

没有预先的 train/val,按类别分层切分(默认 8:2,固定种子)。逐类切分是必须的: others 只有 65 张,整体随机切很容易让它在验证集里一张不剩。

用自监督骨架初始化

前面自监督训练出的骨架可以直接拿来初始化——这正是聚类阶段产物的下游用途:

model:
  backbone: {name: torchvision, arch: resnet101}   # 必须与自监督时一致
  init_from: runs/simclr_r101/simclr_backbone.pth
  frozen_backbone: false    # true = 冻结骨架只训分类头(线性探针)

只加载 backbone,分类头仍随机初始化。frozen_backbone: true 可以单独衡量 自监督特征本身有多好。

从分类到定位:生成框标签

分类模型只学过"这是什么缺陷",没学过"缺陷在哪"。归因方法能把分类分数反推回空间位置, 阈值化 + 取连通域外接矩形,就得到弱监督的框标签。另外还有一条不依赖模型的传统 CV 路线:

# Grad-CAM(默认,全量 11k 张约 2 分钟)
uv run python scripts/generate_boxes.py --out runs/boxes_gradcam --method gradcam

# SHAP Partition(更精细,全量约 1.7 小时)
uv run python scripts/generate_boxes.py --out runs/boxes_shap --method shap

# 传统 CV 背景差分(最快,且不需要分类模型)
uv run python scripts/generate_boxes.py --out runs/boxes_cv --method cv

默认对全量数据生成;--limit / --per-class 只用于快速试参数。

产出是自成一体的 YOLO 数据集

runs/<name>/
├── images/<类别>/<图名>.png    # 硬链接到原图,不占额外磁盘
├── labels/<类别>/<图名>.txt    # YOLO:class_id cx cy w h(归一化)
├── heatmaps/<类别>/<图名>.png  # 热力图叠加原图 + 框,供人工核对
├── data.yaml
└── summary.json

images/ 与 labels/ 的对应关系符合本仓库的 YOLO 约定,所以 FilmDefectDataset 能直接读——这一步把分类的产出接回了检测训练链路。

热力图与框的类别默认取图片所在目录的人工标注(--target label)。类别的真值是人工归类, 分类模型只用来定位;改用 --target predicted(模型 top1 预测)时,模型分错的图框的类别也会跟着错。 两者只在预测与目录不一致的图上有区别——boxes_shap 那次是 0.85%,约 95 张。 summary.json 里的 prediction_matches_folder_label 记录两者的一致率(实测 0.987)。

三种方法的取舍

gradcam shap cv
依赖分类模型 是 是 否
分辨率 受限于特征图(layer3 @128² 只有 8×8) 与输入同分辨率(128×128) 原图分辨率
单张耗时 ~0.01 秒 ~0.55 秒 ~0.011 秒
小目标 几像素的黑点在 8×8 上不可分辨 能贴合到几十像素的目标 贴合好
找的是什么 模型据以判断类别的区域 同左,更精细 偏离局部背景的区域

实测框面积占比(越小越贴合缺陷本体):blackdot 0.019(shap) vs 0.269(gradcam)、 mosquito 0.030 vs 0.173、whitedot 0.076 vs 0.253。小目标类别上 SHAP 明显更准, 代价是慢 50 倍。SHAP 需要 uv sync --group explain。

传统 CV:背景差分(--method cv)

实现在 engine/classical.py:大半径方框均值估背景 → 残差(可选只看 比背景暗/亮的一侧)→ 小半径平滑 → 判定有无缺陷 → 按峰值归一化,最后交给和另外两种方法 同一套 cam_to_boxes 出框。只依赖 numpy。

它与归因方法有三点本质差别:

  • 不需要分类模型,--checkpoint 都不用给,分类器还没训练出来就能先跑一版框; 也因此不受分类精度影响。
  • 类别只能来自目录(--target label),它自己不判类。
  • 找的是"看起来和周围不一样"的区域,不是"模型据以判断类别的区域"。对 blackdot / whitedot / mosquito 这类局部目标效果好;ripple / wrinkle 这种铺满整图的 纹理类本来就没有局部目标,它给不出有意义的框。

有无缺陷是显式判定的,这点和归因方法不同。残差峰值要同时高于绝对对比度下限 (--cv-min-contrast,默认 0.05≈13 个灰阶)和噪声水平的 15 倍,否则直接返回全零图、出 0 个框。 噪声水平用残差的 MAD 估计。这条判定是必需的:实测纯噪声图的"残差峰值 / MAD"稳定在 9 附近 (最大 10.9,几乎与噪声大小无关),而有一个对比度 0.25 的黑点时是 25~330,两者分得很开; 若只按峰值归一化,一张正常图的噪声会被拉到 1,满图假框。所以正常图天然得到空标签。

调参:--cv-radius(背景半径,0=按图取 min(高,宽)//8;太小会把缺陷自身算进背景)、 --cv-polarity(dark 黑点蚊虫 / bright 白点反光 / both)、--cv-min-contrast、 --smooth-radius(与 shap 共用)、--cv-max-side(分析前缩边长,默认 256——框是归一化的, 缩放不影响标签,但能挡住纯 Python 的连通域标记在大图上变慢)。

已知局限。 背景估计在图像边界处会对强烈的明暗阶跃产生响应,所以贴着边框的 明暗交界可能被框出来。对 edge_clean / edge_dirty 这种本来就在拍边缘的类别,这是真实信号; 对其他类别则是假框,复核时留意。

调参

  • --max-boxes(默认 1):一张裁片通常只对应一处缺陷,多出来的框多半是噪声
  • --threshold(默认 0.4):热力图相对最大值的阈值
  • --min-area-ratio:缺省随方法而变(gradcam 0.02、shap 0.002)。两者归因图分辨率 差 16 倍,同一个小黑点在 128×128 上只占百来像素,沿用 0.02(=328 像素)会被整个滤掉
  • --peak-ratio(默认 0.75):次要连通域的峰值需达全局最大值的这个比例,滤掉弱响应假框
  • --padding(默认 0.1):框四周外扩。归因倾向只覆盖最具判别性的部位,往往偏小
  • --stage(gradcam,默认 6=layer3):layer4 只有 4×4 太粗,layer2 响应太散
  • --no-heatmaps / --no-images:只要标签时可跳过

不是所有类别都可定位。 normal/black/edge_clean 描述的是整图属性而非局部缺陷, 它们的框没有物理意义。真正有用的是 blackdot/whitedot/mosquito/grease 这些 有局部目标的类别。

检测模型训练(YOLO)

generate_boxes.py 产出的是弱监督框,要先在 Label Studio 里人工修正,再划分 train/val 训练:

# 1) 生成 Label Studio 复核材料到 runs/boxes_shap_label_studio/
uv run python scripts/prepare_label_studio.py --source runs/boxes_shap
#                                             先看计划加 --dry-run

# 2) 在 Label Studio 里修正标签(导入步骤见下文,脚本结束时也会打印)

# 3) 导出修好的标签,划分 train/val 到 data/boxes_shap/(导出脚本待写)

# 4) 训练
uv run python detect.py --config configs/detect/film.yaml

# 只评估 / 续训
uv run python detect.py --config configs/detect/film.yaml --eval-only \
    --checkpoint runs/detect/film/weights/best.pt
uv run python detect.py --config configs/detect/film.yaml --resume

准备 Label Studio 复核材料

runs/boxes_shap_label_studio/    # 同时也是 Label Studio 的 DOCUMENT_ROOT
├── images/<类别>_<序号>.png     # 一律 PNG:PNG 原图硬链接,BMP 等无损转换
├── labels/<类别>_<序号>.txt
├── classes.txt                  # 每行一个类别名,行号即 class_id
├── name_map.csv                 # 新文件名 ↔ 源文件路径
├── summary.json                 # 逐类数量、空标签处置、落盘方式
├── tasks.json                   # 每张图一个任务,框作为 predictions
└── tasks.label_config.xml       # 各类别的 RectangleLabels

不写 data.yaml、不划分 train/val。 这是给人复核的材料,不是训练集,没有 data.yaml 也就不会被误拿去训练。复核期间 Label Studio 直接从这个目录读图,不要删除或 --overwrite 它。 detect.py 若检测到 data.yaml 的 train 与 val 相同(generate_boxes.py 的产出就是这样), 也会告警:验证指标算的是训练集本身,mAP 虚高,早停与 best.pt 的选择失真。

图片统一成 PNG。 pearlescent / blue / uv 的原始数据是 BMP(共 10,814 张), generate_boxes.py 按原后缀链接,于是产出里混着两种格式。这一步按文件头 (不是后缀)判断:真 PNG 照旧硬链接不占磁盘,其余无损转 PNG——不缩放像素,归一化的框坐标 原样有效;CMYK 这类 PNG 存不下的模式转 RGB。转换无法硬链接,会占额外磁盘,--dry-run 会先报出要转多少张。统一后缀后 a.png 与另一目录的 a.bmp 会撞名,脚本直接报错。

按类别重新编号。 源文件名五花八门(反面纵黑点_1_550mmX1561.71m_20260416160441_0.png、 pearlescent_404.bmp……),输出统一改成 <类别>_<序号>.png,如 blackdot_0042.png。 类别以标签第一列为准,文件名永远不会与标签矛盾;序号在类内按源路径排序,同一份来源重跑, 同一张图拿到的名字不变。原文件名里的正反面、位置、时间戳等信息都在 name_map.csv 里 (带 BOM,Excel 直接打开中文不乱码)。想保留原名加 --keep-names。

classes.txt 是 label-studio-converter 读类别用的(也兼容 LabelImg / X-AnyLabeling), 行号就是 class_id,内容与来源 data.yaml 严格一致。

导入 Label Studio

图片 URL 必须按 Label Studio 实际的 DOCUMENT_ROOT 生成。 任务里写的是 /data/local-files/?d=<相对路径>/images/<图名>.png,Label Studio 把 d 拼到 LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT 后面找文件,相对路径错了就是整个项目 404。 脚本缺省读这个环境变量,也可以用 --document-root 显式指定:

# DOCUMENT_ROOT 统一放标注项目,产出仍留在 runs/
uv run python scripts/prepare_label_studio.py --source runs/boxes_shap \
    --document-root /home/deep/labeling
情形 做法 图片 URL
没设 DOCUMENT_ROOT 输出目录自己当 DOCUMENT_ROOT ?d=images/…
输出目录在 DOCUMENT_ROOT 之下 直接用 ?d=<相对路径>/images/…
输出目录在 DOCUMENT_ROOT 之外 在 DOCUMENT_ROOT 里建同名软链接指向它 ?d=boxes_shap_label_studio/images/…

DOCUMENT_ROOT 里已有同名的真目录(比如手动复制过去的一份)时脚本会报错,不会覆盖它。 Local files 存储的绝对路径填 <DOCUMENT_ROOT>/<访问路径>/images,脚本打印的步骤里已经算好。

脚本结束时会打印带实际路径的导入步骤:设置 LABEL_STUDIO_LOCAL_FILES_* 环境变量启动、 粘贴 labeling config、添加 Local files 存储(只保存不要 Sync,否则每张图多出一个空任务)、 导入 tasks.json。

几个取舍:

  • 为什么是扁平结构:label-studio-converter import yolo 只认 <input>/{images,labels,classes.txt}, 且 os.listdir 不递归,看不见子目录。
  • 软链接可以用:Label Studio(1.23.0 核对过源码)拼文件路径用 safe_join(abspath), 存储路径做 normpath 后按字符串前缀比较,都不解析软链接,所以 DOCUMENT_ROOT 里放一个 指向 runs/ 的软链接就能访问,不必复制。存储路径也必须是 DOCUMENT_ROOT 的子目录, 与 DOCUMENT_ROOT 相同会被拒绝,这正是要填到 …/images 的原因。
  • 框作为 predictions 而不是 annotations:它们是模型生成的。审核员修改提交后才成为 annotation, 哪些审过一目了然,导出时也只导出 annotation,没审过的框不会混进去。
  • 不传 --image-dims:图片尺寸有 3,223 种,传了会把同一个宽高写进每条标注。
  • converter 自带的导入提示漏了环境变量(它在循环里改写了 image_root_url,与默认值比较失败), 所以它的输出被静音,改由脚本打印。

复核重点是 normal。 SHAP 在每张图上都框了东西,normal 有 4,295 个框、占 38%。 若 normal 表示无缺陷,这些框要删掉,否则检测器会学着"检测正常"。

没有框的图怎么处理

Grad-CAM 没框出东西的有 339 张(占 3%)。它们默认保留为负样本,对压低误检有用。

但要当心:空标签在 YOLO 里的语义是"这张图什么都没有"——一个明确的负样本断言, 而不是"有缺陷但没能定位出来"。而这 339 张里只有 69 张真的来自 normal,其余 270 张是人工 标成缺陷、Grad-CAM 却失手的图:

来源类别 空标签 占该类
grease 95 7.8%
blackdot 82 11.3%
normal 69 2.6%
strip_h 28 2.9%
zigzag 23 2.4%
black 17 24.3%
其余 9 类 25 <3%

照搬的话,blackdot 有 11% 的正样本、black 有 24% 会以"此图无缺陷"的身份进入训练, 等于在教模型抑制真实缺陷。用 --empty-policy 按来源类别分别处置:

uv run python scripts/prepare_label_studio.py --source runs/boxes_gradcam \
    --empty-policy drop normal=keep black=full edge_clean=full
  • keep(默认)——当负样本保留。只有真正无缺陷的类别该用
  • drop——丢弃。缺陷类没定位出来的属于信息缺失,不该冒充负样本
  • full——退化成整图框。适合 black/edge_clean 这种描述整图属性、缺陷本就铺满全图的类别

第一个裸词是全局默认,后面 类别=策略 逐类覆盖。类别名或策略写错会直接报错而不是静默忽略。 没有显式指定策略、又落到 keep 的类别会打一条告警,提醒你确认它是否真的无缺陷。

拍平之后不同类别目录下的同名文件会互相覆盖,脚本会直接报错而不是静默丢数据; 真遇上就用 scripts/export_clusters.py --prefix-parent 重新导出。

模型与增强的取舍

  • 默认 YOLO26 的 n 档:CLAUDE.md 的"只用 convolution 模型"只约束聚类与数据集生成阶段, 检测这一步不受限,所以用最新的 YOLO26——它是端到端的(NMS-free),推理侧少一道后处理。 选最小的 n:数据只有 1.1 万张且标签是弱监督的,模型再大容易把标注噪声一起学进去。
  • 关掉 mosaic / mixup:它们把四张图拼一起,而我们每张图只有一个弱监督框,拼接后小目标被 缩得更小,对 blackdot/whitedot 这类本就只占千分之几面积的类别是净伤害。
  • 关掉 hsv_h:薄膜图近乎单色,色相扰动没有意义。翻转照旧开着(缺陷没有固定朝向)。
  • 冒烟测试用 --set train.epochs=1 train.fraction=0.02 train.imgsz=320 先把链路跑通。

指标怎么读。 框来自 Grad-CAM/SHAP 的弱监督标注,不是人工真值。mAP 衡量的是"检测器 能不能复现归因热力区域",不等于"缺陷定位有多准"。真值上的精度仍需人工抽检。

可视化

train.swanlab.enabled: true 时接 SwanLab 回调,训练曲线实时可看(swanlab watch 开本地看板)。 ultralytics 自己还会在 runs/detect/film/ 下写:

  • results.png / results.csv——各项 loss 与 mAP 曲线
  • confusion_matrix.png、confusion_matrix_normalized.png
  • BoxPR_curve.png / BoxP_curve.png / BoxR_curve.png / BoxF1_curve.png
  • train_batch*.jpg、val_batch*_labels.jpg 与 val_batch*_pred.jpg——标注与预测的直观对照
  • metrics.json——mAP50-95 / mAP50 / P / R 与逐类 mAP50

许可提醒。 ultralytics 是 AGPL-3.0。模型若要随产品分发,需要开源整套代码或另购商业授权; 仅内部使用不受影响。

方法

  • 骨架:tiny_conv(自研深度可分离卷积,默认)或 torchvision(resnet / mobilenet / shufflenet)。
  • neck:nearest 上采样 + 3×3 卷积,把特征恢复到 stride 4,兼顾小缺陷定位与速度。
  • 定位头:anchor-free 中心点(CenterNet 风格)——heatmap + wh + offset, 推理端用 3×3 max-pool 做峰值抑制,无 NMS 循环,便于导出到 TensorRT / OpenVINO。
  • 分类头:全局池化 + 全连接,输出多热标签(一张图可同时含多种缺陷,全 0 表示 OK 片)。
  • 指标:定位 precision / recall / F1 / mean IoU(IoU 阈值匹配),分类 micro-F1 与整图完全正确率。 train.monitor 决定用哪个指标挑选 best.pth。

配置支持 _base_ 继承与命令行 --set a.b.c=value 覆盖,见 utils/config.py。 新增骨架 / 数据集只需用注册表装饰器注册,再在 yaml 里按 name 引用。

开发

uv run pytest         # 单元测试
uv run ruff check .   # 静态检查
S
Description
No description provided
Readme
4.1 MiB
Languages
Python 93.8%
C++ 3%
Shell 1.6%
Mermaid 1.5%