Passer au contenu

Chargement de YOLOv5 depuis PyTorch Hub

📚 Ce guide explique comment charger YOLOv5 🚀 à partir de PyTorch Hub sur https://pytorch.org/hub/ultralytics_yolov5.

Avant de commencer

Installer requirements.txt dans un environnement Python>=3.8.0, incluant PyTorch>=1.8. Les modèles et les jeux de données sont téléchargés automatiquement à partir de la dernière version de YOLOv5.

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

💡 Conseil de pro : le clonage de https://github.com/ultralytics/yolov5 n'est pas requis 😃

Charger YOLOv5 avec PyTorch Hub

Exemple simple

Cet exemple charge un modèle YOLOv5s pré-entraîné depuis PyTorch Hub comme model et transmet une image pour l'inférence. 'yolov5s' est le modèle YOLOv5 le plus léger et le plus rapide. Pour plus de détails sur tous les modèles disponibles, veuillez consulter le 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

Exemple détaillé

Cet exemple montre inférence par lots avec PIL et module DNN d'OpenCV sources d'images. results peut être imprimé to console, enregistré à runs/hub, affiché à l'écran dans les environnements pris en charge, et renvoyé sous forme de tenseur ou 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

Résultats de l'inférence YOLO sur zidane.jpg Résultats de l'inférence YOLO sur bus.jpg

Pour toutes les options d'inférence, voir YOLOv5 AutoShape() transférer méthode.

Paramètres d'inférence

Les modèles YOLOv5 contiennent divers attributs d'inférence tels que le seuil de confiance, le seuil IoU, etc. qui peuvent être définis par :

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

Dispositif

Les modèles peuvent être transférés vers n'importe quel appareil après leur création :

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

Les modèles peuvent également être créés directement sur n'importe quel device:

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

💡 Conseil de pro : les images d'entrée sont automatiquement transférées vers le bon périphérique de modèle avant l'inférence.

Masquer les sorties

Les modèles peuvent être chargés silencieusement avec _verbose=False:

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

Nombre de canaux d'entrée

Pour charger un modèle YOLOv5s pré-entraîné avec 4 canaux d'entrée au lieu des 3 par défaut :

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

Dans ce cas, le modèle sera composé de poids pré-entraînés à l'exception de la toute première couche d'entrée, qui n'a plus la même forme que la couche d'entrée pré-entraînée. La couche d'entrée restera initialisée par des poids aléatoires.

Nombre de classes

Pour charger un modèle YOLOv5s pré-entraîné avec 10 classes de sortie au lieu des 80 par défaut :

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

Dans ce cas, le modèle sera composé de poids pré-entraînés à l'exception de les couches de sortie, qui n'ont plus la même forme que les couches de sortie pré-entraînées. Les couches de sortie resteront initialisées par des poids aléatoires.

Forcer le rechargement

Si vous rencontrez des problèmes avec les étapes ci-dessus, définissez force_reload=True peut aider en supprimant le cache existant et en forçant un nouveau téléchargement de la dernière version de YOLOv5 depuis PyTorch Hub.

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

Inférence de capture d'écran

Pour exécuter l’inférence sur l’écran de votre bureau :

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)

Inférence multi-GPU

Les modèles YOLOv5 peuvent être chargés sur plusieurs GPU en parallèle avec l'inférence multithread :

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

Entraînement

Pour charger un modèle YOLOv5 pour l'entraînement plutôt que pour l'inférence, définissez autoshape=False. Pour charger un modèle avec des poids initialisés aléatoirement (pour entraîner à partir de zéro), utilisez pretrained=False. Vous devez fournir votre propre script d'entraînement dans ce cas. Alternativement, consultez notre YOLOv5 Tutoriel d'entraînement des données personnalisées pour l'entraînement du modèle.

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

Résultats Base64

Pour une utilisation avec les services API. Voir l'exemple d'API REST Flask pour plus de détails.

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

Résultats du rognage

Les résultats peuvent être retournés et enregistrés sous forme de crops de détection :

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

Résultats Pandas

Les résultats peuvent être retournés sous forme de DataFrames Pandas :

results = model(im)  # inference
results.pandas().xyxy[0]  # Pandas DataFrame
Sortie Pandas (cliquez pour développer) 
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

Résultats triés

Les résultats peuvent être triés par colonne, par exemple pour trier la détection des chiffres des plaques d’immatriculation de gauche à droite (axe des x) :

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

Résultats recadrés en boîte

Les résultats peuvent être retournés et enregistrés sous forme de crops de détection :

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

Résultats JSON

Les résultats peuvent être retournés au format JSON une fois convertis en .pandas() dataframes en utilisant le .to_json() méthode. Le format JSON peut être modifié en utilisant le orient argument. Voir pandas .to_json() documentation pour plus de détails.

results = model(ims)  # inference
results.pandas().xyxy[0].to_json(orient="records")  # JSON img1 predictions
Sortie JSON (cliquez pour développer) 
[
    {
        "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"
    }
]

Modèles personnalisés

Cet exemple charge un modèle personnalisé à 20 classes VOCmodèle YOLOv5s pré-entraîné 'best.pt' avec 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

Modèles TensorRT, ONNX et OpenVINO

PyTorch Hub prend en charge l'inférence sur la plupart des formats d'exportation YOLOv5, y compris les modèles entraînés personnalisés. Consultez le didacticiel d'exportation TFLite, ONNX, CoreML, TensorRT pour plus de détails sur l'exportation des modèles.

💡 Conseil de pro : TensorRT peut être jusqu'à 2 à 5 fois plus rapide que PyTorch sur les benchmarks GPU. 💡 Conseil de pro : ONNX et OpenVINO peuvent être jusqu'à 2 à 3 fois plus rapides que PyTorch sur les benchmarks 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

Environnements pris en charge

Ultralytics propose une gamme d'environnements prêts à l'emploi, chacun préinstallé avec les dépendances essentielles telles que CUDA, CUDNN, Python et PyTorch, pour démarrer vos projets.

État du projet

YOLOv5 CI

Ce badge indique que tous les tests d'intégration continue (CI) YOLOv5 GitHub Actions réussissent. Ces tests CI vérifient rigoureusement la fonctionnalité et les performances de YOLOv5 à travers divers aspects clés : l'entraînement, la validation, l'inférence, l'exportation et les benchmarks. Ils garantissent un fonctionnement cohérent et fiable sur macOS, Windows et Ubuntu, avec des tests effectués toutes les 24 heures et à chaque nouveau commit.



📅 Créé il y a 1 an ✏️ Mis à jour il y a 6 mois
glenn-jocher UltralyticsAssistant RizwanMunawar Burhan-Q

Commentaires