PyTorch Hub

📚 This guide explains how to load YOLOv5 🚀 from PyTorch Hub at https://pytorch.org/hub/ultralytics_yolov5.

Antes de começares

Instala requirements.txt em um Python>=3.8.0 incluindo PyTorch>=1.8. Os modelos e conjuntos de dados são transferidos automaticamente a partir daversão mais recente de YOLOv5 .

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

💡 ProTip: Não é necessário clonar https://github.com/ultralytics / yolov5 😃

Carrega YOLOv5 com PyTorch Hub

Exemplo simples

Este exemplo carrega um modelo YOLOv5s pré-treinado do PyTorch Hub como model e passa uma imagem para inferência. 'yolov5s' é o modelo mais leve e mais rápido do YOLOv5 . Para mais informações sobre todos os modelos disponíveis, consulta a LEIA-ME.

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

Exemplo detalhado

Este exemplo mostra-te inferência agrupada com PIL e OpenCV fontes de imagem. results pode ser impresso para a consola, guardado para runs/hub, mostrado para ser analisado em ambientes suportados, e devolvido como tensores ou pandas quadros de dados.

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 resultados da inferência sobre zidane.jpg

Para todas as opções de inferência, consulta YOLOv5 AutoShape() avança método.

Definições de inferência

YOLOv5 contêm vários atributos de inferência, como o limiar de confiança, o limiar IoU, etc., que podem ser definidos por:

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

Os modelos podem ser transferidos para qualquer dispositivo após a criação:

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

Os modelos também podem ser criados diretamente em qualquer device:

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

Dica: As imagens de entrada são automaticamente transferidas para o dispositivo modelo correto antes da inferência.

Silencia as saídas

Os modelos podem ser carregados silenciosamente com _verbose=False:

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

Canais de entrada

Para carregar um modelo YOLOv5s pré-treinado com 4 canais de entrada em vez dos 3 predefinidos:

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

Neste caso, o modelo será composto por pesos pré-treinados , exceto na primeira camada de entrada, que já não tem a mesma forma que a camada de entrada pré-treinada. A camada de entrada continuará a ser inicializada por pesos aleatórios.

Número de classes

Para carregar um modelo YOLOv5s pré-treinado com 10 classes de saída em vez das 80 predefinidas:

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

Nesse caso, o modelo será composto de pesos pré-treinados , exceto para as camadas de saída, que não têm mais a mesma forma que as camadas de saída pré-treinadas. As camadas de saída continuarão a ser inicializadas por pesos aleatórios.

Forçar recarga

Se tiveres problemas com os passos acima, define force_reload=True pode ajudar, eliminando a cache existente e forçando uma nova transferência da versão mais recente de YOLOv5 a partir do PyTorch Hub.

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

Inferência de capturas de ecrã

Para executar a inferência no ecrã do teu ambiente de trabalho:

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)

Multi-GPU Inferência

YOLOv5 pode ser carregado em várias GPUs em paralelo com inferência encadeada:

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

Formação

Para carregar um modelo YOLOv5 para treino em vez de inferência, define autoshape=False. Para carregar um modelo com pesos inicializados aleatoriamente (para treinar do zero) usa pretrained=False. Neste caso, tens de fornecer o teu próprio guião de formação. Em alternativa, consulta o nosso YOLOv5 Tutorial para treinar dados personalizados para o treino do modelo.

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

Resultados Base64

Para utilização com serviços de API. Consulta https://github.com/ultralytics / yolov5/pull/2291 e o exemplo da API REST do Flask para obteres mais informações.

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

Resultados cortados

Os resultados podem ser devolvidos e guardados como culturas de deteção:

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

Resultados Pandas

Os resultados podem ser devolvidos como Pandas DataFrames:

results = model(im)  # inference
results.pandas().xyxy[0]  # Pandas DataFrame

Saída de Pandas (clica para expandir)

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

Resultados ordenados

Os resultados podem ser ordenados por coluna, ou seja, para ordenar a deteção de dígitos de matrículas da esquerda para a direita (eixo x):

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

Resultados cortados em caixa

Os resultados podem ser devolvidos e guardados como culturas de deteção:

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

Resultados JSON

Os resultados podem ser devolvidos no formato JSON depois de convertidos para .pandas() utilizando os quadros de dados .to_json() método. O formato JSON pode ser modificado utilizando o método orient argumento. Vê pandas .to_json() documentação para mais pormenores.

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

Saída JSON (clica para expandir)

[
    {
        "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"
    }
]

Modelos personalizados

Este exemplo carrega uma classe 20 personalizada COV-modelo YOLOv5s treinado 'best.pt' com 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 Modelos

PyTorch O Hub suporta a inferência na maioria dos formatos de exportação YOLOv5 , incluindo modelos treinados personalizados. Consulte o tutorial de exportação TFLite, ONNX, CoreML, TensorRT para obter detalhes sobre a exportação de modelos.

💡 ProTip: TensorRT pode ser até 2-5X mais rápido do que PyTorch em GPU benchmarks ProTip: ONNX e OpenVINO pode ser até 2-3X mais rápido do que PyTorch em CPU benchmarks

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

Ambientes suportados

Ultralytics fornece uma gama de ambientes prontos a utilizar, cada um pré-instalado com dependências essenciais, tais como CUDA, CUDNN, Python, e PyTorchpara dar início aos teus projectos.

Grátis GPU Notebooks:
Google Nuvem: Guia de início rápido do GCP
Amazon: Guia de início rápido do AWS
Azure: Guia de início rápido do AzureML
Docker: Guia de início rápido do Docker

Estado do projeto

Este emblema indica que todos os testes de Integração Contínua (CI) do YOLOv5 GitHub Actions foram aprovados com êxito. Esses testes de CI verificam rigorosamente a funcionalidade e o desempenho do YOLOv5 em vários aspectos principais: treinamento, validação, inferência, exportação e benchmarks. Eles garantem uma operação consistente e confiável no macOS, Windows e Ubuntu, com testes realizados a cada 24 horas e a cada novo commit.

📅 Created 11 months ago ✏️ Updated 18 days ago