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 安装方法
- 克隆 YOLOv5 仓库:
git clone https://github.com/ultralytics/yolov5.git
cd yolov5- 安装依赖:
pip install -r requirements.txt3. YOLOv5 核心概念
3.1 模型架构
YOLOv5 有多种规格的模型,从 YOLOv5n(最小)到 YOLOv5x(最大),以适应不同的应用场景:
- YOLOv5n:最小模型,速度最快,适合边缘设备
- YOLOv5s:小模型,平衡速度和精度
- YOLOv5m:中模型,更高的精度
- YOLOv5l:大模型,精度更高
- YOLOv5x:最大模型,精度最高
3.2 工作原理
YOLOv5 的工作原理包括以下步骤:
- 特征提取:使用 CSPDarknet 作为骨干网络提取图像特征
- 特征融合:使用 PANet 结构融合不同尺度的特征
- 预测:在多个尺度上预测边界框和类别概率
- 非极大值抑制:过滤重叠的边界框,保留最优结果
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 数据集准备
- 创建数据集目录结构:
dataset/
├── images/
│ ├── train/
│ └── val/
└── labels/
├── train/
└── val/标注数据:使用 LabelImg 等工具标注数据,生成 YOLO 格式的标注文件。
创建配置文件:创建
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.pt6. 模型导出和部署
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 coreml6.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. 进一步学习资源
- YOLOv5 官方文档:https://docs.ultralytics.com/
- YOLOv5 GitHub 仓库:https://github.com/ultralytics/yolov5
- YOLOv5 教程:https://ultralytics.com/yolov5
- 目标检测论文:YOLO 系列论文和相关研究