YOLO算法实战：从零构建AI SOP检测模型的代码详解

1次阅读

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

在工厂流水线上，标准作业程序（SOP）的合规性检测长期面临三大挑战：

复杂环境干扰：生产线光照变化、设备反光、工人遮挡等因素导致传统 OpenCV 模板匹配准确率不足 60%
微小目标检测：螺丝是否拧紧、标签粘贴位置等关键细节往往只占图像 3%-5% 像素
实时性要求：传统方法处理单帧需 200-300ms，而流水线节拍通常要求 100ms 内响应

对比实验数据：
| 方法 | 准确率 | 推理速度(FPS) | 硬件成本 |
|—————-|——–|—————|———-|
| Haar 特征分类器 | 58% | 8 | 低 |
| HOG+SVM | 63% | 15 | 中 |
| YOLOv8n | 89% | 45 | 中 |

2023 年主流版本特性对比：

YOLOv5 优势：
成熟的社区生态（GitHub 34k stars）
更友好的 ONNX 导出支持
适合嵌入式部署的 n / s 版本
YOLOv8 亮点：
引入 Anchor-Free 检测头
自带分类和分割任务支持
验证集 mAP 提升约 5%

选型建议：新手推荐 YOLOv8，因为其：
– 预训练模型对小目标检测更优（COCO 上 APs 提升 7.2%）
– 简化了数据增强配置
– 内置超参数进化算法

常见标注格式转换示例（VOC 转 YOLO）：

import xml.etree.ElementTree as ET
import os

def voc_to_yolo(xml_path, class_list):
    tree = ET.parse(xml_path)
    root = tree.getroot()

    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    yolo_lines = []
    for obj in root.iter('object'):
        cls = obj.find('name').text
        cls_id = class_list.index(cls)

        bndbox = obj.find('bndbox')
        xmin = int(bndbox.find('xmin').text)
        ymin = int(bndbox.find('ymin').text)
        xmax = int(bndbox.find('xmax').text)
        ymax = int(bndbox.find('ymax').text)

        # 转换为中心点 + 宽高比例
        x_center = ((xmin + xmax) / 2) / w
        y_center = ((ymin + ymax) / 2) / h
        width = (xmax - xmin) / w
        height = (ymax - ymin) / h

        yolo_lines.append(f"{cls_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")

    return yolo_lines

典型 data.yaml 配置示例：

train: ../dataset/images/train
val: ../dataset/images/val

# 类别数量与名称
nc: 5
names: ['helmet', 'gloves', 'machine', 'screw', 'label']

# 超参数优化建议
augment: True  # 自动启用 Mosaic 增强
fliplr: 0.5   # 水平翻转概率
mixup: 0.2    # 图像混合比例

启动训练命令：

yolo task=detect mode=train model=yolov8n.pt data=data.yaml epochs=100 imgsz=640 batch=16

方案一：FP16 量化加速

from ultralytics import YOLO

# 加载训练好的模型
model = YOLO('best.pt')

# 导出 FP16 格式的 ONNX
model.export(format='onnx', half=True, dynamic=False)

方案二：TensorRT 部署

import tensorrt as trt

# 创建 logger
logger = trt.Logger(trt.Logger.WARNING)

# 构建引擎
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)

# 解析 ONNX 模型
with open('model_fp16.onnx', 'rb') as f:
    parser.parse(f.read())

# 配置优化参数
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)

# 序列化引擎
engine = builder.build_serialized_network(network, config)
with open('engine.trt', 'wb') as f:
    f.write(engine)

高频问题解决方案：

类别不平衡：
使用 Focal Loss 替代 CE Loss

过采样少数类样本

# 在数据加载时加权采样
dataset = LoadImagesAndLabels(..., oversample_thres=0.3)

跨摄像头泛化：
训练时添加不同 ISP 模拟（伽马校正、白平衡扰动）
测试时启用 TTA（Test Time Augmentation）
边缘设备内存优化：
采用切片推理（将大图分割为 512×512 区块）

限制并发推理线程数

// 在 Jetson 上设置 GPU 频率
sudo jetson_clocks --fan

建议尝试以下改进并对比指标变化：

添加 CBAM 注意力模块：

class CBAM(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.ca = ChannelAttention(channels)
        self.sa = SpatialAttention()

    def forward(self, x):
        x = self.ca(x) * x
        x = self.sa(x) * x
        return x

更换损失函数：
用 CIoU 代替 GIoU
增加 Objectness 分支权重
实验记录模板：
| 改进方案 | mAP@0.5 | 参数量(M) | 推理时延(ms) |
|—————-|———|———–|————–|
| Baseline | 0.82 | 3.1 | 15.2 |
| +CBAM | 0.85 | 3.3 | 16.8 |
| +CIoU | 0.84 | 3.1 | 15.4 |