Carregando YOLOv5 a partir do PyTorch Hub

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

Antes de começares

Instala os requirements.txt num ambiente Python>=3.8.0, incluindo PyTorch>=1.8. Modelos e datasets são descarregados automaticamente a partir da última release do YOLOv5.

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

Clonar o ultralytics/yolov5 não é necessário — o PyTorch Hub busca o código automaticamente.

Carrega o YOLOv5 com o 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, por favor vê 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 a inferência em lote com fontes de imagem PIL e OpenCV. Os results podem ser impressos na consola, guardados em runs/hub, mostrados no ecrã em ambientes suportados, e devolvidos como tensores ou dataframes pandas.

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.

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 num 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 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 predefinidos:

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 já não tem a mesma forma que a camada de entrada pré-treinada. A camada de entrada permanecerá inicializada com 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)

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

Forçar Recarregamento

Se encontrares problemas com os passos acima, definir force_reload=True pode ajudar ao descartar a cache existente e forçar uma nova transferência da última versão do YOLOv5 a partir do PyTorch Hub. Cópias em cache residem em ~/.cache/torch/hub; eliminar essa pasta alcança o mesmo efeito.

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

Inferência de Captura de Ecrã

Para executar a inferência no ecrã do teu computador:

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 múltiplos GPUs em paralelo com inferência por threads:

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 treino em vez de inferência, define autoshape=False. Para carregar um modelo com pesos inicializados aleatoriamente (para treinar do zero), usa pretrained=False. Deves fornecer o teu próprio script de treino neste caso. Alternativamente, vê o nosso Tutorial de Treino com Dados Personalizados do YOLOv5 para treino de modelos.

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 uso com serviços de API. Vê o exemplo de Flask REST API para 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 Cortados

Os resultados podem ser devolvidos e guardados como recortes 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
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

Resultados Ordenados

Os resultados podem ser ordenados por coluna, isto é, 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 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"
    }
]

Modelos Personalizados

Este exemplo carrega um modelo YOLOv5s personalizado de 20 classes treinado em VOC '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 do YOLOv5, incluindo modelos treinados personalizados. Vê o tutorial de Exportação TFLite, ONNX, CoreML, TensorRT para detalhes sobre a exportação de modelos.

Dica
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 fornece uma gama de ambientes prontos a usar, cada um pré-instalado com dependências essenciais como CUDA, CUDNN, Python e PyTorch, para dar início aos teus projetos.

Status do Projeto

YOLOv5 CI

Este selo indica que todos os testes de Integração Contínua (CI) do GitHub Actions do YOLOv5 estão passando com sucesso. Esses testes de CI verificam rigorosamente a funcionalidade e o desempenho do YOLOv5 em vários aspectos fundamentais: 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.

Comentários