Cargando YOLOv5 desde PyTorch Hub
📚 Esta guía explica cómo cargar YOLOv5 🚀 desde PyTorch Hub en https://pytorch.org/hub/ultralytics_yolov5.
Antes de empezar
Instala requirements.txt en un entorno Python>=3.8.0, incluyendo PyTorch>=1.8. Los modelos y los conjuntos de datos se descargan automáticamente desde la última versión de YOLOv5.
pip install -r https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt
💡 Consejo profesional: No es necesario clonar https://github.com/ultralytics/yolov5 😃
Cargar YOLOv5 con PyTorch Hub
Ejemplo Simple
Este ejemplo carga un modelo YOLOv5s preentrenado desde PyTorch Hub como model
y pasa una imagen para la inferencia. 'yolov5s'
es el modelo YOLOv5 más ligero y rápido. Para obtener detalles sobre todos los modelos disponibles, consulta el 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
Ejemplo Detallado
Este ejemplo muestra inferencia por lotes con PIL y OpenCV fuentes de imagen. results
puede ser impreso en la consola, guardado a datos runs/hub
, mostrado en la pantalla en entornos compatibles, y devuelto como tensores o 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
Para todas las opciones de inferencia, consulta YOLOv5 AutoShape()
forward method.
Ajustes de Inferencia
Los modelos YOLOv5 contienen varios atributos de inferencia, como el umbral de confianza, el umbral de IoU, etc., que se pueden establecer mediante:
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
Los modelos se pueden transferir a cualquier dispositivo después de su creación:
model.cpu() # CPU
model.cuda() # GPU
model.to(device) # i.e. device=torch.device(0)
Los modelos también se pueden crear directamente en cualquier device
:
model = torch.hub.load("ultralytics/yolov5", "yolov5s", device="cpu") # load on CPU
💡 ProTip: Las imágenes de entrada se transfieren automáticamente al dispositivo del modelo correcto antes de la inferencia.
Silenciar Salidas
Los modelos se pueden cargar de forma silenciosa con _verbose=False
:
model = torch.hub.load("ultralytics/yolov5", "yolov5s", _verbose=False) # load silently
Canales de Entrada
Para cargar un modelo YOLOv5s preentrenado con 4 canales de entrada en lugar de los 3 predeterminados:
model = torch.hub.load("ultralytics/yolov5", "yolov5s", channels=4)
En este caso, el modelo estará compuesto por pesos preentrenados excepto por la primera capa de entrada, que ya no tiene la misma forma que la capa de entrada preentrenada. La capa de entrada permanecerá inicializada con pesos aleatorios.
Número de Clases
Para cargar un modelo YOLOv5s preentrenado con 10 clases de salida en lugar de las 80 predeterminadas:
model = torch.hub.load("ultralytics/yolov5", "yolov5s", classes=10)
En este caso, el modelo estará compuesto por pesos preentrenados excepto por las capas de salida, que ya no tienen la misma forma que las capas de salida preentrenadas. Las capas de salida permanecerán inicializadas con pesos aleatorios.
Forzar Recarga
Si tiene problemas con los pasos anteriores, establecer force_reload=True
puede ayudar descartando la caché existente y forzando una nueva descarga de la última versión de YOLOv5 desde PyTorch Hub.
model = torch.hub.load("ultralytics/yolov5", "yolov5s", force_reload=True) # force reload
Inferencia de Captura de Pantalla
Para ejecutar la inferencia en la pantalla de su escritorio:
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)
Inferencia Multi-GPU
Los modelos YOLOv5 se pueden cargar en múltiples GPUs en paralelo con inferencia por hilos:
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()
Entrenamiento
Para cargar un modelo YOLOv5 para entrenamiento en lugar de inferencia, establezca autoshape=False
. Para cargar un modelo con pesos inicializados aleatoriamente (para entrenar desde cero) use pretrained=False
. Debe proporcionar su propio script de entrenamiento en este caso. Alternativamente, consulte nuestro YOLOv5 Tutorial de Entrenamiento con Datos Personalizados para el entrenamiento del 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 en Base64
Para usar con servicios API. Consulta el ejemplo de Flask REST API para obtener más detalles.
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
Los resultados se pueden devolver y guardar como recortes de detección:
results = model(im) # inference
crops = results.crop(save=True) # cropped detections dictionary
Resultados en Pandas
Los resultados se pueden devolver como DataFrames de Pandas:
results = model(im) # inference
results.pandas().xyxy[0] # Pandas DataFrame
Salida de Pandas (haga clic 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
Los resultados se pueden ordenar por columna, es decir, para ordenar la detección de dígitos de la matrícula de izquierda a derecha (eje x):
results = model(im) # inference
results.pandas().xyxy[0].sort_values("xmin") # sorted left-right
Resultados Recortados de los Cuadros Delimitadores
Los resultados se pueden devolver y guardar como recortes de detección:
results = model(im) # inference
crops = results.crop(save=True) # cropped detections dictionary
Resultados JSON
Los resultados pueden devolverse en formato JSON una vez convertidos a .pandas()
dataframes utilizando el método .to_json()
. El formato JSON se puede modificar utilizando el argumento orient
. Consulte pandas .to_json()
documentación para obtener más detalles.
results = model(ims) # inference
results.pandas().xyxy[0].to_json(orient="records") # JSON img1 predictions
Salida JSON (haga clic 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 ejemplo carga un modelo YOLOv5s personalizado de 20 clases VOCentrenado con 'best.pt'
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 y OpenVINO
PyTorch Hub admite la inferencia en la mayoría de los formatos de exportación de YOLOv5, incluidos los modelos entrenados personalizados. Consulte el tutorial de exportación a TFLite, ONNX, CoreML, TensorRT para obtener más detalles sobre la exportación de modelos.
💡 ProTip: TensorRT puede ser hasta 2-5 veces más rápido que PyTorch en benchmarks de GPU. 💡 ProTip: ONNX y OpenVINO pueden ser hasta 2-3 veces más rápidos que PyTorch en 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
Entornos Compatibles
Ultralytics proporciona una gama de entornos listos para usar, cada uno preinstalado con dependencias esenciales como CUDA, CUDNN, Python y PyTorch, para poner en marcha sus proyectos.
- Notebooks GPU gratuitos:
- Google Cloud: Guía de inicio rápido de GCP
- Amazon: Guía de inicio rápido de AWS
- Azure: Guía de inicio rápido de AzureML
- Docker: Guía de inicio rápido de Docker
Estado del Proyecto
Esta insignia indica que todas las pruebas de Integración Continua (CI) de YOLOv5 GitHub Actions están pasando exitosamente. Estas pruebas de CI verifican rigurosamente la funcionalidad y el rendimiento de YOLOv5 en varios aspectos clave: entrenamiento, validación, inferencia, exportación y benchmarks. Aseguran un funcionamiento consistente y confiable en macOS, Windows y Ubuntu, con pruebas realizadas cada 24 horas y con cada nuevo commit.