YOLOv5 von PyTorch Hub laden

📚 Dieser Leitfaden erklärt, wie man YOLOv5 🚀 von PyTorch Hub unter https://pytorch.org/hub/ultralytics_yolov5 lädt.

Bevor Sie beginnen

Installieren Sie requirements.txt in einer Python>=3.8.0-Umgebung, einschließlich PyTorch>=1.8. Modelle und Datensätze werden automatisch vom neuesten YOLOv5 Release heruntergeladen.

pip install -r https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt

💡 Pro-Tipp: Das Klonen von https://github.com/ultralytics/yolov5 ist nicht erforderlich 😃

YOLOv5 mit PyTorch Hub laden

Einfaches Beispiel

Dieses Beispiel lädt ein vortrainiertes YOLOv5s-Modell von PyTorch Hub als model und übergibt ein Bild zur Inferenz. 'yolov5s' ist das leichteste und schnellste YOLOv5-Modell. Einzelheiten zu allen verfügbaren Modellen finden Sie im 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

Detailliertes Beispiel

Dieses Beispiel zeigt Batch-Inferenz mit PIL und OpenCV Bildquellen. results kann sein gedruckt zur Konsole, gespeichert zu runs/hub, gezeigt auf dem Bildschirm in unterstützten Umgebungen und zurückgegeben als Tensors oder 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 Inferenz-Ergebnisse auf zidane.jpg YOLO Inferenz-Ergebnisse auf bus.jpg

Alle Inferenzoptionen finden Sie unter YOLOv5 AutoShape() Vorwärts Methode.

Inferenz-Einstellungen

YOLOv5 Modelle enthalten verschiedene Inferenzattribute wie Konfidenzschwellenwert, IoU-Schwellenwert usw., die wie folgt festgelegt werden können:

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

Gerät

Modelle können nach der Erstellung auf jedes Gerät übertragen werden:

model.cpu()  # CPU
model.cuda()  # GPU
model.to(device)  # i.e. device=torch.device(0)

Modelle können auch direkt auf jedem Gerät erstellt werden device:

model = torch.hub.load("ultralytics/yolov5", "yolov5s", device="cpu")  # load on CPU

💡 Pro-Tipp: Eingabebilder werden vor der Inferenz automatisch auf das korrekte Modellgerät übertragen.

Ausgaben stummschalten

Modelle können lautlos geladen werden mit _verbose=False:

model = torch.hub.load("ultralytics/yolov5", "yolov5s", _verbose=False)  # load silently

Eingangskanäle

So lädst du ein vortrainiertes YOLOv5s-Modell mit 4 Eingangskanälen anstelle der standardmäßigen 3:

model = torch.hub.load("ultralytics/yolov5", "yolov5s", channels=4)

In diesem Fall besteht das Modell aus vortrainierten Gewichten, mit Ausnahme der allerersten Eingabeschicht, die nicht mehr die gleiche Form wie die vortrainierte Eingabeschicht hat. Die Eingabeschicht wird weiterhin mit zufälligen Gewichten initialisiert.

Anzahl der Klassen

So lädst du ein vortrainiertes YOLOv5s-Modell mit 10 Ausgabeklassen anstelle der standardmäßigen 80:

model = torch.hub.load("ultralytics/yolov5", "yolov5s", classes=10)

In diesem Fall besteht das Modell aus vortrainierten Gewichten mit Ausnahme der Ausgabeschichten, die nicht mehr die gleiche Form wie die vortrainierten Ausgabeschichten haben. Die Ausgabeschichten werden weiterhin mit zufälligen Gewichten initialisiert.

Neu laden erzwingen

Wenn Sie bei den obigen Schritten auf Probleme stoßen, kann das Setzen von force_reload=True kann helfen, indem der bestehende Cache verworfen und ein neuer Download der neuesten YOLOv5-Version vom PyTorch Hub erzwungen wird.

model = torch.hub.load("ultralytics/yolov5", "yolov5s", force_reload=True)  # force reload

Screenshot-Inferenz

So führen Sie eine Inferenz auf Ihrem Desktop-Bildschirm aus:

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 Inferenz

YOLOv5 Modelle können parallel auf mehreren GPUs mit Threading-basierter Inferenz geladen werden:

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()

Training

Um ein YOLOv5-Modell für das Training anstelle der Inferenz zu laden, setze autoshape=False. Um ein Modell mit zufällig initialisierten Gewichten zu laden (um von Grund auf neu zu trainieren), verwenden Sie pretrained=False. In diesem Fall müssen Sie Ihr eigenes Trainingsskript bereitstellen. Alternativ können Sie sich unser YOLOv5 ansehen Tutorial zum Trainieren benutzerdefinierter Daten für das Modelltraining.

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-Ergebnisse

Zur Verwendung mit API-Diensten. Siehe Flask REST API Beispiel für Details.

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

Zugeschnittene Ergebnisse

Ergebnisse können zurückgegeben und als Erkennungsausschnitte gespeichert werden:

results = model(im)  # inference
crops = results.crop(save=True)  # cropped detections dictionary

Pandas Ergebnisse

Ergebnisse können als Pandas DataFrames zurückgegeben werden:

results = model(im)  # inference
results.pandas().xyxy[0]  # Pandas DataFrame

Pandas-Ausgabe (zum Erweitern anklicken)

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

Sortierte Ergebnisse

Ergebnisse können nach Spalten sortiert werden, z. B. um die Erkennung von Kennzeichen-Ziffern von links nach rechts (X-Achse) zu sortieren:

results = model(im)  # inference
results.pandas().xyxy[0].sort_values("xmin")  # sorted left-right

Ergebnisse mit zugeschnittenen Boxen

Ergebnisse können zurückgegeben und als Erkennungsausschnitte gespeichert werden:

results = model(im)  # inference
crops = results.crop(save=True)  # cropped detections dictionary

JSON-Ergebnisse

Ergebnisse können im JSON-Format zurückgegeben werden, sobald sie konvertiert wurden in .pandas() DataFrames mit dem .to_json() Methode. Das JSON-Format kann mit dem folgenden Befehl geändert werden: orient Argument. Siehe pandas .to_json() Dokumentation für Details.

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

JSON-Ausgabe (zum Erweitern anklicken)

[
    {
        "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"
    }
]

Benutzerdefinierte Modelle

Dieses Beispiel lädt eine benutzerdefinierte 20-Klassen VOC-trainiertes YOLOv5s-Modell 'best.pt' mit 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- und OpenVINO-Modelle

PyTorch Hub unterstützt die Inferenz auf den meisten YOLOv5-Exportformaten, einschließlich benutzerdefinierter trainierter Modelle. Siehe TFLite, ONNX, CoreML, TensorRT Export Tutorial für Details zum Exportieren von Modellen.

💡 Pro-Tipp: TensorRT kann auf GPU-Benchmarks bis zu 2-5x schneller sein als PyTorch. 💡 Pro-Tipp: ONNX und OpenVINO können auf CPU-Benchmarks bis zu 2-3x schneller sein als PyTorch.

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

Unterstützte Umgebungen

Ultralytics bietet eine Reihe von gebrauchsfertigen Umgebungen, die jeweils mit wichtigen Abhängigkeiten wie CUDA, CUDNN, Python und PyTorch vorinstalliert sind, um Ihre Projekte zu starten.

Kostenlose GPU Notebooks:
Google Cloud: GCP Quickstart-Anleitung
Amazon: AWS Quickstart-Anleitung
Azure: AzureML Quickstart-Anleitung
Docker: Docker Quickstart-Anleitung

Projektstatus

Dieses Badge zeigt an, dass alle YOLOv5 GitHub Actions Continuous Integration (CI)-Tests erfolgreich bestanden wurden. Diese CI-Tests überprüfen rigoros die Funktionalität und Leistung von YOLOv5 in verschiedenen Schlüsselbereichen: Training, Validierung, Inferenz, Export und Benchmarks. Sie gewährleisten einen konsistenten und zuverlässigen Betrieb unter macOS, Windows und Ubuntu, wobei die Tests alle 24 Stunden und bei jedem neuen Commit durchgeführt werden.

📅 Vor 1 Jahr erstellt ✏️ Vor 6 Monaten aktualisiert