Carregando YOLOv5 do PyTorch Hub

📚 Este guia explica como carregar o YOLOv5 🚀 do PyTorch Hub em https://pytorch.org/hub/ultralytics_yolov5.

Antes de Começar

Instale o requirements.txt em um ambiente Python>=3.8.0, incluindo PyTorch>=1.8. Os modelos e os conjuntos de dados são baixados automaticamente da versão mais recente do YOLOv5 release.

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

💡 Dica Pro: Clonar https://github.com/ultralytics/yolov5 não é necessário 😃

Carregue 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 YOLOv5 mais leve e rápido. Para detalhes sobre todos os modelos disponíveis, consulte o 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

Exemplo Detalhado

Este exemplo mostra inferência em lote com PIL e OpenCV fontes de imagem. results pode ser impresso para o console, salvo para runs/hub, mostrou para tela em ambientes suportados e retornado como tensores ou 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

Resultados da inferência YOLO em zidane.jpg

Para todas as opções de inferência, consulte YOLOv5 AutoShape() Avançar método.

Configurações de Inferência

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

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

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 Pro: As imagens de entrada são automaticamente transferidas para o dispositivo de modelo correto antes da inferência.

Silenciar 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 padrão:

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

Neste caso, o modelo será composto por pesos pré-treinados exceto a primeira camada de entrada, que não tem mais o mesmo formato da camada de entrada pré-treinada. A camada de entrada permanecerá 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 padrão:

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

Neste caso, o modelo será composto por pesos pré-treinados exceto as camadas de saída, que não têm mais o mesmo formato das camadas de saída pré-treinadas. As camadas de saída permanecerão inicializadas por pesos aleatórios.

Recarregar à Força

Se você tiver problemas com as etapas acima, definir force_reload=True pode ajudar descartando o cache existente e forçando um novo download da versão mais recente do YOLOv5 do PyTorch Hub.

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

Inferência de Captura de Tela

Para executar a inferência na tela do seu 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)

Inferência Multi-GPU

Os modelos YOLOv5 podem ser carregados 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()

Treinamento

Para carregar um modelo YOLOv5 para treinamento em vez de inferência, defina autoshape=False. Para carregar um modelo com pesos inicializados aleatoriamente (para treinar do zero), use pretrained=False. Você deve fornecer seu próprio script de treinamento neste caso. Alternativamente, consulte nosso YOLOv5 Tutorial de Treinamento de Dados Personalizados para o treinamento 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 em Base64

Para uso com serviços de API. Consulte o exemplo de Flask REST API para obter detalhes.

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

Resultados Recortados

Os resultados podem ser retornados e salvos como recortes de detecção:

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

Resultados Pandas

Os resultados podem ser retornados como DataFrames Pandas:

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

Saída Pandas (clique 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 classificados por coluna, ou seja, para classificar a detecção de dígitos da placa de licença da esquerda para a direita (eixo x):

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

Resultados com Caixas Recortadas

Os resultados podem ser retornados e salvos como recortes de detecção:

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

Resultados JSON

Os resultados podem ser retornados em formato JSON, uma vez convertidos para .pandas() dataframes usando o .to_json() método. O formato JSON pode ser modificado usando o orient argumento. Veja pandas .to_json() documentação para mais detalhes.

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

Saída JSON (clique 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 um modelo personalizado de 20 classes. VOC-modelo YOLOv5s pré-treinado 'best.pt' com o 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

Modelos TensorRT, ONNX e OpenVINO

O PyTorch Hub suporta 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.

💡 Dica Pro: TensorRT pode ser até 2 a 5 vezes mais rápido que PyTorch em benchmarks de GPU 💡 Dica Pro: ONNX e OpenVINO podem ser até 2 a 3 vezes mais rápidos que PyTorch em benchmarks de 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

Ambientes suportados

A Ultralytics oferece uma variedade de ambientes prontos para uso, cada um pré-instalado com dependências essenciais como CUDA, CUDNN, Python e PyTorch, para dar o pontapé inicial em seus projetos.

Notebooks GPU Gratuitos:
Google Cloud: Guia de Início Rápido do GCP
Amazon: Guia de Início Rápido AWS
Azure: Guia de Início Rápido AzureML
Docker: Guia de Início Rápido do Docker

Estado do Projeto

Este selo indica que todos os testes de Integração Contínua (CI) das Ações do GitHub YOLOv5 estão sendo aprovados com sucesso. 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.

📅 Criado há 1 ano ✏️ Atualizado há 4 meses