Редактировать страницу на GitHub (EN)

Загрузка 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
Совет

Клонирование ultralytics/yolov5 не требуется — PyTorch Hub получает код автоматически.

Загрузка 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 датафреймов.

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 inference results on zidane.jpg YOLO inference results on bus.jpg

For all inference options see YOLOv5 AutoShape() forward method.

Настройки инференса

Модели 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

Устройство (Device)

Модели можно перенести на любое устройство после создания:

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. Кэшированные копии хранятся в ~/.cache/torch/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. Подробности см. в примере 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 Output (click to expand) 
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

Результаты в формате JSON

Results can be returned in JSON format once converted to .pandas() dataframes using the .to_json() method. The JSON format can be modified using the orient argument. See pandas .to_json() documentation for details.

results = model(ims)  # inference
results.pandas().xyxy[0].to_json(orient="records")  # JSON img1 predictions
JSON Output (click to expand) 
[
    {
        "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"
    }
]

Пользовательские модели

This example loads a custom 20-class VOC-trained YOLOv5s model 'best.pt' with 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 для получения подробной информации об экспорте моделей.

Совет
  • TensorRT может быть в 2–5 раз быстрее, чем PyTorch, согласно бенчмаркам GPU.
  • 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 часа и при каждом новом коммите.

Contributors
Редактировать страницу на GitHub (EN)
GLglenn-jocher13PDpderrenger1RIRizwanMunawar1BUBurhan-Q1

Комментарии