Загрузка 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

Результаты вывода YOLO на zidane.jpg Результаты вывода YOLO на bus.jpg

Все опции инференса смотрите в 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 часа и при каждом новом коммите.

📅 Создано 1 год назад ✏️ Обновлено 5 месяцев назад