通过 PyTorch Hub 加载 YOLOv5

一、准备

anaconda安装，之后在其内安装的Python3.8

PYTORCH安装请点这里

# CUDA 10.2 GPU采用这个
conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=10.2 -c pytorch

# CUDA 11.1
conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=11.1 -c pytorch -c conda-forge

# CPU Only 我使用的这个
conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cpuonly

如果 pip安装

pip install torch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0

二、简单示例

这里使用轻量级yolov5s模型，测试单个图片，并输出结果

import torch
# Model
# yolov5s表示主目录下的yolov5s.pt，而且必须是主目录下
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # or yolov5m, yolov5l, yolov5x, custom
# 这里，换成自已的模型，调用best.pt
# model = torch.hub.load('ultralytics/yolov5',  'custom', path='best.pt')
model.conf = 0.52  # confidence threshold (0-1)
model.iou = 0.45  # NMS IoU threshold (0-1)
# (optional list) filter by class, i.e. = [0, 15, 16] for persons, cats and dogs
model.classes = None
model.classes = [0, 15, 16,17] # 17表示是马，0是人

# Images
img = 'https://ultralytics.com/images/zidane.jpg'  # or file, Path, PIL, OpenCV, numpy, list
img = 'img_t/bus.jpg'
# Inference
results = model(img)

# Results
results.print()  # or .show(), .save(), .crop(), .pandas(), etc.
print(results.xyxy[0])  # img1 predictions (tensor)
print('----------------')
print(results.pandas().xyxy[0])  # img1 predictions (pandas)

results.save()
#results.crop() # 截取检测的像素后，生成单一图片
#results.pandas()

输出：

YOLOv5s summary: 213 layers, 7225885 parameters, 0 gradients, 16.5 GFLOPs
Adding AutoShape... 
tensor([[6.71788e+02, 3.95372e+02, 8.10000e+02, 8.78361e+02, 8.96172e-01, 0.00000e+00],
        [2.20657e+02, 4.08141e+02, 3.46167e+02, 8.67381e+02, 8.70248e-01, 0.00000e+00],
        [4.92508e+01, 3.89990e+02, 2.48078e+02, 9.12459e+02, 8.51563e-01, 0.00000e+00],
        [4.54152e-02, 5.52411e+02, 6.78823e+01, 8.75375e+02, 5.34941e-01, 0.00000e+00]])
----------------
         xmin        ymin        xmax        ymax  confidence  class    name
0  671.787903  395.372070  810.000000  878.361328    0.896172      0  person
1  220.657059  408.140961  346.167328  867.381165    0.870248      0  person
2   49.250809  389.990479  248.078201  912.458679    0.851563      0  person
3    0.045415  552.411316   67.882317  875.374634    0.534941      0  person
image 1/1: 1080x810 4 persons
Speed: 35.9ms pre-process, 314.1ms inference, 53.9ms NMS per image at shape (1, 3, 640, 480)
Saved 1 image to runs\detect\exp4

三、详细示例

代码中同时使用PIL和OpenCV，识别结果保存在runs/detect目录下。

import torch
from PIL import Image
import cv2
# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # or yolov5m, yolov5l, yolov5x, custom

# Images
img1 = Image.open('img_t/zidane.jpg')  # PIL image
img2 = cv2.imread('img_t/bus.jpg')[:, :, ::-1]  # OpenCV image (BGR to RGB)
imgs = [img1, img2]  # batch of images

# Inference
results = model(imgs, size=640)  # includes NMS

# Results
results.print()  # or .show(), .save(), .crop(), .pandas(), etc.
print(results.xyxy[0])  # img1 predictions (tensor)
print('----------------')
print(results.pandas().xyxy[0])  # img1 predictions (pandas)

results.save()

输出结果：

Speed: 25.4ms pre-process, 312.3ms inference, 3.5ms NMS per image at shape (2, 3, 640, 640)
tensor([[7.42897e+02, 4.79784e+01, 1.14114e+03, 7.16823e+02, 8.80724e-01, 0.00000e+00],
        [4.42039e+02, 4.37350e+02, 4.96720e+02, 7.09883e+02, 6.87266e-01, 2.70000e+01],
        [1.25215e+02, 1.93607e+02, 7.10841e+02, 7.13070e+02, 6.42236e-01, 0.00000e+00],
        [9.82893e+02, 3.08405e+02, 1.02733e+03, 4.20230e+02, 2.63046e-01, 2.70000e+01]])
----------------
         xmin        ymin         xmax        ymax  confidence  class    name
0  742.896973   47.978394  1141.142212  716.823120    0.880724      0  person
1  442.038971  437.349854   496.719574  709.882935    0.687266     27     tie
2  125.215118  193.606750   710.841187  713.070251    0.642236      0  person
3  982.893250  308.404541  1027.329590  420.230103    0.263046     27     tie
Saved 2 images to runs\detect\exp3

图片：

四、参数设置

这里参数主要是指置信度阈值，NMS LOU阈值，类筛选器等模型属性参数。

model.conf = 0.25  # confidence threshold (0-1)
model.iou = 0.45  # NMS IoU threshold (0-1)
model.classes = None  # (optional list) filter by class, i.e. = [0, 15, 16] for persons, cats and dogs

results = model(imgs, size=320)  # custom inference size

五、输入通道设置

加载YOLOv5s模型输入通道默认值为3，可以通过以下方式修改。

# 这里将通道数设置为4
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', channels=4)

六、分类数设置

YOLOv5模型默认分类数为80，可以通过以下方式修改。

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', classes=10)

七、强制重新加载

可以使用force_reload=True帮助清理缓存并且强制更新下载最新YOLOv5版本。

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', force_reload=True)  # force reload

八、训练

加载YOLOv5模型是为了训练而不是检测，可以设置autoshape=False。
加载模型并随即初始化权值可以设置pretrained=False。

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', autoshape=False)  # load pretrained
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', autoshape=False, pretrained=False)  # load scratch

九、Base64结果

示例如下：

import base64
from io import BytesIO
import torch
from PIL import Image
import cv2
import numpy as np

imgs = 'img_t/bus.jpg'
img1 = Image.open('img_t/zidane.jpg')  # PIL image
img2 = cv2.imread('img_t/bus.jpg')[:, :, ::-1]
# imgs = [img1,img2]
results = model(imgs)  # inference

#print(results.imgs)  # array of original images (as np array) passed to model for inference
#print(results.render())  # updates results.imgs with boxes and labels
for img in results.imgs:
    buffered = BytesIO()
    img_base64 = Image.fromarray(img)
    img_base64.save(buffered, format="JPEG")
    base64_data = base64.b64encode(buffered.getvalue()).decode('utf-8')
    # 此时的图片是原图，没有任何标志
    print(base64_data)  # base64 encoded image with results

十、JSON结果

示例如下：

results = model(imgs)  # inference

results.pandas().xyxy[0].to_json(orient="records")  # JSON img1 predictions

[{"xmin":671.787902832,"ymin":395.3720703125,"xmax":810.0,"ymax":878.361328125,"confidence":0.8961722255,"class":0,"name":"person"},{"xmin":220.6570587158,"ymin":408.1409606934,"xmax":346.1673278809,"ymax":867.3811645508,"confidence":0.8702477217,"class":0,"name":"person"},{"xmin":49.2508087158,"ymin":389.9904785156,"xmax":248.0782012939,"ymax":912.4586791992,"confidence":0.8515626192,"class":0,"name":"person"},{"xmin":12.6506881714,"ymin":223.3784332275,"xmax":809.7072143555,"ymax":788.5164794922,"confidence":0.8493340015,"class":5,"name":"bus"},{"xmin":0.045415163,"ymin":552.411315918,"xmax":67.8823165894,"ymax":875.3746337891,"confidence":0.5349411964,"class":0,"name":"person"}]

参考：保存模型与调用