Ultralytics YOLO27:
Get Started

Cargar 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, 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#

Este ejemplo carga un modelo YOLOv5s preentrenado desde PyTorch Hub como model y le pasa una imagen para realizar inferencias. 'yolov5s' es un modelo YOLOv5 pequeño y rápido. Consulta el README para obtener información sobre todos los modelos disponibles.

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 inferencias por lotes con imágenes de PIL y OpenCV. results se puede imprimir en la consola, guardar en runs/hub, mostrar en pantalla en los entornos compatibles y devolver como tensores o marcos de datos 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

Consulta el método forward de YOLOv5 AutoShape() para ver todas las opciones de inferencia method.

Configuración de inferencia#

Los modelos YOLOv5 incluyen varios atributos de inferencia, como el umbral de confianza, el umbral de IoU, etc., que se pueden definir 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 en la primera capa de entrada, que ya no tendrá la misma forma que la capa de entrada preentrenada. La capa de entrada se inicializará 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 en las capas de salida, que ya no tendrán la misma forma que las capas de salida preentrenadas. Las capas de salida se inicializará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 realizar inferencias, define autoshape=False. Para cargar un modelo con pesos inicializados aleatoriamente (para entrenarlo desde cero), usa pretrained=False. En este caso, debes proporcionar tu propio script de entrenamiento. También puedes consultar nuestro Tutorial para entrenar con datos personalizados de YOLOv5 para entrenar modelos.

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#

Para usar con servicios API. Consulta el ejemplo de Flask REST API para obtener más información.

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) la detección de dígitos de una matrícula:

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 la inferencia en la mayoría de los formatos de exportación de YOLOv5, incluidos los modelos entrenados con datos personalizados. Consulta el tutorial de exportación a TFLite, ONNX, CoreML y TensorRT para obtener más información sobre cómo exportar modelos.

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 una variedad de 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
GLglenn-jocher3RAraimbekovm1PDpderrenger1

Comentarios