Swin-UNet-Tiny 工业缺陷分割模型
基于 Swin Transformer 的 U-Net 语义分割模型，用于工业缺陷检测。
模型架构
Encoder (Swin Transformer + Patch Merging)
Patch Embed (4×4) → 56×56×96
Stage 1: Swin Blocks × 2 → 56×56×96 ──────┐
PatchMerging → 28×28×192 │
Stage 2: Swin Blocks × 2 → 28×28×192 ────┐ │
PatchMerging → 14×14×384 │ │
Stage 3: Swin Blocks × 6 → 14×14×384 ──┐ │ │
PatchMerging → 7×7×768 │ │ │
Stage 4 (Bottleneck): 7×7×768 │ │ │
│ │ │
Decoder (Patch Expanding + Swin Blocks) │ │ │
Expand → 14×14×384 + skip ─────────────┘ │ │
Expand → 28×28×192 + cat + proj ─────────┘ │
Expand → 56×56×96 + cat + proj ──────────────┘
Final Expand → 112×112×48
Bilinear → 224×224 → Conv Head → 分割图
训练配置
| 配置项 |
值 |
| 模型 |
Swin-UNet-Tiny |
| 输入 |
224×224×4 (RGB + 颜色差异通道) |
| 输出 |
224×224×1 (缺陷概率图) |
| 预训练 |
ImageNet-1K (swin_tiny_patch4_window7_224) |
| 数据集 |
MVTec AD + VisA (混合训练) |
| 损失函数 |
Focal Loss + Dice Loss |
| 优化器 |
AdamW (lr=1.5e-4, weight_decay=5e-4) |
| 调度器 |
Cosine Annealing + Warmup (5 epochs) |
| Batch Size |
64 |
| Epochs |
200 (早停 patience=30, 实际 137 epoch) |
| 混合精度 |
AMP |
| 梯度裁剪 |
1.0 |
评估结果
MVTec AD 测试集 (723 张)
| 指标 |
值 |
| Dice (缺陷) |
0.5476 |
| IoU |
0.4621 |
| Pixel AP |
0.8253 |
| PR-AUC |
0.9500 |
| Precision |
0.7603 |
| Recall |
0.5385 |
| Normal Accuracy |
0.9302 |
VisA 测试集 (9861 张)
| 指标 |
值 |
| Dice (缺陷) |
0.0160 |
| IoU |
0.0122 |
| Pixel AP |
0.2833 |
| PR-AUC |
0.8979 |
| Normal Accuracy |
0.9717 |
注：VisA 上效果较差，可能原因：(1) VisA 缺陷形态更多样（PCB、管道等），224×224 分辨率丢失细节；(2) 训练时正常样本远多于缺陷，模型偏向预测正常；(3) VisA 的颜色差异通道分布与 MVTec 不一致。
MVTec 各类别详细结果
| 类别 |
Dice |
IoU |
DefDice |
Precision |
Recall |
| bottle |
0.2723 |
0.2296 |
0.6912 |
0.3645 |
0.2459 |
| cable |
0.1452 |
0.1249 |
0.5885 |
0.2081 |
0.1346 |
| capsule |
0.0861 |
0.0543 |
0.1762 |
0.3313 |
0.0599 |
| carpet |
0.2494 |
0.1923 |
0.6372 |
0.2346 |
0.2984 |
| grid |
0.1400 |
0.0991 |
0.3851 |
0.2669 |
0.1102 |
| hazelnut |
0.2196 |
0.1922 |
0.8471 |
0.2285 |
0.2179 |
| leather |
0.2683 |
0.2140 |
0.7201 |
0.2775 |
0.2861 |
| metal_nut |
0.3767 |
0.3319 |
0.8129 |
0.4013 |
0.3589 |
| pill |
0.4086 |
0.3703 |
0.7749 |
0.4249 |
0.4098 |
| screw |
0.0660 |
0.0486 |
0.1787 |
0.1223 |
0.0532 |
| tile |
0.2453 |
0.2056 |
0.7214 |
0.3131 |
0.2097 |
| toothbrush |
0.2034 |
0.1629 |
0.6103 |
0.2008 |
0.2245 |
| transistor |
0.0511 |
0.0447 |
0.4341 |
0.0642 |
0.0475 |
| wood |
0.3192 |
0.2732 |
0.8247 |
0.3063 |
0.3398 |
| zipper |
0.3561 |
0.3072 |
0.8308 |
0.3604 |
0.3606 |
模型参数
| 项目 |
值 |
| 总参数量 |
~41M |
| 编码器 |
Swin Transformer Tiny (预训练) |
| 解码器 |
Patch Expanding + Swin Blocks |
| 输入通道 |
4 (RGB + 颜色差异) |
| 输入尺寸 |
224×224 |
使用方法
推理
import torch
from models.swin_unet import swin_unet_tiny
# 加载模型
model = swin_unet_tiny(image_size=224, n_classes=1, in_chans=4)
checkpoint = torch.load('best.pth', map_location='cpu')
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()
# 推理
x = torch.randn(1, 4, 224, 224) # RGB + 颜色差异通道
with torch.no_grad():
pred = model(x) # (1, 1, 224, 224)
prob = torch.sigmoid(pred) # 缺陷概率图
颜色差异通道计算
import cv2
import numpy as np
def compute_color_diff(image_rgb):
"""计算颜色差异通道"""
lab = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2LAB).astype(np.float32)
mean_lab = lab.mean(axis=(0, 1))
diff = np.sqrt(((lab - mean_lab) ** 2).sum(axis=2))
return (diff / (diff.max() + 1e-7)).astype(np.float32)
训练
# MVTec + VisA 混合训练
bash train_unet_autodl.sh swin_unet
# 单独评估
python evaluate_unet.py \
--weights best.pth \
--dataset mvtec \
--data-root data/mvtec_anomaly_detection \
--img-size 224 224 \
--visualize
项目代码
模型实现在 models/swin_unet.py，训练脚本为 train_unet.py，评估脚本为 evaluate_unet.py。
环境依赖
- Python 3.8+
- PyTorch 2.0+
- CUDA 11.8+
局限性
- VisA 数据集上效果较差，需要进一步优化
- 224×224 分辨率对小缺陷可能不够
- 颜色差异通道在不同数据集间分布不一致