Chargement de YOLOv5 depuis PyTorch Hub#
📚 Ce guide explique comment charger YOLOv5 🚀 depuis le Hub PyTorch sur https://pytorch.org/hub/ultralytics_yolov5.
Avant de commencer#
Installe requirements.txt dans un environnement Python>=3.8.0, incluant PyTorch>=1.8. Les models et les datasets se téléchargent automatiquement à partir de la dernière release de YOLOv5.
pip install -r https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txtIl n'est pas nécessaire de cloner ultralytics/yolov5 — PyTorch Hub récupère le code automatiquement.
Charger YOLOv5 avec PyTorch Hub#
Exemple simple#
Cet exemple charge un modèle YOLOv5s pré-entraîné depuis PyTorch Hub en tant que 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, consulte 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 tieExemple détaillé#
Cet exemple montre une inférence par lots avec des sources d'images PIL et OpenCV. results peut être affiché dans la console, sauvegardé dans runs/hub, montré à l'écran sur les environnements pris en charge, et renvoyé sous forme de tensors ou de dataframes 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
Pour toutes les options d'inférence, consulte la méthode AutoShape() forward de YOLOv5.
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 sizeAppareil#
Les modèles peuvent être transférés vers n'importe quel périphérique 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 CPULes images d'entrée sont automatiquement transférées vers le périphérique correct du modèle avant l'inférence.
Sorties silencieuses#
Les modèles peuvent être chargés silencieusement avec _verbose=False :
model = torch.hub.load("ultralytics/yolov5", "yolov5s", _verbose=False) # load silentlyCanaux 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 sauf pour 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 sauf pour les couches de sortie, qui n'ont plus la même forme que les couches de sortie pré-entraînée. Les couches de sortie resteront initialisées par des poids aléatoires.
Forcer le rechargement#
Si tu rencontres des problèmes avec les étapes ci-dessus, définir 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. Les copies en cache se trouvent dans ~/.cache/torch/hub ; supprimer ce dossier produit le même effet.
model = torch.hub.load("ultralytics/yolov5", "yolov5s", force_reload=True) # force reloadInférence sur capture d'écran#
Pour exécuter l'inférence sur ton écran de 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 une inférence threadée :
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éfinis autoshape=False. Pour charger un modèle avec des poids initialisés aléatoirement (pour s'entraîner à partir de zéro), utilise pretrained=False. Tu dois fournir ton propre script d'entraînement dans ce cas. Sinon, consulte notre Tutoriel d'entraînement sur des données personnalisées de YOLOv5 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 scratchRésultats Base64#
Pour une utilisation avec des services API. Consulte l'exemple Flask REST API 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 resultsRésultats rognés#
Les résultats peuvent être retournés et enregistrés sous forme de détections rognées :
results = model(im) # inference
crops = results.crop(save=True) # cropped detections dictionaryRésultats Pandas#
Les résultats peuvent être renvoyés sous forme de Pandas DataFrames :
results = model(im) # inference
results.pandas().xyxy[0] # Pandas DataFramePandas 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 tieRésultats triés#
Les résultats peuvent être triés par colonne, par exemple pour trier la détection de chiffres sur plaque d'immatriculation de gauche à droite (axe x) :
results = model(im) # inference
results.pandas().xyxy[0].sort_values("xmin") # sorted left-rightRésultats JSON#
Les résultats peuvent être renvoyés au format JSON une fois convertis en dataframes .pandas() à l'aide de la méthode .to_json(). Le format JSON peut être modifié à l'aide de l'argument orient. Consulte la documentation de pandas .to_json() pour plus de détails.
results = model(ims) # inference
results.pandas().xyxy[0].to_json(orient="records") # JSON img1 predictionsJSON 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"
}
]Modèles personnalisés#
Cet exemple charge un modèle YOLOv5s personnalisé entraîné sur VOC à 20 classes '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 repoModèles TensorRT, ONNX et OpenVINO#
PyTorch Hub prend en charge l'inférence sur la plupart des formats d'exportation de YOLOv5, y compris les modèles entraînés sur mesure. Consulte le tutoriel d'exportation TFLite, ONNX, CoreML, TensorRT pour plus de détails sur l'exportation des modèles.
- TensorRT peut être 2 à 5 fois plus rapide que PyTorch sur les benchmarks GPU.
- ONNX et OpenVINO peuvent être 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/") # PaddlePaddleEnvironnements pris en charge#
Ultralytics propose une gamme d'environnements prêts à l'emploi, chacun préinstallé avec des dépendances essentielles telles que CUDA, CUDNN, Python et PyTorch, pour donner un élan à tes projets.
- Notebooks GPU gratuits :
- Google Cloud : Guide de démarrage rapide GCP
- Amazon : Guide de démarrage rapide AWS
- Azure : Guide de démarrage rapide AzureML
- Docker : Guide de démarrage rapide Docker
Statut du projet#
Ce badge indique que tous les tests d'intégration continue (CI) YOLOv5 GitHub Actions s'exécutent avec succès. Ces tests CI vérifient rigoureusement la fonctionnalité et les performances de YOLOv5 sur divers aspects clés : entraînement, validation, inférence, exportation et 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.