Загрузка YOLOv5 из PyTorch Hub
📚 В этом руководстве объясняется, как загрузить YOLOv5 🚀 из PyTorch Hub по адресу https://pytorch.org/hub/ultralytics_yolov5.
Прежде чем начать
Установите requirements.txt в среде Python>=3.8.0, включая PyTorch>=1.8. Модели и наборы данных загружаются автоматически из последнего релиза YOLOv5.
pip install -r https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt
💡 ProTip: Клонирование https://github.com/ultralytics/yolov5 не требуется 😃
Загрузите YOLOv5 с помощью PyTorch Hub
Простой пример
Этот пример загружает предварительно обученную модель YOLOv5s из PyTorch Hub как model
и передает изображение для инференса. 'yolov5s'
это самая легкая и быстрая модель YOLOv5. Подробную информацию обо всех доступных моделях см. в README.
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
Подробный пример
В этом примере показано Пакетный вывод (batched inference) с PIL и OpenCV Источники изображений. results
может быть напечатано в консоль, сохранено в runs/hub
, показал на экран в поддерживаемых средах и возвращается как тензоры или pandas dataframes.
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 AutoShape()
вперед метод.
Настройки инференса
Модели YOLOv5 содержат различные атрибуты вывода, такие как порог уверенности, порог IoU и т. д., которые можно установить с помощью:
model.conf = 0.25 # NMS confidence threshold
model.iou = 0.45 # NMS IoU threshold
model.agnostic = False # NMS class-agnostic
model.multi_label = False # NMS multiple labels per box
model.classes = None # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs
model.max_det = 1000 # maximum number of detections per image
model.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
💡 ProTip: Входные изображения автоматически передаются на правильное устройство модели перед выводом.
Отключить вывод
Модели можно загружать в фоновом режиме с помощью _verbose=False
:
model = torch.hub.load("ultralytics/yolov5", "yolov5s", _verbose=False) # load silently
Входные каналы
Чтобы загрузить предварительно обученную модель YOLOv5s с 4 входными каналами вместо стандартных 3:
model = torch.hub.load("ultralytics/yolov5", "yolov5s", channels=4)
В этом случае модель будет состоять из предварительно обученных весов за исключением самого первого входного слоя, который больше не имеет той же формы, что и предварительно обученный входной слой. Входной слой останется инициализированным случайными весами.
Количество классов
Чтобы загрузить предварительно обученную модель YOLOv5s с 10 выходными классами вместо стандартных 80:
model = torch.hub.load("ultralytics/yolov5", "yolov5s", classes=10)
В этом случае модель будет состоять из предварительно обученных весов за исключением выходных слоев, которые больше не имеют той же формы, что и предварительно обученные выходные слои. Выходные слои останутся инициализированными случайными весами.
Принудительная перезагрузка
Если у вас возникли проблемы с вышеуказанными шагами, установка force_reload=True
может помочь, отбросив существующий кеш и принудительно загрузив последнюю версию YOLOv5 из PyTorch Hub.
model = torch.hub.load("ultralytics/yolov5", "yolov5s", force_reload=True) # force reload
Вывод скриншота
Чтобы запустить inference на экране вашего рабочего стола:
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)
Multi-GPU Inference
Модели YOLOv5 можно загружать на несколько GPU параллельно с многопоточным выводом:
import threading
import torch
def run(model, im):
"""Performs inference on an image using a given model and saves the output; model must support `.save()` method."""
results = model(im)
results.save()
# Models
model0 = torch.hub.load("ultralytics/yolov5", "yolov5s", device=0)
model1 = torch.hub.load("ultralytics/yolov5", "yolov5s", device=1)
# Inference
threading.Thread(target=run, args=[model0, "https://ultralytics.com/images/zidane.jpg"], daemon=True).start()
threading.Thread(target=run, args=[model1, "https://ultralytics.com/images/bus.jpg"], daemon=True).start()
Обучение
Чтобы загрузить модель YOLOv5 для обучения, а не для инференса, установите: autoshape=False
. Чтобы загрузить модель со случайно инициализированными весами (для обучения с нуля), используйте pretrained=False
В этом случае вы должны предоставить свой собственный скрипт обучения. В качестве альтернативы, ознакомьтесь с нашим YOLOv5 Руководство по обучению на пользовательских данных для обучения модели.
import torch
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. Подробности смотрите в примере Flask REST API.
import base64
from io import BytesIO
from PIL import Image
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 (нажмите, чтобы развернуть)
print(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
Отсортированные результаты
Результаты можно сортировать по столбцам, например, для сортировки обнаружения цифр номерного знака слева направо (по оси 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 Результаты
Результаты могут быть возвращены в формате JSON после преобразования в .pandas()
dataframes, используя .to_json()
метод. JSON-формат можно изменить с помощью orient
аргумент. См. pandas .to_json()
документацию для получения подробной информации.
results = model(ims) # inference
results.pandas().xyxy[0].to_json(orient="records") # JSON img1 predictions
JSON Output (нажмите, чтобы развернуть)
[
{
"xmin": 749.5,
"ymin": 43.5,
"xmax": 1148.0,
"ymax": 704.5,
"confidence": 0.8740234375,
"class": 0,
"name": "person"
},
{
"xmin": 433.5,
"ymin": 433.5,
"xmax": 517.5,
"ymax": 714.5,
"confidence": 0.6879882812,
"class": 27,
"name": "tie"
},
{
"xmin": 115.25,
"ymin": 195.75,
"xmax": 1096.0,
"ymax": 708.0,
"confidence": 0.6254882812,
"class": 0,
"name": "person"
},
{
"xmin": 986.0,
"ymin": 304.0,
"xmax": 1028.0,
"ymax": 420.0,
"confidence": 0.2873535156,
"class": 27,
"name": "tie"
}
]
Пользовательские модели
Этот пример загружает пользовательский 20-классовый VOC-обученная модель YOLOv5s 'best.pt'
с PyTorch Hub.
import torch
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 Export tutorial.
💡 ProTip: TensorRT может быть в 2-5 раз быстрее, чем PyTorch, в бенчмарках GPU. 💡 ProTip: ONNX и OpenVINO могут быть в 2-3 раза быстрее, чем PyTorch, в бенчмарках CPU
import torch
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.pt") # PyTorch
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.torchscript") # TorchScript
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.onnx") # ONNX
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s_openvino_model/") # OpenVINO
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.engine") # TensorRT
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.mlmodel") # CoreML (macOS-only)
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.tflite") # TFLite
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s_paddle_model/") # PaddlePaddle
Поддерживаемые среды
Ultralytics предоставляет ряд готовых к использованию сред, каждая из которых предварительно установлена с необходимыми зависимостями, такими как CUDA, CUDNN, Python и PyTorch, чтобы дать старт вашим проектам.
- Бесплатные блокноты GPU:
- Google Cloud: Краткое руководство по GCP
- Amazon: Краткое руководство по AWS
- Azure: Краткое руководство по AzureML
- Docker: Краткое руководство по Docker
Статус проекта
Этот значок указывает на то, что все тесты непрерывной интеграции (CI) YOLOv5 GitHub Actions успешно пройдены. Эти тесты CI тщательно проверяют функциональность и производительность YOLOv5 по различным ключевым аспектам: обучение, валидация, вывод, экспорт и бенчмарки. Они обеспечивают стабильную и надежную работу в macOS, Windows и Ubuntu, при этом тесты проводятся каждые 24 часа и при каждом новом коммите.