【YoLo】从 PyTorch Hub 加载 YOLOv5

使用 PyTorch Hub 加载 YOLOv5

此示例从 PyTorch Hub 加载预训练的 YOLOv5s 模型，model并传递图像进行推理。'yolov5s'是最轻最快的 YOLOv5 型号。

import torch

# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# Image
im = 'https://ultralytics.com/images/zidane.jpg'

# Inference
results = model(im)

results.pandas().xyxy[0]
#      xmin    ymin    xmax   ymax  confidence  class    name
# 0  749.50   43.50  1148.0  704.5    0.874023      0  person
# 1  433.50  433.50   517.5  714.5    0.687988     27     tie
# 2  114.75  195.75  1095.0  708.0    0.624512      0  person
# 3  986.00  304.00  1028.0  420.0    0.286865     27     tie

详细示例

此示例显示了PIL和OpenCV图像源的批量推理。可以打印到控制台，保存到，在支持的环境中显示到屏幕上，并以张量或pandas数据帧的形式返回。resultsruns/hub

import cv2
import torch
from PIL import Image

# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# Images
for f in 'zidane.jpg', 'bus.jpg':
    torch.hub.download_url_to_file('https://ultralytics.com/images/' + f, f)  # download 2 images
im1 = Image.open('zidane.jpg')  # PIL image
im2 = cv2.imread('bus.jpg')[..., ::-1]  # OpenCV image (BGR to RGB)

# Inference
results = model([im1, im2], size=640) # batch of images

# Results
results.print()  
results.save()  # or .show()

results.xyxy[0]  # im1 predictions (tensor)
results.pandas().xyxy[0]  # im1 predictions (pandas)
#      xmin    ymin    xmax   ymax  confidence  class    name
# 0  749.50   43.50  1148.0  704.5    0.874023      0  person
# 1  433.50  433.50   517.5  714.5    0.687988     27     tie
# 2  114.75  195.75  1095.0  708.0    0.624512      0  person
# 3  986.00  304.00  1028.0  420.0    0.286865     27     tie

 

推理设置

YOLOv5 模型包含各种推理属性，例如置信度阈值、IoU 阈值等，可以通过以下方式设置：

model.conf = 0.25  # NMS confidence threshold
      iou = 0.45  # NMS IoU threshold
      agnostic = False  # NMS class-agnostic
      multi_label = False  # NMS multiple labels per box
      classes = None  # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs
      max_det = 1000  # maximum number of detections per image
      amp = False  # Automatic Mixed Precision (AMP) inference

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

设备

模型创建后可以转移到任何设备：

model.cpu()  # CPU
model.cuda()  # GPU
model.to(device)  # i.e. device=torch.device(0)

模型也可以直接在任何device：

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', device='cpu')  # load on CPU

💡专业提示：在推理之前，输入图像会自动传输到正确的模型设备。

静默加载输出

模型可以静默加载_verbose=False：

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', _verbose=False)  # load silently

输入通道

要加载具有 4 个输入通道而不是默认的 3 个输入通道的预训练 YOLOv5s 模型：

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', channels=4)

在这种情况下，模型将由预训练的权重组成，除了第一个输入层，它不再与预训练的输入层具有相同的形状。输入层将保持由随机权重初始化。

输出类别数

要加载具有 10 个输出类而不是默认的 80 个输出类的预训练 YOLOv5s 模型：

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

在这种情况下，模型将由预训练的权重组成，除了输出层，它们不再与预训练的输出层具有相同的形状。输出层将保持由随机权重初始化。

强制重新加载

如果您在上述步骤中遇到问题，force_reload=True通过丢弃现有缓存并强制从 PyTorch Hub 重新下载最新的 YOLOv5 版本，设置可能会有所帮助。

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

截图推理

要在桌面屏幕上运行推理：

import torch
from PIL import ImageGrab

# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# Image
im = ImageGrab.grab()  # take a screenshot

# Inference
results = model(im)

训练

要加载 YOLOv5 模型进行训练而不是推理，请设置autoshape=False. 要加载具有随机初始化权重的模型（从头开始训练），请使用pretrained=False. 在这种情况下，您必须提供自己的训练脚本。或者，请参阅我们的 YOLOv5训练自定义数据教程以进行模型训练。

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 结果

用于 API 服务。有关详细信息，请参阅#2291和Flask REST API示例。

results = model(im)  # inference

results.ims # array of original images (as np array) passed to model for inference
results.render()  # updates results.ims with boxes and labels
for im in results.ims:
    buffered = BytesIO()
    im_base64 = Image.fromarray(im)
    im_base64.save(buffered, format="JPEG")
    print(base64.b64encode(buffered.getvalue()).decode('utf-8'))  # base64 encoded image with results

裁剪结果

结果可以返回并保存为检测作物：

results = model(im)  # inference
crops = results.crop(save=True)  # cropped detections dictionary

Pandas结果

结果可以作为Pandas DataFrames返回：

results = model(im)  # inference
results.pandas().xyxy[0]  # Pandas DataFrame

Pandas输出（点击展开）

排序结果

结果可以按列排序，即从左到右（x轴）排序车牌数字检测：

results = model(im)  # inference
results.pandas().xyxy[0].sort_values('xmin')  # sorted left-right

箱形裁剪结果

结果可以返回并保存为检测作物：

results = model(im)  # inference
crops = results.crop(save=True)  # cropped detections dictionary

JSON 结果

.pandas()使用该.to_json()方法转换为数据帧后，可以以 JSON 格式返回结果。可以使用orient参数修改 JSON 格式。有关详细信息，请参阅熊猫.to_json() 文档。

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

JSON 输出（点击展开）

定制模型

此示例使用 PyTorch Hub 加载自定义 20 类VOC训练的YOLOv5s 模型'best.pt'。

model = torch.hub.load('ultralytics/yolov5', 'custom', path='path/to/best.pt')  # local model
model = torch.hub.load('path/to/yolov5', 'custom', path='path/to/best.pt', source='local')  # local repo

TensorRT、ONNX 和 OpenVINO 模型

PyTorch Hub 支持对大多数 YOLOv5 导出格式进行推理，包括自定义训练模型。有关导出模型的详细信息，请参阅TFLite、ONNX、CoreML、TensorRT 导出教程。

💡专业提示：在GPU 基准测试中， TensorRT可能比 PyTorch 快 2-5 倍
💡专业提示：在CPU 基准测试中， ONNX和OpenVINO可能比 PyTorch 快 2-3 倍

model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s.pt')  # PyTorch
                                                            'yolov5s.torchscript')  # TorchScript
                                                            'yolov5s.onnx')  # ONNX
                                                            'yolov5s_openvino_model/')  # OpenVINO
                                                            'yolov5s.engine')  # TensorRT
                                                            'yolov5s.mlmodel')  # CoreML (macOS-only)
                                                            'yolov5s.tflite')  # TFLite
                                                            'yolov5s_paddle_model/')  # PaddlePaddle

环境

YOLOv5 可以在以下任何经过验证的最新环境中运行（预装所有依赖项，包括CUDA / CUDNN、Python和PyTorch）：

• 带有免费 GPU 的Google Colab 和 Kaggle笔记本： 
• 谷歌云深度学习虚拟机。请参阅GCP 快速入门指南
• 亚马逊深度学习 AMI。请参阅AWS 快速入门指南
• 泊坞窗图像。请参阅Docker 快速入门指南