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 https://github.com/ultralytics/yolov5 non è richiesto 😃

Carica 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 dettagli su tutti i modelli disponibili, consultare 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

Questo esempio mostra inferenza a batch con PIL e OpenCV fonti di immagini. results può essere stampato alla console, salvato a runs/hub, mostrato per visualizzare su ambienti supportati e restituito come tensor oppure pandas dataframe.

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

Risultati dell'inferenza YOLO su zidane.jpg

Per tutte le opzioni di inferenza, vedere YOLOv5 AutoShape() inoltrare metodo.

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 a 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 automaticamente trasferite al dispositivo del modello corretto prima dell'inferenza.

Silenzia 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 pre-addestrato con 4 canali di input anziché i 3 predefiniti:

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

In questo caso, il modello sarà composto da pesi pre-addestrati ad eccezione del primo livello di input, che non ha più la stessa forma del livello di input pre-addestrato. Il livello di input rimarrà inizializzato con pesi casuali.

Numero di Classi

Per caricare un modello YOLOv5s pre-addestrato con 10 classi di output anziché le 80 predefinite:

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

In questo caso, il modello sarà composto da pesi pre-addestrati ad eccezione dei livelli di output, che non hanno più la stessa forma dei livelli di output pre-addestrati. I livelli di output rimarranno inizializzati con pesi casuali.

Forza il ricaricamento

Se riscontri problemi con i passaggi precedenti, impostando force_reload=True può essere d'aiuto scartando la cache esistente e forzando un nuovo download dell'ultima versione di YOLOv5 da PyTorch Hub.

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

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

Training

Per caricare un modello YOLOv5 per il training anziché per l'inferenza, impostare 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 YOLOv5 Tutorial sull'addestramento di dati personalizzati 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 i servizi API. Vedi l'esempio di 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

Output Pandas (clicca per espandere)

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 della targa da sinistra a destra (asse x):

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

Risultati ritagliati del riquadro

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

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

Risultati JSON

I risultati possono essere restituiti in formato JSON una volta convertiti in .pandas() dataframe utilizzando .to_json() metodo. Il formato JSON può essere modificato usando il orient argomento. Vedi pandas .to_json() documentazione per i dettagli.

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

Output JSON (clicca per espandere)

[
    {
        "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 personalizzato a 20 classi VOC-modello YOLOv5s addestrato '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 personalizzati. Consulta il tutorial sull'esportazione in TFLite, ONNX, CoreML, TensorRT per i dettagli sull'esportazione dei modelli.

💡 Suggerimento: TensorRT può essere fino a 2-5 volte più veloce di PyTorch nei benchmark GPU. 💡 Suggerimento: ONNX e OpenVINO possono essere fino a 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 gamma di ambienti pronti all'uso, ciascuno preinstallato con dipendenze essenziali come CUDA, CUDNN, Python e PyTorch, per avviare i tuoi progetti.

Notebook GPU gratuiti:
Google Cloud: Guida rapida GCP
Amazon: Guida rapida AWS
Azure: Guida rapida ad AzureML
Docker: Guida rapida a Docker

Stato del progetto

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

📅 Creato 1 anno fa ✏️ Aggiornato 6 mesi fa