PyTorch Втулка

📚 Это руководство объясняет, как загрузить 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

💡 Совет: клонирование https://github.com/ultralytics / yolov5 не обязательно 😃

Нагрузка YOLOv5 с концентратором PyTorch

Простой пример

Этот пример загружает предварительно обученную модель 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

Подробный пример

В этом примере показано пакетный вывод с PIL и OpenCV источники изображения. 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

YOLO результаты умозаключений по зидану.jpg

Обо всех опциях вывода см. YOLOv5 AutoShape() вперед метод.

Настройки вывода

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

Входные каналы

Чтобы загрузить предварительно обученную модель 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

Вывод по скриншоту

Чтобы запустить вывод на экране рабочего стола:

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)

МультиGPU Выводы

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-сервисами. Подробнее см. https://github.com/ ultralytics/ yolov5/pull/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

Pandas Output (нажмите, чтобы развернуть)

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() кадры данных с помощью .to_json() метод. Формат JSON может быть изменен с помощью функции orient аргумент. См. pandas .to_json() документация для деталей.

results = model(ims)  # inference
results.pandas().xyxy[0].to_json(orient="records")  # JSON img1 predictions

Вывод в формате JSON (нажмите, чтобы развернуть)

[
    {
        "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 .

💡 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 Облако: Руководство по быстрому запуску GCP
Amazon: Руководство по быстрому запуску AWS
Azure: Руководство по быстрому запуску AzureML
Docker: Руководство по быстрому запуску Docker

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

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

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