【YoLo】从 PyTorch Hub 加载 YOLOv5_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