【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 scratchBase64 结果
用于 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 DataFrames返回：
results = model(im) # inference
results.pandas().xyxy[0] # Pandas DataFrame熊猫输出（点击展开）
排序结果
结果可以按列排序，即从左到右（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 dictionaryJSON 结果
.pandas()使用该.to_json()方法转换为数据帧后，可以以 JSON 格式返回结果。可以使用orient参数修改 JSON 格式。有关详细信息，请参阅熊猫.to_json() 文档。
results = model(ims) # inference
results.pandas().xyxy[0].to_json(orient="records") # JSON img1 predictionsJSON 输出（点击展开）
定制模型
此示例使用 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 repoTensorRT、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 快速入门指南