基于深度学习的工业表面缺陷检测

Situation（背景）

工业产线上的产品表面缺陷检测传统上依赖人工目检，存在效率低、一致性差、疲劳漏检等问题。本项目面向工业质检场景，利用深度学习模型对产品表面图像进行自动缺陷分类（OK / NG），替代人工目检。

数据集为产线实拍图像，存在严重的类别不平衡：正常样本（OK）远多于缺陷样本（NG）。

Task（任务）

在类别不平衡的数据集上训练一个高精度的二分类模型
设计数据增强策略，扩充稀缺的 NG 样本
选择合适的损失函数处理类别不平衡问题
评估模型在鲁棒性测试集上的泛化能力

Action（行动）

模型选型：ResNet50 迁移学习

采用 ImageNet 预训练的 ResNet50 作为骨干网络，替换最后的全连接层为二分类输出（num_classes=2）。迁移学习利用大规模数据集上学到的通用视觉特征，在小样本工业场景下收敛更快、泛化更好。

损失函数：Focal Loss + 类别权重

针对 OK/NG 类别不平衡问题，使用 sklearn.compute_class_weight('balanced') 计算类别权重，并自定义 Focal Loss：

class FocalLoss(nn.Module):
    def __init__(self, alpha=None, gamma=2.0):
        ...
    def forward(self, inputs, targets):
        ce_loss = F.cross_entropy(inputs, targets, reduction='none')
        pt = torch.exp(-ce_loss)
        focal_term = (1 - pt) ** self.gamma
        if self.alpha is not None:
            alpha_factor = self.alpha[targets].to(inputs.device)
            focal_term = alpha_factor * focal_term
        loss = focal_term * ce_loss
        return loss.mean()

Focal Loss 通过 (1 − p_t)^γ 因子自动降低易分样本的权重，使模型聚焦于难分的缺陷样本，gamma=2.0 时效果最佳。

数据增强：在线 + 离线双轨策略

在线增强（训练时随机施加）：

transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ColorJitter(0.1, 0.1, 0.1),
    transforms.RandomAffine(degrees=0, translate=(0.05, 0.05), scale=(0.95, 1.05)),
    transforms.ToTensor(),
    transforms.RandomErasing(p=0.5, scale=(0.01, 0.1), ratio=(0.3, 3.3)),
])

离线增强（扩充 NG 样本集）：对 NG 样本施加 ±15° 随机旋转 + 高斯噪声（σ = 15），生成增强副本加入训练集：

def transform_image(image):
    angle = random.uniform(-15, 15)
    rotated = image.rotate(angle)
    noisy = add_gaussian_noise(rotated, mean=0, std=15)
    return noisy

训练策略

优化器：Adam，初始学习率 10⁻⁵
学习率调度：CosineAnnealingLR，T_max = 100，η_min = 10⁻⁶
早停：patience=10，监控验证集损失
验证集划分：分层采样（stratified），80/20 划分，保持类别比例
TensorBoard：记录训练/验证损失与准确率曲线

鲁棒性评估

构建独立的 robust_val 测试集，评估模型在未见过的数据分布上的泛化能力，输出 precision、recall、F1 等详细指标。

Result（结果）

模型成功收敛，在验证集上达到高准确率
Focal Loss 有效缓解了类别不平衡导致的 NG 召回率偏低问题
离线数据增强显著提升了 NG 样本的分类性能
鲁棒性测试集上模型保持了良好的泛化能力

关键设计决策

决策	原因
ResNet50 而非 ResNet18	更深的网络捕捉更细粒度的缺陷纹理特征
Focal Loss 而非加权 CE	自适应聚焦难分样本，比固定权重更灵活
CosineAnnealingLR	平滑衰减避免学习率阶跃造成的训练震荡
RandomErasing	模拟遮挡，提升模型对局部缺陷的鲁棒性

项目仓库

代码与模型权重已开源：GitHub - SerenaLina/ResNet18

包含：训练脚本（train.py）、增强训练脚本（enhance_train.py）、测试脚本（test.py）、数据增强工具（enhance_comparison_gen.py）、数据分布分析（calculate.py）、训练曲线与可视化图表。

Tian`s Blog