Caricamento di YOLOv5 da PyTorch Hub

📚 Questa guida spiega come caricare YOLOv5 🚀 da PyTorch Hub all'indirizzo https://pytorch.org/hub/ultralytics_yolov5.

Prima di iniziare

Installa requirements.txt in un ambiente Python>=3.8.0, incluso PyTorch>=1.8. Modelli e dataset vengono scaricati automaticamente dall'ultima release di YOLOv5.

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

Clonare ultralytics/yolov5 non è richiesto: PyTorch Hub recupera il codice automaticamente.

Caricare YOLOv5 con PyTorch Hub

Esempio semplice

Questo esempio carica un modello YOLOv5s preaddestrato da PyTorch Hub come model e passa un'immagine per l'inferenza. 'yolov5s' è il modello YOLOv5 più leggero e veloce. Per i dettagli su tutti i modelli disponibili, consulta il 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

Esempio dettagliato

This example shows batched inference with PIL and OpenCV image sources. results can be printed to console, saved to runs/hub, showed to screen on supported environments, and returned as tensors or 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 inference results on zidane.jpg YOLO inference results on bus.jpg

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

Impostazioni di inferenza

I modelli YOLOv5 contengono vari attributi di inferenza come soglia di confidenza, soglia IoU, ecc., che possono essere impostati tramite:

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

Dispositivo

I modelli possono essere trasferiti su qualsiasi dispositivo dopo la creazione:

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

I modelli possono anche essere creati direttamente su qualsiasi device:

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

Le immagini di input vengono trasferite automaticamente al dispositivo del modello corretto prima dell'inferenza.

Silenziare gli output

I modelli possono essere caricati silenziosamente con _verbose=False:

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

Canali di input

Per caricare un modello YOLOv5s preaddestrato con 4 canali di input invece dei 3 predefiniti:

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

In questo caso il modello sarà composto da pesi preaddestrati tranne che per il primissimo layer di input, che non ha più la stessa forma del layer di input preaddestrato. Il layer di input rimarrà inizializzato con pesi casuali.

Numero di classi

Per caricare un modello YOLOv5s preaddestrato con 10 classi di output invece delle 80 predefinite:

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

In questo caso il modello sarà composto da pesi preaddestrati tranne che per i layer di output, che non hanno più la stessa forma dei layer di output preaddestrati. I layer di output rimarranno inizializzati con pesi casuali.

Forzare il ricaricamento

Se riscontri problemi con i passaggi precedenti, impostare force_reload=True può essere d'aiuto scartando la cache esistente e forzando un nuovo download dell'ultima versione di YOLOv5 da PyTorch Hub. Le copie in cache si trovano in ~/.cache/torch/hub; eliminare quella cartella ottiene lo stesso effetto.

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

Inferenza da screenshot

Per eseguire l'inferenza sullo schermo del tuo desktop:

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)

Inferenza Multi-GPU

I modelli YOLOv5 possono essere caricati su più GPU in parallelo con inferenza 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()

Addestramento

Per caricare un modello YOLOv5 per l'addestramento anziché per l'inferenza, imposta autoshape=False. Per caricare un modello con pesi inizializzati casualmente (per addestrare da zero) usa pretrained=False. In questo caso devi fornire il tuo script di addestramento. In alternativa, consulta il nostro Tutorial sull'addestramento di dati personalizzati per YOLOv5 per l'addestramento del modello.

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

Risultati Base64

Per l'uso con servizi API. Vedi l'esempio Flask REST API per i dettagli.

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

Risultati ritagliati

I risultati possono essere restituiti e salvati come ritagli di rilevamento:

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

Risultati Pandas

I risultati possono essere restituiti come 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

Risultati ordinati

I risultati possono essere ordinati per colonna, ad esempio per ordinare il rilevamento delle cifre delle targhe da sinistra a destra (asse x):

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

Risultati 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"
    }
]

Modelli personalizzati

Questo esempio carica un modello YOLOv5s personalizzato addestrato su 20 classi VOC chiamato 'best.pt' con 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

Modelli TensorRT, ONNX e OpenVINO

PyTorch Hub supporta l'inferenza sulla maggior parte dei formati di esportazione YOLOv5, inclusi i modelli addestrati in modo personalizzato. Vedi il tutorial sull'esportazione TFLite, ONNX, CoreML, TensorRT per i dettagli sull'esportazione dei modelli.

Suggerimento
  • TensorRT può essere 2-5 volte più veloce di PyTorch nei benchmark GPU.
  • ONNX e OpenVINO possono essere 2-3 volte più veloci di PyTorch nei benchmark 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

Ambienti supportati

Ultralytics fornisce una serie di ambienti pronti all'uso, ciascuno preinstallato con dipendenze essenziali come CUDA, CUDNN, Python e PyTorch, per avviare rapidamente i tuoi progetti.

Stato del progetto

YOLOv5 CI

Questo badge indica che tutti i test di Continuous Integration (CI) di YOLOv5 GitHub Actions vengono superati con successo. Questi test CI verificano rigorosamente la funzionalità e le prestazioni di YOLOv5 su vari aspetti chiave: training, validation, inference, export e benchmarks. Garantiscono un funzionamento coerente e affidabile su macOS, Windows e Ubuntu, con test condotti ogni 24 ore e ad ogni nuovo commit.

Commenti