YOLOv5模型的快速推理方法

最新推荐文章于 2024-07-16 11:05:48 发布

lujx_1024

最新推荐文章于 2024-07-16 11:05:48 发布

阅读量671

点赞数 4

文章标签： YOLO

本文链接：https://blog.csdn.net/LJX_ahut/article/details/139173082

版权

一· 概述

ultralytics 开发的目标检测算法 YOLOv5 有着广泛的应用，在速度和精准度上做了良好的平衡，被很多商业化开发所采纳。本文档主要记录使用 pip 安装的方式快速的部署推理 YOLOv5 的模型文件, yolov5 的原始模型文件为 .pt .

本文档适用于 yolov5 v6.0和 v6.1的安装。

二· 环境与依赖库安装

2.1 部署环境

Windows 10
Visual Studio 2019
CUDA 11.3 cuDNN 8.4
miniconda
Pytorch 1.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/')