从数据到部署：基于PyTorch与UNet的钢材缺陷检测实战

1. 为什么钢材缺陷检测需要AI在钢铁生产线上肉眼检查表面缺陷就像在高速公路上找蚂蚁——不仅效率低下而且容易出错。我曾在某钢厂见过质检员每天要检查上千张钢板下班时眼睛布满血丝的场景。传统机器视觉方法虽然能检测明显缺陷但对于细微裂纹或复杂纹理背景下的缺陷误检率常常超过30%。UNet网络就像给计算机装上了工业显微镜人脑判断的组合能力。它的编码器部分像经验丰富的老师傅能快速抓住缺陷的典型特征解码器部分则像精确的绘图师能把缺陷轮廓像素级还原。实测下来这套方案在Kaggle钢材数据集上能达到89%的Dice系数比传统方法提升近40%。2. 从零搭建数据管道2.1 获取与理解原始数据Kaggle上的Severstal钢材数据集就像个缺陷动物园包含四种典型问题红斑Red像钢板表面的锈迹划痕Scratches类似钥匙刮擦的细线坑洞Pits不规则的黑色凹陷夹杂Inclusions材料内部的杂质斑点下载后建议这样组织目录steel_defect/ ├── train_images/ # 原始图像(1600x256像素) ├── train_masks/ # 人工标注的缺陷掩膜 ├── train.csv # 每张图缺陷类型和位置 └── class_dict.csv # 颜色编码对照表2.2 数据增强的工业级技巧在钢厂实地考察时发现缺陷可能出现任何角度和位置。因此我们的增强策略要模拟产线实际情况import albumentations as A transform A.Compose([ A.RandomRotate90(), # 钢板可能任意角度进入产线 A.Flip(p0.5), # 正反面都可能出现缺陷 A.RandomBrightnessContrast(p0.2), # 模拟光照变化 A.GaussNoise(var_limit(10,50)), # 模仿相机噪声 A.Cutout(max_h_size20, max_w_size20, num_holes5) # 模拟粉尘干扰 ])特别要注意的是增强时图像和掩膜必须同步变换。有次我忘记设置is_maskTrue参数导致训练时模型完全学偏这个坑希望大家避开。3. 改造UNet适应钢铁场景3.1 网络结构的工程优化原始UNet在医学图像上表现优异但直接用于金属表面检测会遇到两个问题细小裂纹容易被忽略反光区域产生大量误报我们的改进方案class SteelUNet(nn.Module): def __init__(self): super().__init__() # 使用ResNet34作为编码器预训练权重加速收敛 self.encoder timm.create_model(resnet34, features_onlyTrue, pretrainedTrue) # 解码器加入注意力门控 self.decoder UnetDecoder( encoder_channels[64, 64, 128, 256, 512], decoder_channels[256, 128, 64, 32, 16], attention_typescse # 空间-通道注意力 ) # 输出层改用3x3卷积1x1卷积组合 self.seg_head nn.Sequential( nn.Conv2d(16, 16, 3, padding1), nn.ReLU(), nn.Conv2d(16, 4, 1) # 4类缺陷 )3.2 损失函数的实战选择经过多次实验发现组合损失效果最好def loss_function(pred, target): bce_loss F.binary_cross_entropy_with_logits(pred, target) dice_loss 1 - dice_score(pred.sigmoid(), target) return 0.7*bce_loss 0.3*dice_loss在验证集上的表现损失组合mIoU推理速度(FPS)纯BCE0.7245BCEDice(7:3)0.8143FocalDice0.79404. 训练中的避坑指南4.1 学习率与批大小的秘密钢材图像的高分辨率导致显存吃紧。经过测试在RTX 3090上推荐配置批大小(batch_size): 8初始学习率: 3e-4使用梯度累积(steps4)模拟更大batchoptimizer AdamW(model.parameters(), lr3e-4, weight_decay1e-3) scheduler CosineAnnealingLR(optimizer, T_max50, eta_min1e-5)4.2 早停策略的工业考量不同于学术实验工业场景更关注稳定性early_stopping EarlyStopping( patience10, delta0.001, modemax, # 监控验证集IoU checkpoint_pathbest_model.pth )训练曲线显示通常在50-60轮后达到平台期。有次我强行训练到200轮结果过拟合导致产线上误检率飙升不得不回滚模型版本。5. 部署时的工程细节5.1 模型轻量化实战使用TensorRT加速的步骤# 转换ONNX格式 torch.onnx.export(model, dummy_input, unet.onnx, opset_version11, do_constant_foldingTrue) # TensorRT优化 trtexec --onnxunet.onnx --saveEngineunet.engine \ --fp16 --workspace4096优化前后对比版本精度(mIoU)延迟(ms)显存占用原始PyTorch0.81582.1GBTensorRT-FP320.80221.4GBTensorRT-FP160.79150.9GB5.2 PyQt5界面的工业设计产线工人需要的是一键操作的简洁界面class DefectApp(QMainWindow): def __init__(self): super().__init__() self.model load_engine(unet.engine) # 加载TRT模型 self.initUI() def initUI(self): # 大按钮更适合戴手套操作 self.btn_load QPushButton(加载钢板图像, self) self.btn_load.setFixedSize(200, 60) self.btn_load.clicked.connect(self.load_image) # 双视图显示 self.lbl_orig QLabel(self) self.lbl_result QLabel(self) # 缺陷统计面板 self.table QTableWidget(4, 3, self) self.table.setHorizontalHeaderLabels([类型, 面积, 严重度])6. 产线实测经验分享在北方某钢厂部署时遇到几个意外情况冬季车间温度低至-15℃导致GPU风扇停转。后来加了工业级防寒机箱。电磁干扰造成图像传输丢帧。改用光纤传输后解决。工人误触保存按钮覆盖好样本。增加了二次确认对话框。目前系统已稳定运行8个月平均检测速度达到每秒3.2米钢板缺陷检出率从人工的82%提升到95%每年为钢厂节省质检成本约240万元。