Vai al contenuto

PyTorch Hub

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

Prima di iniziare

Installa il file requirements.txt in un file Python>=3.8.0 incluso PyTorch>=1.8. I modelli e i dataset vengono scaricati automaticamente dall'ultimarelease di YOLOv5 .

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

💡 Suggerimento: non è necessario clonare https://github.com/ultralytics / yolov5 😃

Carica YOLOv5 con PyTorch Hub

Esempio semplice

Questo esempio carica un modello YOLOv5s pre-addestrato da PyTorch Hub come model e passa un'immagine per l'inferenza. 'yolov5s' è il modello YOLOv5 più leggero e veloce. Per maggiori dettagli su tutti i modelli disponibili, consulta il sito 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 in batch con PIL e OpenCV fonti delle immagini. results può essere stampato alla console, salvato a runs/hub, mostrato per la proiezione sugli ambienti supportati e restituito come tensori o panda 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 risultati dell'inferenza su zidane.jpg YOLO risultati dell'inferenza sul bus.jpg

Per tutte le opzioni di inferenza vedi YOLOv5 AutoShape() in avanti metodo.

Impostazioni dell'inferenza

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

model.conf = 0.25  # NMS confidence threshold
iou = 0.45  # NMS IoU threshold
agnostic = False  # NMS class-agnostic
multi_label = False  # NMS multiple labels per box
classes = None  # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs
max_det = 1000  # maximum number of detections per image
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 in ingresso vengono trasferite automaticamente al dispositivo modello corretto prima dell'inferenza.

Uscite di silenzio

I modelli possono essere caricati silenziosamente con _verbose=False:

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

Canali di ingresso

Per caricare un modello YOLOv5s pre-addestrato con 4 canali di ingresso 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 primo strato di input, che non avrà più la stessa forma dello strato di input preaddestrato. Lo strato di input rimarrà inizializzato con pesi casuali.

Numero di classi

Per caricare un modello YOLOv5s pre-addestrato 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 gli strati di uscita, che non avranno più la stessa forma degli strati di uscita preaddestrati. Gli strati di uscita rimarranno inizializzati con pesi casuali.

Ricarica forzata

Se dovessi riscontrare dei problemi con i passaggi sopra descritti, imposta force_reload=True può essere d'aiuto scaricando 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 dello screenshot

Per eseguire l'inferenza sullo schermo del 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

YOLOv5 i modelli possono essere caricati su più GPU in parallelo con l'inferenza in thread:

import torch
import threading


def run(model, im):
    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()

Formazione

Per caricare un modello YOLOv5 per l'addestramento piuttosto che per l'inferenza, imposta autoshape=False. Per caricare un modello con pesi inizializzati in modo casuale (per addestrarlo da zero) usa pretrained=False. In questo caso devi fornire il tuo script di formazione. In alternativa, consulta il nostro YOLOv5 Esercitazione sui dati personalizzati di Train per la formazione 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

Da utilizzare con i servizi API. Vedi https://github.com/ultralytics/yolov5/pull/2291 e l'esempio di API REST di Flask per maggiori dettagli.

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 colture di rilevamento:

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

Risultati dei panda

I risultati possono essere restituiti come DataFrames di Pandas:

results = model(im)  # inference
results.pandas().xyxy[0]  # Pandas DataFrame
Output di 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 in un riquadro

I risultati possono essere restituiti e salvati come colture 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() utilizzando i dataframe .to_json() metodo. Il formato JSON può essere modificato utilizzando il metodo 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 una classe 20 personalizzata 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

TensorRT, ONNX e OpenVINO Modelli

PyTorch Hub supporta l'inferenza sulla maggior parte dei formati di esportazione di YOLOv5 , compresi i modelli addestrati in modo personalizzato. Per maggiori dettagli sull'esportazione dei modelli, consulta il tutorial TFLite, ONNX, CoreML, TensorRT Export.

💡 Suggerimento: TensorRT può essere fino a 2-5 volte più veloce di PyTorch su benchmark della GPU 💡 ProTip: ONNX e OpenVINO può essere fino a 2-3 volte più veloce di PyTorch nei benchmark per CPU. benchmark della 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 offre una serie di ambienti pronti all'uso, ciascuno preinstallato con le dipendenze essenziali come CUDA, CUDNN, Python, e PyTorchper dare il via ai tuoi progetti.

Stato del progetto

YOLOv5 CI

Questo badge indica che tutti i test di YOLOv5 GitHub Actions Continuous Integration (CI) sono stati superati con successo. Questi test CI verificano rigorosamente la funzionalità e le prestazioni di YOLOv5 in vari aspetti chiave: formazione, validazione, inferenza, esportazione e benchmark. Assicurano un funzionamento coerente e affidabile su macOS, Windows e Ubuntu, con test condotti ogni 24 ore e su ogni nuovo commit.



Creato 2023-11-12, Aggiornato 2023-12-03
Autori: glenn-jocher (3)

Commenti