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.txtClonar 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 tieExemplo 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
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 sizeDispositivo
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 CPUAs 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 silentlyCanais 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 reloadInferê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 scratchResultados 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 resultsResultados Cortados
Os resultados podem ser devolvidos e guardados como recortes de deteção:
results = model(im) # inference
crops = results.crop(save=True) # cropped detections dictionaryResultados Pandas
Os resultados podem ser devolvidos como Pandas DataFrames:
results = model(im) # inference
results.pandas().xyxy[0] # Pandas DataFramePandas 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 tieResultados 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-rightResultados 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 predictionsJSON 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 repoModelos 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.
- O TensorRT pode ser 2-5× mais rápido que o PyTorch em benchmarks de GPU.
- O ONNX e OpenVINO podem ser 2-3× mais rápidos que o 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/") # PaddlePaddleAmbientes 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.
- Notebooks com GPU Gratuita:
- Google Cloud: Guia de Início Rápido GCP
- Amazon: Guia de Início Rápido AWS
- Azure: Guia de Início Rápido AzureML
- Docker: Guia de Início Rápido Docker
Status do Projeto
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.