YOLOv5 目标检测模型入门

1. YOLOv5 简介

YOLOv5 是一种实时目标检测模型,由 Ultralytics 公司开发,以其速度和准确性著称。YOLO(You Only Look Once)是一种单阶段目标检测算法,它将目标检测问题转化为回归问题,直接从图像中预测边界框和类别概率。

1.1 YOLOv5 的主要特点

  • 实时目标检测:速度快,适合实时应用场景
  • 易于训练和部署:提供完整的训练和部署工具链
  • 支持多种设备和平台:可在 CPU、GPU、边缘设备上运行
  • 提供预训练模型:包含多种规格的预训练模型,适用于不同场景
  • 活跃的社区支持:持续更新和改进

1.2 YOLOv5 的应用场景

  • 安防监控:实时监控和识别人员、车辆等
  • 自动驾驶:识别道路上的车辆、行人、交通标志等
  • 零售分析:货架商品识别和库存管理
  • 医疗影像:医学图像中的目标检测
  • 无人机视觉:空中目标识别和跟踪

2. 安装 YOLOv5

2.1 环境要求

  • Python 3.7 或更高版本
  • PyTorch 1.7 或更高版本
  • CUDA 支持(推荐,用于 GPU 加速)
  • 其他依赖库:numpy, matplotlib, opencv-python 等

2.2 安装方法

  1. 克隆 YOLOv5 仓库:
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
  1. 安装依赖:
pip install -r requirements.txt

3. YOLOv5 核心概念

3.1 模型架构

YOLOv5 有多种规格的模型,从 YOLOv5n(最小)到 YOLOv5x(最大),以适应不同的应用场景:

  • YOLOv5n:最小模型,速度最快,适合边缘设备
  • YOLOv5s:小模型,平衡速度和精度
  • YOLOv5m:中模型,更高的精度
  • YOLOv5l:大模型,精度更高
  • YOLOv5x:最大模型,精度最高

3.2 工作原理

YOLOv5 的工作原理包括以下步骤:

  1. 特征提取:使用 CSPDarknet 作为骨干网络提取图像特征
  2. 特征融合:使用 PANet 结构融合不同尺度的特征
  3. 预测:在多个尺度上预测边界框和类别概率
  4. 非极大值抑制:过滤重叠的边界框,保留最优结果

3.3 输出格式

YOLOv5 的输出格式为:

[center_x, center_y, width, height, confidence, class1, class2, ...]

其中:

  • center_x, center_y:边界框中心点坐标
  • width, height:边界框宽度和高度
  • confidence:边界框包含目标的置信度
  • class1, class2, ...:各个类别的概率

4. 使用预训练模型进行目标检测

4.1 图像检测

使用预训练模型对图像进行目标检测:

import torch
from PIL import Image
import numpy as np
import cv2

# 加载模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # 选择模型规格

# 加载图像
img = 'https://ultralytics.com/images/zidane.jpg'

# 进行检测
results = model(img)

# 显示结果
results.show()  # 显示检测结果
results.save()  # 保存检测结果

# 获取检测数据
detections = results.pandas().xyxy[0]  # 获取检测结果的DataFrame
print(detections)

4.2 视频检测

使用预训练模型对视频进行目标检测:

import torch
import cv2

# 加载模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# 打开视频
cap = cv2.VideoCapture('video.mp4')

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
    
    # 进行检测
    results = model(frame)
    
    # 获取检测结果
    detected_frame = np.squeeze(results.render())
    
    # 显示结果
    cv2.imshow('YOLOv5 Detection', detected_frame)
    
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

4.3 实时摄像头检测

使用预训练模型进行实时摄像头目标检测:

import torch
import cv2

# 加载模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# 打开摄像头
cap = cv2.VideoCapture(0)  # 0 表示默认摄像头

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
    
    # 进行检测
    results = model(frame)
    
    # 获取检测结果
    detected_frame = np.squeeze(results.render())
    
    # 显示结果
    cv2.imshow('YOLOv5 Real-time Detection', detected_frame)
    
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

5. 自定义数据集训练

5.1 数据集准备

  1. 创建数据集目录结构
dataset/
├── images/
│   ├── train/
│   └── val/
└── labels/
    ├── train/
    └── val/
  1. 标注数据:使用 LabelImg 等工具标注数据,生成 YOLO 格式的标注文件。

  2. 创建配置文件:创建 data.yaml 文件,指定数据集路径和类别信息:

path: ../dataset  # 数据集路径
train: images/train  # 训练集图像路径
val: images/val  # 验证集图像路径

nc: 2  # 类别数量
names: ['person', 'car']  # 类别名称

5.2 训练模型

使用 YOLOv5 提供的训练脚本进行训练:

# 训练命令
python train.py --img 640 --batch 16 --epochs 50 --data data.yaml --weights yolov5s.pt

参数说明:

  • --img:输入图像尺寸
  • --batch:批量大小
  • --epochs:训练轮数
  • --data:数据集配置文件
  • --weights:预训练权重文件

5.3 模型评估

训练完成后,可以评估模型性能:

# 评估命令
python val.py --img 640 --data data.yaml --weights runs/train/exp/weights/best.pt

6. 模型导出和部署

6.1 导出模型

YOLOv5 支持导出为多种格式,用于不同的部署场景:

# 导出为 ONNX 格式
python export.py --weights runs/train/exp/weights/best.pt --include onnx

# 导出为 TensorRT 格式
python export.py --weights runs/train/exp/weights/best.pt --include engine

# 导出为 CoreML 格式
python export.py --weights runs/train/exp/weights/best.pt --include coreml

6.2 部署到不同平台

6.2.1 部署到 Web 应用

使用 ONNX Runtime Web 部署到浏览器:

// 加载模型
const session = await ort.InferenceSession.create('yolov5s.onnx');

// 预处理图像
const preprocessedImage = preprocess(image);

// 进行推理
const inputs = { images: preprocessedImage };
const outputs = await session.run(inputs);

// 后处理结果
const results = postprocess(outputs);

6.2.2 部署到移动设备

使用 CoreML 部署到 iOS 设备:

// 加载模型
let model = try VNCoreMLModel(for: YOLOv5().model)

// 创建请求
let request = VNCoreMLRequest(model: model) { request, error in
    // 处理检测结果
    guard let results = request.results as? [VNRecognizedObjectObservation] else { return }
    // 处理检测到的目标
}

// 执行请求
let handler = VNImageRequestHandler(cgImage: image.cgImage!)
try handler.perform([request])

6.2.3 部署到边缘设备

使用 TensorRT 部署到 NVIDIA 边缘设备:

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit

# 加载 TensorRT 引擎
with open('yolov5s.engine', 'rb') as f:
    engine_data = f.read()

runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
engine = runtime.deserialize_cuda_engine(engine_data)

# 创建执行上下文
context = engine.create_execution_context()

# 分配内存
inputs = []
outputs = []
bindings = []

for binding in range(engine.num_bindings):
    size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size * np.dtype(np.float32).itemsize
    device_mem = cuda.mem_alloc(size)
    bindings.append(int(device_mem))
    if engine.binding_is_input(binding):
        inputs.append(device_mem)
    else:
        outputs.append(device_mem)

# 执行推理
stream = cuda.Stream()
cuda.memcpy_htod_async(inputs[0], preprocessed_image, stream)
context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
cuda.memcpy_dtoh_async(output, outputs[0], stream)
stream.synchronize()

7. 实用技巧

7.1 模型选择

根据应用场景选择合适的模型:

  • 边缘设备:YOLOv5n 或 YOLOv5s
  • 服务器端:YOLOv5m、YOLOv5l 或 YOLOv5x
  • 实时应用:YOLOv5n 或 YOLOv5s
  • 高精度要求:YOLOv5l 或 YOLOv5x

7.2 性能优化

  • 批量大小:根据 GPU 内存调整批量大小
  • 图像尺寸:根据目标大小调整输入图像尺寸
  • 置信度阈值:根据应用场景调整置信度阈值
  • 非极大值抑制阈值:调整 NMS 阈值以减少重叠框

7.3 数据增强

YOLOv5 内置了丰富的数据增强策略:

  • 随机缩放:随机缩放图像大小
  • 随机翻转:水平和垂直翻转
  • 随机裁剪:随机裁剪图像
  • 颜色变换:调整亮度、对比度、饱和度等
  • 马赛克增强:将多张图像拼接成一张

8. 常见问题解决

8.1 检测精度问题

  • 问题:检测精度低
  • 解决方案
    • 增加训练数据量
    • 使用更大的模型
    • 调整超参数
    • 改进数据标注质量

8.2 检测速度问题

  • 问题:检测速度慢
  • 解决方案
    • 使用更小的模型
    • 减小输入图像尺寸
    • 使用 GPU 加速
    • 导出为优化格式(如 TensorRT)

8.3 部署问题

  • 问题:部署失败
  • 解决方案
    • 检查依赖库版本
    • 确保硬件兼容性
    • 按照官方文档正确导出模型

9. 总结

YOLOv5 是一种强大的实时目标检测模型,它通过单阶段检测方法实现了速度和精度的良好平衡。无论是使用预训练模型进行快速部署,还是训练自定义模型以适应特定场景,YOLOv5 都提供了完整的工具链和文档支持。

通过本教程的学习,你应该已经掌握了 YOLOv5 的核心概念和基本使用方法,可以开始使用 YOLOv5 进行自己的目标检测项目开发。

10. 进一步学习资源