Saltar para o conteúdo

PyTorch Centro

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

Antes de começar

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

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

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

Carregar 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 pormenores sobre todos os modelos disponíveis, consulte a LEIAME.

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 pormenorizado

Este exemplo mostra inferência agrupada com PIL e OpenCV fontes de imagem. results pode ser impresso para a consola, guardado para runs/hub, apresentado para o ecrã 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 YOLO resultados da inferência no autocarro.jpg

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

Definições de inferência

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

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.

Saídas de silêncio

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 para a 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)

Neste caso, o modelo será composto por pesos pré-treinados , exceto para 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 continuarão a ser inicializadas por pesos aleatórios.

Forçar recarga

Se tiver problemas com os passos acima, defina 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 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 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()

Formação

Para carregar um modelo YOLOv5 para treino em vez de inferência, defina autoshape=False. Para carregar um modelo com pesos inicializados aleatoriamente (para treinar do zero) use pretrained=False. Neste caso, o utilizador deve fornecer o seu próprio guião de formação. Em alternativa, consulte 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. Consulte https://github.com/ultralytics / yolov5/pull/2291 e o exemplo da API REST do Flask para obter detalhes.

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 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 (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 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() dataframes utilizando o .to_json() método. O formato JSON pode ser modificado utilizando o método orient argumento. Ver 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 (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 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 parâmetros de referência

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 seus projectos.

Estado do projeto

YOLOv5 IC

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.

📅C riado há 1 ano ✏️ Atualizado há 1 mês

Comentários