Pytorch Hub 两行代码搞定YOLOv5推理

OpenCV学堂

发布于 2023-08-22 13:16:39

1.3K0

微信公众号：OpenCV学堂 关注获取更多计算机视觉与深度学习知识

Pytorch Hub介绍

Pytorch Hub是一个帮助研究者实现模型再现、快速推理验证的预训练模型库与一套相关的API框架。支持远程从github上下载指定模型、上传与分享训练好的模型、支持从本地加载预训练模型、自定义模型。支持模型远程加载与本地推理、当前Pytorch Hub已经对接到Torchvision、YOLOv5、YOLOv8、pytorchvideo等视觉框架

支持一键远程加载各种模型，主要包括

支持根据URL下载指定模型到本地文件夹

YOLOv5支持

首先需要安装下面的依赖包支持

pip install -r https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt

这个时候不需要再把YOLOv5的代码拉取到本地了，两行代码即可实现YOLOv5模型的推理，直接运行下面的代码即可：

支持多张图像推理：

支持本地自定义对象检测模型推理：

支持多个GPU推理模式

不同设备之间切换支持

推理参数支持设置

相关源码贴图如下：

import torch
import cv2 as cv

# load image data
img = "data/images/zidane.jpg"

# 加载本地预训练模型
model = torch.hub.load(repo_or_dir='D:/python/yolov5-7.0-attention/', model='yolov5s', source='local')

# load image
# img = "D:/bird_test/test004.png"

# 加载本地自定义模型
# model = torch.hub.load('D:/python/yolov5-7.0-attention/', 'custom', path='uav_bird.pt', source='local')
results = model(img)

# 显示
frame = results.render()[0]
bgr = cv.cvtColor(frame, cv.COLOR_RGB2BGR)
cv.imshow("Pytorch Hub + YOLOv5 Custom Object Detection", bgr)
cv.waitKey(0)