前言

上一篇文章讲了怎么用预训练的 YOLO 模型做检测。如果你的检测对象不在 COCO 的 80 个类别里——比如你想检测 PCB 板上的焊点缺陷、仓库里的特定货品、或者工地上的安全帽佩戴——就需要自己训练模型。

这篇文章从头到尾走一遍流程：标注数据、训练模型、评估效果、导出部署。

一、场景定义

在开始标注之前，先把问题定义清楚：

  • 要检测哪些类别？（比如：安全帽、反光背心、安全鞋）
  • 一个画面里通常有几个目标？（确定数据采集量）
  • 有没有特殊情况？（遮挡、光线差、小目标）
  • 精度要求多高？（实时场景 mAP50 > 0.9，辅助分析 0.7 可能就够用）

二、数据标注

工具选择

工具 特点 适合
LabelImg 轻量，单机运行 个人、少量数据
Label Studio Web 界面，团队协作 团队标注
Roboflow 在线平台，含数据增强和版本管理 快速原型、小团队
CVAT 功能最全，支持视频标注 专业标注团队

推荐个人和小团队用 LabelImg 或 Roboflow。

标注流程（以 LabelImg 为例）

1
2
pip install labelImg
labelImg
  1. “Open Dir” 打开图片文件夹
  2. “Change Save Dir” 设置标注文件保存位置
  3. 按 W 开始画框
  4. 输入类别名
  5. 按 Ctrl+S 保存（生成 XML 文件）
  6. 下一张 D，上一张 A

标注完成后，用一个小脚本把 LabelImg 的 Pascal VOC XML 格式转成 YOLO 格式：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import os, xml.etree.ElementTree as ET

def voc_to_yolo(xml_path, class_map, output_path):
    tree = ET.parse(xml_path)
    root = tree.getroot()
    
    # 从 XML 的 size 节点中动态获取图片的真实宽和高，支持多尺度图片数据集
    size = root.find("size")
    img_w = int(size.find("width").text)
    img_h = int(size.find("height").text)
    
    with open(output_path, "w") as f:
        for obj in root.findall("object"):
            cls = class_map[obj.find("name").text]
            bbox = obj.find("bndbox")
            xmin = float(bbox.find("xmin").text)
            ymin = float(bbox.find("ymin").text)
            xmax = float(bbox.find("xmax").text)
            ymax = float(bbox.find("ymax").text)
            # 转 YOLO 归一化格式
            x_center = ((xmin + xmax) / 2) / img_w
            y_center = ((ymin + ymax) / 2) / img_h
            width = (xmax - xmin) / img_w
            height = (ymax - ymin) / img_h
            f.write(f"{cls} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")

三、数据集组织

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
dataset/
├── images/
│   ├── train/
│   │   ├── 001.jpg
│   │   └── ...
│   └── val/
│       ├── 100.jpg
│       └── ...
├── labels/
│   ├── train/
│   │   ├── 001.txt    # 每行: class_id x_center y_center w h
│   │   └── ...
│   └── val/
│       ├── 100.txt
│       └── ...
└── data.yaml

data.yaml：

1
2
3
4
5
6
7
8
9
10
path: ./dataset           # 数据集根目录
train: images/train       # 训练集路径（相对 path）
val: images/val           # 验证集路径

names:
  0: helmet               # 安全帽
  1: vest                  # 反光背心
  2: shoes                 # 安全鞋

nc: 3                      # 类别数

数据量建议：

数据规模 每类图片数 效果预期
最少 50-100 张 能跑，但容易误检漏检
推荐 300-500 张 基本可用
充裕 1000+ 张 效果稳定

四、开始训练

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
from ultralytics import YOLO

# 从头训练
model = YOLO("yolo11n.pt")   # 基于预训练权重初始化

results = model.train(
    data="dataset/data.yaml",
    epochs=100,              # 训练轮数
    imgsz=640,               # 输入图片尺寸
    batch=16,                # 批大小（显存不够就减）
    lr0=0.01,                # 初始学习率
    lrf=0.01,                # 最终学习率 = lr0 * lrf
    optimizer="AdamW",
    patience=15,             # 15 个 epoch 没提升就早停
    device=0,                # GPU 编号，CPU 填 "cpu"
    workers=4,               # 数据加载线程数
    save=True,               # 保存最佳模型
    project="./runs",        # 输出目录
    name="helmet-detection", # 实验名称
    exist_ok=True,
    # 数据增强
    hsv_h=0.015,             # HSV 色相变化
    hsv_s=0.7,               # HSV 饱和度变化
    hsv_v=0.4,               # HSV 亮度变化
    degrees=10.0,            # 随机旋转角度
    translate=0.1,           # 随机平移
    scale=0.5,               # 随机缩放
    flipud=0.0,              # 上下翻转概率
    fliplr=0.5,              # 左右翻转概率
    mosaic=1.0,              # Mosaic 增强
)

训练过程中会输出每个 epoch 的 loss 和 mAP 等指标。训练结束后，./runs/helmet-detection/weights/best.pt 就是表现最好的模型。

五、评估模型

1
2
3
4
5
6
7
8
9
10
11
12
13
model = YOLO("./runs/helmet-detection/weights/best.pt")

# 在验证集上评估
metrics = model.val()
print(f"mAP50: {metrics.box.map50:.3f}")
print(f"mAP50-95: {metrics.box.map:.3f}")

# 在单张图片上预测并可视化
results = model("test_images/sample.jpg")
results[0].show()

# 批量测试
results = model("test_images/", save=True, conf=0.5)

mAP 解读：

mAP50 含义
> 0.95 非常好，接近人类水平
0.85-0.95 生产可用
0.70-0.85 基本可用，某些场景仍需改进
< 0.70 需要继续优化（加数据/调整训练）

六、常见问题与调优

模型漏检严重：

  • 检查训练数据中是否包含了这类漏检场景（光线暗的、角度偏的、遮挡的）
  • 增加这类难例的样本数量
  • 加大 epochs

误检（把背景当目标）太多：

  • 调高 conf 阈值（默认 0.25，可以调到 0.4-0.5）
  • 在训练数据中增加不包含目标的”负样本”图片

小目标检测不到：

  • 调大 imgsz（640 → 1280，注意会增加显存占用和训练时间）
  • 不用 Mosaic 增强（mosaic=0），小目标容易在拼接后变得更小

训练 loss 不下降：

  • 检查数据标注是否正确（有标注文件为空？坐标超出范围？）
  • 调大或调小学习率
  • 检查 data.yaml 路径和类别数配置

七、模型导出与部署

训练完成后导出为推理格式：

1
2
3
4
5
6
7
8
9
10
11
12
13
model = YOLO("./runs/helmet-detection/weights/best.pt")

# 导出 ONNX（通用推理）
model.export(format="onnx", imgsz=640)

# 导出 TensorRT（GPU 加速推理，速度极快）
model.export(format="engine", imgsz=640, device=0)

# 导出 OpenVINO（Intel CPU）
model.export(format="openvino", imgsz=640)

# 导出 ncnn（移动端）
model.export(format="ncnn", imgsz=640)

在生产环境中跑推理：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
from ultralytics import YOLO
import cv2

model = YOLO("./best.pt")

# 图片检测
results = model("input.jpg", conf=0.5)[0]
for box in results.boxes:
    cls_id = int(box.cls[0])
    name = results.names[cls_id]
    conf = float(box.conf[0])
    print(f"{name}: {conf:.2f}")

# 视频流检测
cap = cv2.VideoCapture(0)  # 摄像头
while True:
    ret, frame = cap.read()
    if not ret:
        break
    results = model(frame, conf=0.5, verbose=False)[0]
    annotated = results.plot()
    cv2.imshow("Detection", annotated)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break
cap.release()

八、流程总结

1
2
3
4
5
6
7
8
9
10
确定要检测的类别和目标数量
  → 采集/收集图片（尽量覆盖各种场景）
  → 用 LabelImg/Label Studio 标注
  → 组织成 YOLO 格式数据集
  → 写 data.yaml
  → 选 YOLO11n/s/m 作为起点，开始训练
  → 评估 mAP，检查误检漏检
  → 调参/加数据/继续训练
  → 导出 ONNX/TensorRT
  → 部署上线

以上是用 YOLO 训练自己模型的标准流程。实际项目中，80% 的时间花在数据采集和标注上，训练本身反而很快。