AI 鹰眼目标检测在工业质检中的实战优化方案

1次阅读

共计 2914 个字符，预计需要花费 8 分钟才能阅读完成。

工业质检是制造业中至关重要的环节，传统的人工检测方式不仅效率低下，而且容易受到主观因素的影响。随着计算机视觉技术的发展，AI 目标检测在工业质检中展现出巨大的潜力。然而，工业场景的特殊性也给目标检测带来了诸多挑战。本文将围绕这些挑战，介绍一套基于 YOLOv5 框架优化的 AI 鹰眼目标检测方案，帮助读者快速落地工业质检项目。

工业质检场景中，目标检测面临以下几个主要挑战：

小目标检测 ：工业零件往往尺寸较小，在图像中占比不足 5%，传统检测算法容易漏检。
密集遮挡 ：零件摆放密集时，相互遮挡严重，导致检测框重叠或误检。
光照变化 ：工厂环境光照不均匀，金属表面反光等问题影响检测效果。
实时性要求 ：生产线高速运转，要求检测算法在毫秒级完成推理。

在众多目标检测框架中，YOLOv5 因其优异的性能和易用性成为工业场景的首选：

Faster R-CNN：两阶段检测器，精度高但速度慢，难以满足实时性要求。
YOLOv3/YOLOv4：速度较快，但模型复杂度高，部署难度大。
YOLOv5：
单阶段检测器，推理速度快（100+ FPS）
模型轻量化，易于部署到边缘设备
社区活跃，文档和工具链完善

我们采用 CBAM（Convolutional Block Attention Module）注意力模块，让模型更关注小目标区域。

import torch
import torch.nn as nn

class CBAM(nn.Module):
    def __init__(self, channels, reduction=16):
        super(CBAM, self).__init__()
        # 通道注意力
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.fc = nn.Sequential(nn.Linear(channels, channels // reduction),
            nn.ReLU(),
            nn.Linear(channels // reduction, channels)
        )
        # 空间注意力
        self.conv = nn.Conv2d(2, 1, kernel_size=7, padding=3)

    def forward(self, x):
        # 通道注意力计算
        avg_out = self.fc(self.avg_pool(x).view(x.size(0), -1))
        max_out = self.fc(self.max_pool(x).view(x.size(0), -1))
        channel_att = torch.sigmoid(avg_out + max_out).unsqueeze(2).unsqueeze(3)
        # 空间注意力计算
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        spatial_att = torch.sigmoid(self.conv(torch.cat([avg_out, max_out], dim=1)))
        # 结合两种注意力
        return x * channel_att * spatial_att

工业场景中目标尺寸差异大，我们改进 PANet 结构，增加浅层特征融合：

在 YOLOv5 原有 3 个检测头基础上，增加一个更浅层的检测头
采用 BiFPN 结构进行跨尺度特征融合
使用 GhostNet 轻量化 backbone 减少计算量

为满足工业设备上的实时推理需求，我们采用以下优化：

训练时量化 ：使用 QAT（Quantization Aware Training）
TensorRT 加速 ：FP16 精度下推理速度提升 2-3 倍
模型剪枝 ：移除冗余卷积核，模型大小减少 40%

针对工业图像特点，我们设计特殊的数据增强：

import albumentations as A

# 工业质检专用数据增强
transform = A.Compose([A.RandomBrightnessContrast(p=0.5),  # 模拟光照变化
    A.MotionBlur(blur_limit=7, p=0.3),  # 模拟运动模糊
    A.GridDistortion(p=0.2),            # 模拟镜头畸变
    A.CoarseDropout(max_holes=8, max_height=16, max_width=16, p=0.3)  # 模拟遮挡
])

针对小目标优化损失函数：

class CustomLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.bce = nn.BCEWithLogitsLoss(reduction='none')
        self.small_obj_weight = 2.0  # 小目标权重加倍

    def forward(self, pred, target):
        # 计算基础损失
        loss = self.bce(pred, target)
        # 小目标加权
        small_mask = (target[..., 4] < 0.01)  # 面积小于 1% 的视为小目标
        loss[small_mask] *= self.small_obj_weight
        return loss.mean()

# 导出 ONNX
model.eval()
dummy_input = torch.randn(1, 3, 640, 640)
torch.onnx.export(model, dummy_input, "model.onnx", 
                  opset_version=11, 
                  input_names=["images"], 
                  output_names=["output"])

# TensorRT 加速 (使用 trtexec 命令)
!trtexec --onnx=model.onnx --fp16 --workspace=2048 --saveEngine=model.engine

在自建工业数据集上的测试结果：