使用 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数据帧的形式返回。results
runs/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 快速入门指南