Classification d'images avec Ultralytics YOLO#

La classification d'images est la plus simple des tâches prises en charge et consiste à classifier une image entière dans l'une d'un ensemble de classes prédéfinies.

La sortie d'un classificateur d'images est une étiquette de classe unique et un score de confiance. La classification d'images est utile lorsque tu as seulement besoin de connaître la classe à laquelle appartient une image et que tu n'as pas besoin de savoir où se trouvent les objets de cette classe ni quelle est leur forme exacte.

Modèles#

Les modèles de classification pré-entraînés YOLO26 sont présentés ici. Les modèles de détection, de segmentation et de pose sont pré-entraînés sur le jeu de données COCO, les modèles Semantic sont pré-entraînés sur Cityscapes, et les modèles de classification sont pré-entraînés sur le jeu de données ImageNet.

Les modèles sont téléchargés automatiquement depuis la dernière version d'Ultralytics lors de la première utilisation.

• Les valeurs acc correspondent aux précisions des modèles sur l'ensemble de validation du jeu de données ImageNet.
  Reproduis avec yolo val classify data=path/to/ImageNet device=0
• Vitesse moyennée sur les images de validation d'ImageNet en utilisant une instance Amazon EC2 P4d.
  Reproduis avec yolo val classify data=path/to/ImageNet batch=1 device=0|cpu
• Les valeurs de Params et FLOPs concernent le modèle fusionné après model.fuse(), qui fusionne les couches Conv et BatchNorm. Les points de contrôle pré-entraînés conservent l'architecture d'entraînement complète et peuvent afficher des nombres plus élevés.

Entraîner (Train)#

Entraîne YOLO26n-cls sur le jeu de données MNIST160 pendant 100 époques pour une taille d'image de 64. Pour une liste complète des arguments disponibles, consulte la page 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)

Format de jeu de données#

Le format du jeu de données de classification YOLO est détaillé dans le Guide des jeux de données. Les jeux de données de classification peuvent également être gérés et étiquetés avec les outils d'annotation d'Ultralytics Platform.

Valider (Val)#

Valide la précision du modèle de classification YOLO26n-cls entraîné sur le jeu de données MNIST160. Aucun argument n'est nécessaire car le model conserve son data et ses arguments d'entraînement en tant qu'attributs du modèle.

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

Prédire (Predict)#

Utilise un modèle YOLO26n-cls entraîné pour exécuter des prédictions sur des images.

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

Consulte tous les détails du mode predict dans la page Prédire.

Sortie des résultats#

La classification d'images renvoie un objet Results par image. Le champ de prédiction principal est result.probs, qui contient le vecteur de probabilité des classes et des assistants pour les meilleures prédictions.

Pour les champs Results spécifiques à chaque tâche, consulte la section Résultats de prédiction par tâche.

Exporter (Export)#

Exporte un modèle YOLO26n-cls vers un format différent tel que 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")

Les formats d'exportation disponibles pour YOLO26-cls se trouvent dans le tableau ci-dessous. Tu peux exporter vers n'importe quel format en utilisant l'argument format, c'est-à-dire format='onnx' ou format='engine'. Tu peux prédire ou valider directement sur des modèles exportés, c'est-à-dire yolo predict model=yolo26n-cls.onnx. Des exemples d'utilisation sont affichés pour ton modèle une fois l'exportation terminée.

Consulte tous les détails de export dans la page Exporter.

FAQ#