Classificação de Imagens com Ultralytics YOLO#

Image classification é a mais simples das tarefas suportadas e envolve a classificação de uma imagem inteira em uma de um conjunto de classes predefinidas.

O resultado de um classificador de imagem é um rótulo de classe único e uma pontuação de confiança. A classificação de imagem é útil quando você precisa apenas saber a qual classe uma imagem pertence e não precisa saber onde os objetos dessa classe estão localizados ou qual é a sua forma exata.

Modelos#

Os modelos YOLO26 Classify pré-treinados são mostrados aqui. Os modelos Detect, Segment e Pose são pré-treinados no conjunto de dados COCO, os modelos Semantic são pré-treinados em Cityscapes, e os modelos Classify são pré-treinados no conjunto de dados ImageNet.

Os Models são baixados automaticamente do release mais recente da Ultralytics no primeiro uso.

• Os valores de acc são as acurácias dos modelos no conjunto de validação do conjunto de dados ImageNet.
  Reproduza com yolo val classify data=path/to/ImageNet device=0
• Speed média calculada nas imagens de validação do ImageNet usando uma instância do Amazon EC2 P4d.
  Reproduza com yolo val classify data=path/to/ImageNet batch=1 device=0|cpu
• Os valores de Params e FLOPs são para o modelo fundido após model.fuse(), que mescla as camadas Conv e BatchNorm. Os checkpoints pré-treinados retêm a arquitetura de treinamento completa e podem mostrar contagens mais altas.

Treinar#

Treina o YOLO26n-cls no conjunto de dados MNIST160 por 100 epochs com um tamanho de imagem de 64. Para ver uma lista completa de argumentos disponíveis, consulte a 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 do conjunto de dados#

O formato do conjunto de dados de classificação YOLO pode ser encontrado em detalhe no Dataset Guide. Os conjuntos de dados de classificação também podem ser geridos e rotulados com as Ultralytics Platform annotation tools.

Validar#

Valida a accuracy do modelo YOLO26n-cls treinado no conjunto de dados MNIST160. Não são necessários argumentos, pois o model retém os seus data e argumentos de treinamento como atributos do 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

Prever#

Use um modelo YOLO26n-cls treinado para executar previsões em imagens.

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 os detalhes do modo predict na página Prever.

Saída de Resultados#

A classificação de imagem retorna um objeto Results por imagem. O campo de previsão principal é result.probs, que contém o vetor de probabilidade de classe e auxiliares para previsões principais.

Para campos de Results específicos de tarefas em cada tarefa, consulta a seção Resultados de Previsão por Tarefa.

Exportar#

Exporte um modelo YOLO26n-cls para um 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")

Os formatos de exportação disponíveis para o YOLO26-cls encontram-se na tabela abaixo. Podes exportar para qualquer formato usando o argumento format, ou seja, format='onnx' ou format='engine'. Podes prever ou validar diretamente em modelos exportados, ou seja, yolo predict model=yolo26n-cls.onnx. Os exemplos de uso são mostrados para o teu modelo após a conclusão da exportação.

Consulta todos os detalhes de export na página Exportar.

FAQ#