---
title: Swin-UNet-Defect-Segmentation
canonical_url: "https://www.modelscope.cn/models/Brilliantccc/Swin-UNet-Defect-Segmentation"
md_url: "https://www.modelscope.cn/models/Brilliantccc/Swin-UNet-Defect-Segmentation.md"
repository: Brilliantccc/Swin-UNet-Defect-Segmentation
last_updated: 2026-09-10
pipeline_tag: semantic-segmentation
tasks:
  - semantic-segmentation
library_name:
  - pytorch
downloads: 0
stars: 0
---

# Swin-UNet-Defect-Segmentation

> Swin-UNet-Defect-Segmentation - Brilliantccc 在 ModelScope 开源的模型。Swin-UNet-Tiny 工业缺陷分割模型

- **Repository**: Brilliantccc/Swin-UNet-Defect-Segmentation
- **Tasks**: semantic-segmentation
- **Downloads**: 0
- **Stars**: 0
- **Last updated**: 2026-09-10

Source: https://www.modelscope.cn/models/Brilliantccc/Swin-UNet-Defect-Segmentation

---

# 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 |

## 使用方法

### 推理

```python
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)  # 缺陷概率图
```

### 颜色差异通道计算

```python
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)
```

### 训练

```bash
# 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 分辨率对小缺陷可能不够
- 颜色差异通道在不同数据集间分布不一致