YOLOv5模型的快速推理方法
一· 概述
ultralytics 开发的目标检测算法 YOLOv5 有着广泛的应用,在速度和精准度上做了良好的平衡,被很多商业化开发所采纳。本文档主要记录使用 pip 安装的方式快速的部署推理 YOLOv5 的模型文件, yolov5 的原始模型文件为 .pt .
本文档适用于 yolov5 v6.0和 v6.1的安装。
二· 环境与依赖库安装
2.1 部署环境
Windows 10Visual Studio 2019CUDA 11.3cuDNN 8.4minicondaPytorch1.7 或 1.8
注:确保除
Pytorch外上述环境已经正确安装,安装过程不再赘述.
2.2 创建虚拟环境并安装依赖
- 创建
conda虚拟环境```shell
conda create -n yolov5_env python=3.8 -y
- 激活虚拟环境
conda activate yolov5_env
- 安装
Pytorch
# 安装指令复制于Pytorch官方网站
# CUDA 11.1
pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html
# CUDA 11.3
conda install pytorch==1.8.1 torchvision==0.9.1 torchaudio==0.8.1 cudatoolkit=11.3 -c pytorch -c conda-forge
注:
Pytorch可根据实际情况安装1.7或1.8版本, 以及对应的CUDA版本。官方网站提供了详细的安装指令。
- 校验
Pytorch安装
python -c "import torch; print(torch.__version__)"
注:如果安装成功,会显示
Pytorch的版本号。如:1.8.1+cu111或1.8.1。
或者执行
python -c "import torch; print(torch.cuda.is_available())"
注:如果安装成功并启用
CUDA,会显示True。
- 安装
yolov5
pip install yolov5==6.1.6
注:
yolov5可根据实际情况安装6.0或6.1版本。pypi提供的小版本号并非和github上的版本号一致,按照大版本号安装即可。
- 校验安装
yolov5 --help
注:如果安装成功,会显示
yolov5的帮助信息。如下所示:
INFO: Showing help with the command 'yolov5 -- --help'.
NAME
yolov5
SYNOPSIS
yolov5 COMMAND
COMMANDS
COMMAND is one of the following:
train
To be called from yolov5.cli
val
detect
export
注:如果执行报错,可能是因为
yolov5的安装路径与当前的命令行所在路径不一致,可以尝试切换到yolov5的安装路径再次执行。不必担心,这不会影响yolov5的使用。
三· 快速推理
由于上述的安装过程,可能在命令行中会因为路径问题导致无法直接使用 yolov5 命令,此处不建议直接使用 yolov5 命令,而是使用 python 脚本的方式进行推理。
3.1 下载模型文件
下载 yolov5 v6.1版本的模型文件,可以在 github 上的 yolov5 releases 页面下载。 这里以 yolov5s.pt 为例。
3.2 推理脚本
import yolov5
from yolov5.models.common import AutoShape, Detections
model_path = 'yolov5s.pt'
# 载入模型
model = yolov5.load(model_path) # yolov5.models.common.AutoShape
# 待检测的图片
img = 'images_D03.png'
# 推理,默认参数
results = model(img) # yolov5.models.common.Detections
# 使用特定尺寸进行推理
# results = model(img, size=1280)
# 数据增强,能够检测出更多的目标,当然也有可能出现误检
# results = model(img, augment=True)
# 检测结果数据解析,所属类别、置信度、目标位置信息
predictions = results.pred[0]
boxes = predictions[:, :4] # x1, x2, y1, y2
scores = predictions[:, 4]
categories = predictions[:, 5]
# 显示检测结果
Detections.show(results,labels=False) # 不显示标签与置信度
Detections.pred
# 保存检测结果图片
results.save(save_dir='results/')
至此,YOLOv5 的快速推理方法已经介绍完毕,可以根据实际情况进行调整。
另一种接口与实现方法(大同小异):
from yolov5 import YOLOv5
import torch
model_path = 'yolov5x_person_det_v2.2.1/weights/best.pt'
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"YOLOv5 is using {device} device")
image_path = 'images_D03_20240506010410_edited_0000000054.png'
model = YOLOv5(model_path, device)
print(type(model))
results = model.predict(image_path)
# 使用特定尺寸进行推理
# results = model.predict(image, size=1280)
# 数据增强,能够检测出更多的目标,当然也有可能出现误检
# results = model.predict(image, augment=True)
# 多个参数一起使用
# results = model.predict(image, size=1280, augment=True)
# 如果需要检测多张图片的话,可以使用列表
# results = model.predict([image1, image2], size=1280, augment=True)
# 检测结果数据解析,所属类别、置信度、目标位置信息
predictions = results.pred[0]
boxes = predictions[:, :4] # x1, x2, y1, y2
scores = predictions[:, 4]
categories = predictions[:, 5]
# 显示检测结果
results.show()
# 保存检测结果
results.save(save_dir='results/')
四· 参考文档
魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。
更多推荐


所有评论(0)