工业视觉实战：从零构建PyTorch+UNet钢材缺陷检测系统

钢铁生产线上，一块钢板以每秒数米的速度通过检测区域。传统质检员需要在高噪音、强光环境下连续工作8小时，肉眼识别发丝粗细的划痕——这种场景正在被AI视觉技术改写。本文将手把手带您实现一个端到端的钢材表面缺陷检测系统，从数据清洗、模型训练到可视化界面开发，完整覆盖工业AI项目的全生命周期。

1. 环境配置与数据准备

工欲善其事，必先利其器。我们先搭建一个可复现的Python环境：

bash复制conda create -n steel_defect python=3.8
conda activate steel_defect
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
pip install albumentations==1.2.1 pytorch-lightning==1.7.7 opencv-python==4.6.0.66

Kaggle钢材数据集包含1.5万张带标注的钢板图像，缺陷类型包括：

数据预处理时需要特别注意三个工业视觉特有的问题：

1. 非均衡样本处理：采用oversampling+undersampling组合策略
2. 小目标增强：对缺陷区域进行随机放大裁剪
3. 环境干扰消除：应用CLAHE算法增强对比度

python复制def apply_clahe(image):
    lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
    cl = clahe.apply(l)
    limg = cv2.merge((cl,a,b))
    return cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)

2. UNet模型深度优化

原始UNet在工业场景下需要针对性改进，我们采用ResNet34作为编码器 backbone，并引入三个关键改进：

1. 深度监督机制：在解码器各层添加辅助损失
2. 空间注意力模块：增强缺陷区域特征响应
3. 混合精度训练：减少显存占用30%

python复制class AttentionBlock(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.query = nn.Conv2d(in_channels, in_channels//8, 1)
        self.key = nn.Conv2d(in_channels, in_channels//8, 1)
        self.value = nn.Conv2d(in_channels, in_channels, 1)
        self.gamma = nn.Parameter(torch.zeros(1))

    def forward(self, x):
        B, C, H, W = x.shape
        q = self.query(x).view(B, -1, H*W).permute(0,2,1)
        k = self.key(x).view(B, -1, H*W)
        v = self.value(x).view(B, -1, H*W)
        
        attn = torch.bmm(q, k)
        attn = F.softmax(attn, dim=-1)
        
        out = torch.bmm(v, attn.permute(0,2,1))
        out = out.view(B, C, H, W)
        return self.gamma * out + x

训练策略上采用分阶段优化：

1. 冻结编码器阶段：只训练解码器头部（10 epochs）
2. 全网络微调阶段：使用余弦退火学习率（50 epochs）
3. 高分辨率精调阶段：输入尺寸提升至1024×1024（5 epochs）

注意：当batch_size=4时，建议使用梯度累积步骤设置为8，可模拟32的batch效果而不爆显存

3. 工业级部署技巧

模型部署到产线需要考虑三个现实约束：

• 实时性要求：单图推理时间<200ms
• 硬件限制：无GPU的工控机环境
• 持续学习：产线新增数据的自动迭代

我们使用TorchScript将模型转换为可脱离Python环境运行的格式：

python复制model = load_trained_model()
scripted_model = torch.jit.script(model)
scripted_model.save("defect_detection.pt")

针对不同硬件平台的优化策略：

4. PyQt5可视化系统开发

一个易用的检测系统需要包含以下功能模块：

1. 实时检测面板：支持摄像头/视频流输入
2. 历史记录查询：按缺陷类型分类检索
3. 报警阈值配置：动态调整灵敏度参数

python复制class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.model = torch.jit.load('defect_detection.pt')
        self.init_ui()
        
    def init_ui(self):
        self.setWindowTitle("钢材缺陷检测系统 v1.0")
        self.image_label = QLabel()
        self.result_label = QLabel()
        
        open_btn = QPushButton("打开图像")
        open_btn.clicked.connect(self.open_image)
        
        layout = QHBoxLayout()
        layout.addWidget(self.image_label)
        layout.addWidget(self.result_label)
        
        main_layout = QVBoxLayout()
        main_layout.addLayout(layout)
        main_layout.addWidget(open_btn)
        
        container = QWidget()
        container.setLayout(main_layout)
        self.setCentralWidget(container)

实际部署时遇到的两个典型问题及解决方案：

1. 金属反光干扰：在摄像头前加装偏振滤镜
2. 运动模糊：与PLC同步触发拍摄时机

5. 持续改进与模型迭代

建立有效的反馈闭环是工业AI项目的关键。我们设计了一套自动化迭代流程：

1. 在线数据标注：质检员修正错误预测结果
2. 增量训练：每周用新数据微调模型
3. A/B测试：新旧模型并行运行对比

模型性能监控指标应包括：

• 假阳性率：误检导致的停产损失
• 召回率：漏检带来的质量风险
• 推理延迟：产线节拍匹配度

python复制def evaluate_model(test_loader):
    model.eval()
    stats = {'TP':0, 'FP':0, 'FN':0}
    
    with torch.no_grad():
        for images, masks in test_loader:
            outputs = model(images)
            preds = (outputs.sigmoid() > 0.5).float()
            
            stats['TP'] += (preds*masks).sum()
            stats['FP'] += (preds*(1-masks)).sum()
            stats['FN'] += ((1-preds)*masks).sum()
    
    precision = stats['TP']/(stats['TP']+stats['FP']+1e-7)
    recall = stats['TP']/(stats['TP']+stats['FN']+1e-7)
    return {'precision': precision.item(), 'recall': recall.item()}

在某个实际产线部署案例中，经过3个月的持续迭代，系统将缺陷检出率从人工检测的82%提升到95.6%，同时将每个钢卷的检测时间从90秒缩短到9秒。