Classificação de imagens com Ultralytics YOLO#

A classificação de imagens é a mais simples das tarefas compatíveis e envolve classificar uma imagem inteira em uma das classes predefinidas.

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

Modelos#

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

Os modelos são baixados automaticamente da versão mais recente do lançamento da Ultralytics na primeira utilização.

• Os valores de acc são as precisões dos modelos no conjunto de validação do ImageNet.
  Reproduza com yolo val classify data=path/to/ImageNet device=0
• A velocidade é calculada como a média sobre as imagens de validação do ImageNet usando uma instância Amazon EC2 P4d.
  Reproduza com yolo val classify data=path/to/ImageNet batch=1 device=0|cpu
• Os valores de Params e FLOPs correspondem ao modelo fundido após model.fuse(), que combina as camadas Conv e BatchNorm. Os checkpoints pré-treinados mantêm toda a arquitetura de treinamento e podem apresentar contagens maiores.

Consulta a pré-visualização não lançada do YOLO27 para ver a velocidade de classificação preliminar e os tamanhos dos modelos.

Treinar#

Treine YOLO26n-cls no conjunto de dados MNIST160 por 100 épocas com tamanho de imagem 64. Para obter uma lista completa dos argumentos disponíveis, consulte a página de Configuração.

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="imagenet", trainer=CustomizedTrainer, epochs=10, imgsz=224, batch=64)
model.val(data="imagenet", validator=CustomizedValidator, imgsz=224, batch=64)

Formato do dataset#

O formato de conjunto de dados de classificação YOLO é descrito em detalhes no Guia de conjuntos de dados. Os conjuntos de dados de classificação também podem ser gerenciados e rotulados com as ferramentas de anotação da Ultralytics Platform.

Validação#

Valide a precisão do modelo YOLO26n-cls treinado no conjunto de dados MNIST160. Nenhum argumento é necessário, pois o model mantém seu data de treinamento e seus argumentos 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

Previsão#

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 Predict.

Saída dos resultados#

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

Para consultar os campos Results específicos da tarefa em todas as tarefas, vê a secção Resultados de Previsão por Tarefa.

Exportação#

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 YOLO26-cls estão na tabela abaixo. Você pode exportar para qualquer formato usando o argumento format, ou seja, format='onnx' ou format='engine'. Você pode prever ou validar diretamente em modelos exportados, ou seja, yolo predict model=yolo26n-cls.onnx. Os exemplos de uso são exibidos para o seu modelo após a conclusão da exportação.

nms=None por padrão utiliza saídas brutas para NMS externo. Define nms=False para selecionar uma cabeça livre de NMS disponível; os formatos não suportados recorrem ao seu caminho de saída nativo. As entradas nms acima identificam formatos que podem incorporar NMS com nms=True.

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

Perguntas frequentes#