Перейти к содержанию

Загрузка 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, чтобы дать старт вашим проектам.

Статус проекта

YOLOv5 CI

Этот значок указывает на то, что все тесты непрерывной интеграции (CI) YOLOv5 GitHub Actions успешно пройдены. Эти тесты CI тщательно проверяют функциональность и производительность YOLOv5 по различным ключевым аспектам: обучение, валидация, вывод, экспорт и бенчмарки. Они обеспечивают стабильную и надежную работу в macOS, Windows и Ubuntu, при этом тесты проводятся каждые 24 часа и при каждом новом коммите.



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

Комментарии