Clasificación de imágenes con Ultralytics YOLO#

Image classification es la más sencilla de las tareas compatibles e implica clasificar una imagen completa en una de entre un conjunto de clases predefinidas.

El resultado de un clasificador de imágenes es una única etiqueta de clase y una puntuación de confianza. La clasificación de imágenes es útil cuando solo necesitas saber a qué clase pertenece una imagen y no necesitas saber dónde están ubicados los objetos de esa clase ni cuál es su forma exacta.

Modelos#

Los modelos Classify preentrenados de YOLO26 se muestran aquí. Los modelos Detect, Segment y Pose están preentrenados en el conjunto de datos COCO, los modelos Semantic están preentrenados en Cityscapes, y los modelos Classify están preentrenados en el conjunto de datos ImageNet.

Los modelos se descargan automáticamente de la última versión de Ultralytics en el primer uso.

• Los valores acc son las precisiones del modelo en el conjunto de validación del conjunto de datos ImageNet.
  Reproduce con yolo val classify data=path/to/ImageNet device=0
• Speed promediada sobre imágenes de validación de ImageNet utilizando una instancia de Amazon EC2 P4d.
  Reproduce con yolo val classify data=path/to/ImageNet batch=1 device=0|cpu
• Los valores de Params y FLOPs son para el modelo fusionado después de model.fuse(), que combina las capas Conv y BatchNorm. Los puntos de control preentrenados conservan toda la arquitectura de entrenamiento y pueden mostrar recuentos más altos.

Entrenar#

Entrena YOLO26n-cls en el conjunto de datos MNIST160 durante 100 epochs con un tamaño de imagen de 64. Para ver una lista completa de los argumentos disponibles, consulta la página Configuration.

from ultralytics import YOLO

# Load a model
model = YOLO("yolo26n-cls.yaml")  # build a new model from YAML
model = YOLO("yolo26n-cls.pt")  # load a pretrained model (recommended for training)
model = YOLO("yolo26n-cls.yaml").load("yolo26n-cls.pt")  # build from YAML and transfer weights

# Train the model
results = model.train(data="mnist160", epochs=100, imgsz=64)

import torch
import torchvision.transforms as T

from ultralytics import YOLO
from ultralytics.data.dataset import ClassificationDataset
from ultralytics.models.yolo.classify import ClassificationTrainer, ClassificationValidator

class CustomizedDataset(ClassificationDataset):
    """A customized dataset class for image classification with enhanced data augmentation transforms."""

    def __init__(self, root: str, args, augment: bool = False, prefix: str = ""):
        """Initialize a customized classification dataset with enhanced data augmentation transforms."""
        super().__init__(root, args, augment, prefix)

        # Add your custom training transforms here
        train_transforms = T.Compose(
            [
                T.Resize((args.imgsz, args.imgsz)),
                T.RandomHorizontalFlip(p=args.fliplr),
                T.RandomVerticalFlip(p=args.flipud),
                T.RandAugment(interpolation=T.InterpolationMode.BILINEAR),
                T.ColorJitter(brightness=args.hsv_v, contrast=args.hsv_v, saturation=args.hsv_s, hue=args.hsv_h),
                T.ToTensor(),
                T.Normalize(mean=torch.tensor(0), std=torch.tensor(1)),
                T.RandomErasing(p=args.erasing, inplace=True),
            ]
        )

        # Add your custom validation transforms here
        val_transforms = T.Compose(
            [
                T.Resize((args.imgsz, args.imgsz)),
                T.ToTensor(),
                T.Normalize(mean=torch.tensor(0), std=torch.tensor(1)),
            ]
        )
        self.torch_transforms = train_transforms if augment else val_transforms

class CustomizedTrainer(ClassificationTrainer):
    """A customized trainer class for YOLO classification models with enhanced dataset handling."""

    def build_dataset(self, img_path: str, mode: str = "train", batch=None):
        """Build a customized dataset for classification training and the validation during training."""
        return CustomizedDataset(root=img_path, args=self.args, augment=mode == "train", prefix=mode)

class CustomizedValidator(ClassificationValidator):
    """A customized validator class for YOLO classification models with enhanced dataset handling."""

    def build_dataset(self, img_path: str):
        """Build a customized dataset for classification standalone validation (no augmentation)."""
        return CustomizedDataset(root=img_path, args=self.args, augment=False, prefix=self.args.split)

model = YOLO("yolo26n-cls.pt")
model.train(data="imagenet1000", trainer=CustomizedTrainer, epochs=10, imgsz=224, batch=64)
model.val(data="imagenet1000", validator=CustomizedValidator, imgsz=224, batch=64)

Formato del conjunto de datos#

El formato del conjunto de datos de clasificación YOLO se puede consultar en detalle en la Dataset Guide. Los conjuntos de datos de clasificación también se pueden gestionar y etiquetar con las herramientas de anotación de Ultralytics Platform.

Validar#

Valida la accuracy del modelo YOLO26n-cls entrenado en el conjunto de datos MNIST160. No se necesitan argumentos, ya que el model conserva su data y sus argumentos de entrenamiento como atributos del modelo.

from ultralytics import YOLO

# Load a model
model = YOLO("yolo26n-cls.pt")  # load an official model
model = YOLO("path/to/best.pt")  # load a custom model

# Validate the model
metrics = model.val()  # no arguments needed, dataset and settings remembered
metrics.top1  # top1 accuracy
metrics.top5  # top5 accuracy

Predecir#

Utiliza un modelo YOLO26n-cls entrenado para ejecutar predicciones en imágenes.

from ultralytics import YOLO

# Load a model
model = YOLO("yolo26n-cls.pt")  # load an official model
model = YOLO("path/to/best.pt")  # load a custom model

# Predict with the model
results = model("https://ultralytics.com/images/bus.jpg")  # predict on an image

# Access the results
for result in results:
    top1 = result.probs.top1  # top predicted class ID
    top1_conf = result.probs.top1conf  # top prediction confidence
    top1_name = result.names[top1]  # top predicted class name

Consulta todos los detalles del modo predict en la página Predict.

Resultado de la salida#

La clasificación de imágenes devuelve un objeto Results por imagen. El campo de predicción principal es result.probs, que contiene el vector de probabilidad de las clases y funciones auxiliares para las predicciones principales.

Para ver los campos de Results específicos de cada tarea en todas ellas, consulta la sección Predict Results by Task.

Exportar#

Exporta un modelo YOLO26n-cls a un formato diferente como ONNX, CoreML, etc.

from ultralytics import YOLO

# Load a model
model = YOLO("yolo26n-cls.pt")  # load an official model
model = YOLO("path/to/best.pt")  # load a custom-trained model

# Export the model
model.export(format="onnx")

Los formatos de exportación disponibles para YOLO26-cls se encuentran en la tabla siguiente. Puedes exportar a cualquier formato usando el argumento format, es decir, format='onnx' o format='engine'. Puedes predecir o validar directamente en modelos exportados, es decir, yolo predict model=yolo26n-cls.onnx. Se muestran ejemplos de uso para tu modelo una vez que finaliza la exportación.

Consulta todos los detalles de export en la página Export.

FAQ#