Ir al contenido

Cargando YOLOv5 desde PyTorch Hub

📚 En esta guía se explica cómo cargar YOLOv5 🚀 de PyTorch Hub en https://pytorch.org/hub/ultralytics_yolov5.

Antes de empezar

Instalar requirements.txt en un Python>=3.8.0 incluyendo PyTorch>=1.8. Los modelos y conjuntos de datos se descargan automáticamente de la últimaversión de YOLOv5 .

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

💡 ProTip: Clonación https://github.com/ultralytics / yolov5 no es necesario 😃

Carga YOLOv5 con PyTorch Hub

Ejemplo sencillo

Este ejemplo carga un modelo YOLOv5s preentrenado de PyTorch Hub como model y pasa una imagen para su inferencia. 'yolov5s' es el modelo más ligero y rápido de YOLOv5 . Para más información sobre todos los modelos disponibles, consulte la página LÉAME.

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 imágenes. results puede ser impreso a la consola, guardado a runs/hub, mostró a la pantalla en entornos compatibles, y devuelto como tensores o pandas marcos de datos.

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 de inferencia sobre zidane.jpg YOLO resultados de inferencia en bus.jpg

Para todas las opciones de inferencia, véase YOLOv5 AutoShape() adelante método.

Ajustes de inferencia

YOLOv5 contienen varios atributos de inferencia, como el umbral de confianza, el umbral de IoU, etc., que pueden establecerse 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 pueden transferirse a cualquier dispositivo tras su creación:

model.cpu()  # CPU
model.cuda()  # GPU
model.to(device)  # i.e. device=torch.device(0)

Los modelos también pueden crearse 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 modelo correcto antes de la inferencia.

Silenciar salidas

Los modelos pueden cargarse silenciosamente 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 la primera capa de entrada, que ya no tendrá la misma forma que la capa de entrada preentrenada. La capa de entrada seguirá 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 las capas de salida, que ya no tendrán la misma forma que las capas de salida preentrenadas. Las capas de salida seguirán inicializadas con pesos aleatorios.

Recarga forzada

Si tiene problemas con los pasos anteriores, configure 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 capturas de pantalla

Para ejecutar la inferencia en la pantalla de tu 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)

Multi-GPU Inferencia

YOLOv5 pueden cargarse en varias GPU en paralelo con inferencia en 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()

Formación

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) utilice pretrained=False. En este caso, deberá proporcionar su propio guión de formación. También puede consultar nuestro YOLOv5 Tutorial Tren de 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 Base64

Para uso con servicios API. Consulte el ejemplo de API REST de Flask 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 pueden devolverse y guardarse como cultivos de detección:

results = model(im)  # inference
crops = results.crop(save=True)  # cropped detections dictionary

Resultados de Pandas

Los resultados pueden ser devueltos como Pandas DataFrames:

results = model(im)  # inference
results.pandas().xyxy[0]  # Pandas DataFrame
Salida de pandas (haga clic para ampliar) 
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 clasificados

Los resultados se pueden ordenar por columnas, es decir, ordenar la detección de dígitos de matrícula de izquierda a derecha (eje x):

results = model(im)  # inference
results.pandas().xyxy[0].sort_values("xmin")  # sorted left-right

Resultados recortados

Los resultados pueden devolverse y guardarse como cultivos 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() utilizando la función .to_json() método. El formato JSON puede modificarse utilizando el método orient argumento. Ver pandas .to_json() documentación para más detalles.

results = model(ims)  # inference
results.pandas().xyxy[0].to_json(orient="records")  # JSON img1 predictions
Salida JSON (haga clic para ampliar) 
[
    {
        "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 una clase 20 personalizada COV-modelo YOLOv5s entrenado 'best.pt' con 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 y OpenVINO Modelos

PyTorch Hub soporta la inferencia en la mayoría de los formatos de exportación de YOLOv5 , incluyendo modelos entrenados a medida. Consulte el tutorial de exportación de TFLite, ONNX, CoreML, TensorRT para obtener más información sobre la exportación de modelos.

💡 ProTip: TensorRT puede ser hasta 2-5X más rápido que PyTorch en GPU puntos de referencia 💡 ProTip: ONNX y OpenVINO puede ser hasta 2-3X más rápido que PyTorch en CPU puntos de referencia

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 serie de entornos listos para usar, cada uno de ellos preinstalado con dependencias esenciales como CUDACUDNN, Pythony PyTorchpara poner en marcha sus proyectos.

Estado del proyecto

YOLOv5 CI

Este distintivo indica que todas las pruebas de integración continua (IC) deYOLOv5 GitHub Actions se han superado con éxito. Estas pruebas de IC comprueban rigurosamente la funcionalidad y el rendimiento de YOLOv5 en varios aspectos clave: formación, validación, inferencia, exportación y puntos de referencia. Garantizan un funcionamiento coherente y fiable en macOS, Windows y Ubuntu, con pruebas realizadas cada 24 horas y en cada nueva confirmación.

Creado hace 1 año ✏️ Actualizado hace 11 días
glenn-jocher UltralíticoAsistente RizwanMunawar Burhan-Q

Comentarios