Ultralytics YOLO27: Únete a la lista de espera.
Empezar

Cargar YOLOv5 desde PyTorch Hub

Esta guía explica cómo cargar YOLOv5 desde PyTorch Hub, el repositorio de modelos preentrenados de PyTorch.

Antes de empezar#

Instala requirements.txt en un entorno Python>=3.8.0, incluido 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

No es necesario clonar ultralytics/yolov5: PyTorch Hub obtiene el código automáticamente.

Carga YOLOv5 con PyTorch Hub#

Ejemplo sencillo#

En este ejemplo se carga desde PyTorch Hub un modelo YOLOv5s preentrenado como model y se le pasa una imagen para inferencia. 'yolov5s' es uno de los modelos pequeños y rápidos entre los puntos de control preentrenados de YOLOv5.

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 la inferencia por lotes con fuentes de imagen PIL y OpenCV, dos de los tipos de entrada que acepta el método forward AutoShape() de YOLOv5. results se puede imprimir en la consola, guardar en runs/hub, mostrar en pantalla en los entornos compatibles y devolver como tensores o DataFrames de 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(f"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

Configuración de inferencia#

Los modelos YOLOv5 contienen varios atributos de inferencia, como el umbral de confianza y el umbral de IoU, que se pueden establecer así:

model.conf = 0.25  # Umbral de confianza de NMS
model.iou = 0.45  # Umbral de IoU de NMS
model.agnostic = False  # NMS independiente de la clase
model.multi_label = False  # Varias etiquetas por cuadro en NMS
model.classes = None  # (lista opcional) filtrar por clase; p. ej., = [0, 15, 16] para personas, gatos y perros de COCO
model.max_det = 1000  # número máximo de detecciones por imagen
model.amp = False  # Inferencia con precisión mixta automática (AMP)

results = model(im, size=320)  # tamaño de inferencia personalizado

Dispositivo#

Los modelos se pueden transferir a cualquier dispositivo después de crearlos:

model.cpu()  # CPU
model.cuda()  # GPU
model.to(device)  # p. ej., device=torch.device(0)

Los modelos también se pueden crear directamente en cualquier device:

model = torch.hub.load("ultralytics/yolov5", "yolov5s", device="cpu")  # cargar en CPU
Consejo

Las imágenes de entrada se transfieren automáticamente al dispositivo correcto del modelo antes de realizar la inferencia.

Silenciar las salidas#

Los modelos se pueden cargar en silencio con _verbose=False:

model = torch.hub.load("ultralytics/yolov5", "yolov5s", _verbose=False)  # cargar en silencio

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, cuya forma ya no coincide con la de la capa de entrada preentrenada. La capa de entrada conservará la inicialización 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, cuya forma ya no coincide con la de las capas de salida preentrenadas. Las capas de salida conservarán la inicialización con pesos aleatorios.

Forzar la recarga#

Si tienes problemas con los pasos anteriores, definir force_reload=True puede ayudar, ya que descarta la caché existente y fuerza la descarga de una versión nueva de YOLOv5 desde PyTorch Hub. Las copias en caché se encuentran en ~/.cache/torch/hub; eliminar esa carpeta produce el mismo efecto.

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

Inferencia mediante captura de pantalla#

Para realizar inferencias en la pantalla de tu ordenador:

import torch
from PIL import ImageGrab

# Model
model = torch.hub.load("ultralytics/yolov5", "yolov5s")

# Image
im = ImageGrab.grab()  # hacer una captura de pantalla

# Inference
results = model(im)

Inferencia con varias GPU#

Los modelos YOLOv5 se pueden cargar en paralelo en varias GPU mediante inferencias con 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 entrenarlo en lugar de usarlo para inferencia, establece autoshape=False. Para cargar un modelo con pesos inicializados aleatoriamente (y entrenarlo desde cero), usa pretrained=False. En este caso, debes proporcionar tu propio script de entrenamiento o entrenar con train.py, como se muestra en el tutorial de entrenamiento con datos personalizados de YOLOv5.

import torch

model = torch.hub.load("ultralytics/yolov5", "yolov5s", autoshape=False)  # cargar preentrenado
model = torch.hub.load("ultralytics/yolov5", "yolov5s", autoshape=False, pretrained=False)  # cargar desde cero

Resultados en Base64#

Los resultados codificados en Base64 son útiles para servicios API, como en el ejemplo de Flask REST API de YOLOv5.

import base64
from io import BytesIO

from PIL import Image

results = model(im)  # inference

results.ims  # matriz de imágenes originales (como np array) pasada al modelo para realizar inferencias
results.render()  # actualiza results.ims con cuadros y etiquetas
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"))  # imagen con resultados codificada en base64

Resultados recortados#

Los resultados se pueden devolver y guardar como recortes de detecciones:

results = model(im)  # inference
crops = results.crop(save=True)  # diccionario de detecciones recortadas

Resultados de Pandas#

Los resultados se pueden devolver como marcos de datos de Pandas:

results = model(im)  # inference
results.pandas().xyxy[0]  # marco de datos de Pandas
Salida de pandas (haz clic para desplegar) 
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; por ejemplo, para ordenar de izquierda a derecha (eje x) las detecciones de dígitos de matrículas:

results = model(im)  # inference
results.pandas().xyxy[0].sort_values("xmin")  # ordenado de izquierda a derecha

Resultados JSON#

Los resultados se pueden devolver en formato JSON una vez convertidos en marcos de datos .pandas() mediante el método .to_json(). El formato JSON se puede modificar con el argumento orient. Consulta la documentación de pandas .to_json() para obtener más información.

results = model(im)  # inference
results.pandas().xyxy[0].to_json(orient="records")  # predicciones JSON de img1
Salida JSON (haz clic para desplegar) 
[
    {
        "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, entrenado con 20 clases de VOC, 'best.pt', con PyTorch Hub.

import torch

model = torch.hub.load("ultralytics/yolov5", "custom", path="path/to/best.pt")  # modelo local
model = torch.hub.load("path/to/yolov5", "custom", path="path/to/best.pt", source="local")  # repositorio local

Modelos TensorRT, ONNX y OpenVINO#

PyTorch Hub admite inferencias con la mayoría de los formatos de exportación de YOLOv5, incluidos los modelos entrenados con datos personalizados.

Consejo
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.mlpackage")  # CoreML (solo macOS)
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 ofrece diversos entornos listos para usar, cada uno con dependencias esenciales preinstaladas, como CUDA, cuDNN, Python y PyTorch, para que puedas poner en marcha tus proyectos.

Estado del proyecto#

YOLOv5 CI

Esta insignia indica que todas las pruebas de integración continua (CI) de GitHub Actions de YOLOv5 se superan correctamente. Estas pruebas de CI comprueban rigurosamente la funcionalidad y el rendimiento de YOLOv5 en varios aspectos clave: entrenamiento, validación, inferencia, exportación y evaluaciones comparativas. Garantizan un funcionamiento uniforme y fiable en macOS, Windows y Ubuntu, con pruebas cada 24 horas y tras cada nueva confirmación.

Preguntas frecuentes#

  • No. torch.hub.load("ultralytics/yolov5", ...) descarga el código en ~/.cache/torch/hub y los pesos la primera vez que se usa. Solo hay que instalar los paquetes de requirements.txt.

  • Usa el punto de entrada custom con la ruta del punto de control: torch.hub.load("ultralytics/yolov5", "custom", path="runs/train/exp/weights/best.pt"). Añade source="local" y una ruta al repositorio local para cargarlo sin acceso a la red.

  • Sí. Pasa el archivo exportado al mismo cargador custom y YOLOv5 seleccionará el backend correspondiente según la extensión; por ejemplo, yolov5s.onnx, yolov5s.engine o yolov5s_openvino_model/.

  • Pasa force_reload=True para descartar el repositorio almacenado en caché y descargar el código más reciente, o elimina manualmente ~/.cache/torch/hub. Así también obtendrás las nuevas versiones de YOLOv5.

  • results.xyxy[0] devuelve un tensor con filas [x1, y1, x2, y2, confidence, class]; results.pandas().xyxy[0] devuelve un DataFrame con los nombres de las clases, y .to_json(orient="records") aplicado a ese DataFrame genera JSON.

Colaboradores
glenn-jocherglenn-jocher4raimbekovmraimbekovm1pderrengerpderrenger1

Comentarios

Haz una pregunta, comparte tus resultados o ayuda a otra persona.