欢迎访问我的网站,希望内容对您有用,感兴趣的可以加入免费知识星球。

pip 安装 YOLOv5

YOLO 迷途小书童 3年前 (2021-09-15) 2741次浏览 0个评论

环境

  • ubuntu 18.04 64bit
  • GTX 1070Ti
  • anaconda with python 3.8
  • pytorch 1.7.1
  • cuda 10.1
  • yolov5 5.0.9

前言

为了方便使用 yolov5 目标检测,有网友已经将其做成了库,提交到了官方的索引库 pypi 上,这样,我们就可以直接使用 pip 进行安装了,真的是太方便了,其项目地址: https://github.com/fcakyon/yolov5-pip

安装

首先创建一个干净的 python 虚拟环境

conda create -n yolov5pip python=3.8
conda activate yolov5pip

如果 python 版本大于3.7,直接使用 pip 安装

pip install yolov5

如果 python 的版本是3.6的话,需要安装指定版本的 numpymatplotlib

pip install "numpy>=1.18.5,<1.20" "matplotlib>=3.2.2,<4"
pip install yolov5

命令行的使用

yolov5 库安装好后,同时会生成 yolov5 的命令行工具,其使用方法与源码中的 detect.py 非常类似,如下

yolov5 detect --source 0  # webcam
                    file.jpg  # image
                    file.mp4  # video
                    path/  # directory
                    path/*.jpg  # glob
                    rtsp://170.93.143.139/rtplive/470011e600ef003a004ee33696235daa  # rtsp stream
                    rtmp://192.168.1.105/live/test  # rtmp stream
                    http://112.50.243.8/PLTV/88888888/224/3221225900/1.m3u8  # http stream

测试发现,cpu 的环境下,图片、视频检测都没问题,但是如果使用 gpu 的话,就会报错了,这可能是目前版本的 bug,已经提交了 issue

YOLOv5 ? 2021-9-15 torch 1.7.1+cu101 CUDA:0 (NVIDIA GeForce GTX 1070 Ti, 8116.4375MB)

Fusing layers...
Model Summary: 283 layers, 7276605 parameters, 0 gradients
Traceback (most recent call last):
File "/home/xugaoxiang/anaconda3/envs/yolov5pip/bin/yolov5", line 8, in
sys.exit(app())
File "/home/xugaoxiang/anaconda3/envs/yolov5pip/lib/python3.8/site-packages/yolov5/cli.py", line 11, in app
fire.Fire(
File "/home/xugaoxiang/anaconda3/envs/yolov5pip/lib/python3.8/site-packages/fire/core.py", line 141, in Fire
component_trace = _Fire(component, args, parsed_flag_args, context, name)
File "/home/xugaoxiang/anaconda3/envs/yolov5pip/lib/python3.8/site-packages/fire/core.py", line 466, in _Fire
component, remaining_args = _CallAndUpdateTrace(
File "/home/xugaoxiang/anaconda3/envs/yolov5pip/lib/python3.8/site-packages/fire/core.py", line 681, in _CallAndUpdateTrace
component = fn(*varargs, **kwargs)
File "/home/xugaoxiang/anaconda3/envs/yolov5pip/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 26, in decorate_context
return func(*args, **kwargs)
File "/home/xugaoxiang/anaconda3/envs/yolov5pip/lib/python3.8/site-packages/yolov5/detect.py", line 120, in run
model(torch.zeros(1, 3, *imgsz).to(device).type_as(next(model.parameters()))) # run once
TypeError: zeros() argument after * must be an iterable, not int

除了 detect 检测之外,yolov5 还支持 trainvalexport

模型训练的命令如下

yolov5 train --data coco.yaml --cfg yolov5s.yaml --weights '' --batch-size 64
                                    yolov5m                                40
                                    yolov5l                                24
                                    yolov5x                                16

具体可以参考之前的博文 YOLOv5模型训练,训练过程可以使用 gpu 进行加速

另外关于 valexport 命令的使用,可以通过 --help 来查看其具体支持的参数,如

yolov5 export --help

在Python中使用

库安装好了,就可以在 python 中调用了,来看个最简单的示例

# 导入模块
import yolov5

# 载入模型
model = yolov5.load('yolov5s.pt')

# 待检测的图片
img = 'https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg'

# 推理,默认参数
# results = model(img)

# 使用特定尺寸进行推理
# 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]

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

# 保存检测结果图片
results.save(save_dir='results/')

将官方提供的 yolov5s 模型下载下来,然后执行上述脚本,检测结果如下

yolov5-pip

由于使用了 augment 参数,检测结果多了个 person,默认参数的检测结果与原版 yolov5 检测结果一致

使用GPU加速检测

首选需要安装 gpu版本的 pytorch,这里选择 1.7.1

pip install torch==1.7.1+cu101 torchvision==0.8.2+cu101 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html

跟原版一样,通过 device 参数来指定 cuda 加速

from yolov5 import YOLOv5

# 模型路径
model_path = 'yolov5s.pt'

device = "cuda" # or "cpu"

# 模型初始化
yolov5 = YOLOv5(model_path, device)

# 待检测图片
image = 'https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg'

# 推理,默认参数
results = yolov5.predict(image)

# 使用特定尺寸进行推理
results = yolov5.predict(image, size=1280)

# 数据增强,能够检测出更多的目标,当然也有可能出现误检
# results = yolov5.predict(image, augment=True)

# 多个参数一起使用
# results = yolov5.predict(image, size=1280, augment=True)

# 如果需要检测多张图片的话,可以使用列表
# results = yolov5.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/')

目前 gpu 的支持还不够完善,像本地视频、网络视频( rtsp 协议)都不支持检测,期待后续版本更新

参考资料

喜欢 (0)

您必须 登录 才能发表评论!