Passer au contenu

Comment entraîner YOLO les données COCO sans conversion

Pourquoi s'entraîner directement sur COCO ?

Annotations dans COCO JSON ce format peut être utilisé directement pour Ultralytics YOLO s'entraîner sans se convertir à .txt les fichiers en premier. Pour ce faire, on crée une sous-classe YOLODataset pour analyser COCO à la volée et l'intégrer dans le pipeline d'entraînement via un modèle d'entraînement personnalisé.

Cette approche fait du COCO la seule source de vérité — non convert_coco() appel, aucune réorganisation de répertoire, aucun fichier d'étiquettes intermédiaire. YOLO26 et tous les autres modèles de détection Ultralytics YOLO sont pris en charge. Les modèles de segmentation et de pose nécessitent des champs d'étiquettes supplémentaires (voir FAQ).

Vous préférez plutôt une conversion ponctuelle ?

Consultez le fichier Guide de conversion COCO vers YOLO pour la norme convert_coco() flux de travail.

Aperçu de l'architecture

Deux classes sont nécessaires :

  1. COCODataset — lit COCO et le convertit boîtes englobantes YOLO en mémoire pendant l'entraînement
  2. COCOTrainer — remplace build_dataset() à utiliser COCODataset au lieu de la valeur par défaut YOLODataset

La mise en œuvre suit le même schéma que la fonction intégrée GroundingDataset, qui lit également les annotations JSON directement. Trois méthodes sont surchargées : get_img_files(), cache_labels(), et get_labels().

Construction de la classe de jeu de données JSON COCO

L'argument COCODataset La classe hérite de YOLODataset et surcharge la logique de chargement des étiquettes. Au lieu de lire .txt à partir d'un répertoire « labels », il ouvre le fichier COCO , parcourt les annotations regroupées par image et convertit chaque cadre de sélection du format COCO [x_min, y_min, width, height] au format YOLO par rapport au centre [x_center, y_center, width, height]. Annotations de foule (iscrowd: 1) et les boîtes de surface nulle sont automatiquement ignorées.

L'argument get_img_files() La méthode renvoie une liste vide car les chemins d'accès aux images sont déterminés à partir du JSON file_name champ à l'intérieur cache_labels(). Les ID de catégorie sont triés et remappés en indices de classe à base zéro, de sorte que les schémas d'ID basés sur 1 (COCO standard) et non contigus fonctionnent correctement.

import json
from collections import defaultdict
from pathlib import Path

import numpy as np

from ultralytics.data.dataset import DATASET_CACHE_VERSION, YOLODataset
from ultralytics.data.utils import get_hash, load_dataset_cache_file, save_dataset_cache_file
from ultralytics.utils import TQDM


class COCODataset(YOLODataset):
    """Dataset that reads COCO JSON annotations directly without conversion to .txt files."""

    def __init__(self, *args, json_file="", **kwargs):
        self.json_file = json_file
        super().__init__(*args, data={"channels": 3}, **kwargs)

    def get_img_files(self, img_path):
        """Image paths are resolved from the JSON file, not from scanning a directory."""
        return []

    def cache_labels(self, path=Path("./labels.cache")):
        """Parse COCO JSON and convert annotations to YOLO format. Results are saved to a .cache file."""
        x = {"labels": []}
        with open(self.json_file) as f:
            coco = json.load(f)

        images = {img["id"]: img for img in coco["images"]}

        # Sort categories by ID and map to 0-indexed classes
        categories = {cat["id"]: i for i, cat in enumerate(sorted(coco["categories"], key=lambda c: c["id"]))}

        img_to_anns = defaultdict(list)
        for ann in coco["annotations"]:
            img_to_anns[ann["image_id"]].append(ann)

        for img_info in TQDM(coco["images"], desc="reading annotations"):
            h, w = img_info["height"], img_info["width"]
            im_file = Path(self.img_path) / img_info["file_name"]
            if not im_file.exists():
                continue

            self.im_files.append(str(im_file))
            bboxes = []
            for ann in img_to_anns.get(img_info["id"], []):
                if ann.get("iscrowd", False):
                    continue
                # COCO: [x, y, w, h] top-left in pixels -> YOLO: [cx, cy, w, h] center normalized
                box = np.array(ann["bbox"], dtype=np.float32)
                box[:2] += box[2:] / 2  # top-left to center
                box[[0, 2]] /= w  # normalize x
                box[[1, 3]] /= h  # normalize y
                if box[2] <= 0 or box[3] <= 0:
                    continue
                cls = categories[ann["category_id"]]
                bboxes.append([cls, *box.tolist()])

            lb = np.array(bboxes, dtype=np.float32) if bboxes else np.zeros((0, 5), dtype=np.float32)
            x["labels"].append(
                {
                    "im_file": str(im_file),
                    "shape": (h, w),
                    "cls": lb[:, 0:1],
                    "bboxes": lb[:, 1:],
                    "segments": [],
                    "normalized": True,
                    "bbox_format": "xywh",
                }
            )
        x["hash"] = get_hash([self.json_file, str(self.img_path)])
        save_dataset_cache_file(self.prefix, path, x, DATASET_CACHE_VERSION)
        return x

    def get_labels(self):
        """Load labels from .cache file if available, otherwise parse JSON and create the cache."""
        cache_path = Path(self.json_file).with_suffix(".cache")
        try:
            cache = load_dataset_cache_file(cache_path)
            assert cache["version"] == DATASET_CACHE_VERSION
            assert cache["hash"] == get_hash([self.json_file, str(self.img_path)])
            self.im_files = [lb["im_file"] for lb in cache["labels"]]
        except (FileNotFoundError, AssertionError, AttributeError, KeyError, ModuleNotFoundError):
            cache = self.cache_labels(cache_path)
        cache.pop("hash", None)
        cache.pop("version", None)
        return cache["labels"]

Les étiquettes analysées sont enregistrées dans un .cache fichier situé à côté du fichier JSON (par exemple instances_train.cache). Lors des exécutions d'entraînement ultérieures, le cache est chargé directement, ignorant l'analyse JSON. Si le fichier JSON change, la vérification du hachage échoue et le cache est automatiquement reconstruit.

Connexion du jeu de données au pipeline d'entraînement

La seule modification à apporter au programme d'entraînement consiste à remplacer build_dataset(). Le défaut DetectionTrainer construit un YOLODataset qui recherche .txt fichiers d'étiquettes. En le remplaçant par COCODataset, l'entraîneur lit plutôt le JSON COCO.

Le chemin d'accès au fichier JSON est extrait d'un fichier personnalisé train_json / val_json champ dans la configuration des données (voir l'étape 3). Pendant l'entraînement, mode="train" décide de train_json; lors de la validation, mode="val" décide de val_json. Si val_json s'il n'est pas défini, le système utilise par défaut train_json.

from ultralytics.models.yolo.detect import DetectionTrainer
from ultralytics.utils import colorstr


class COCOTrainer(DetectionTrainer):
    """Trainer that uses COCODataset for direct COCO JSON training."""

    def build_dataset(self, img_path, mode="train", batch=None):
        json_file = self.data["train_json"] if mode == "train" else self.data.get("val_json", self.data["train_json"])
        return COCODataset(
            img_path=img_path,
            json_file=json_file,
            imgsz=self.args.imgsz,
            batch_size=batch,
            augment=mode == "train",
            hyp=self.args,
            rect=self.args.rect or mode == "val",
            cache=self.args.cache or None,
            single_cls=self.args.single_cls or False,
            stride=int(self.model.stride.max()) if hasattr(self, "model") and self.model else 32,
            pad=0.0 if mode == "train" else 0.5,
            prefix=colorstr(f"{mode}: "),
            task=self.args.task,
            classes=self.args.classes,
            fraction=self.args.fraction if mode == "train" else 1.0,
        )

Configuration de dataset.yaml pour COCO JSON

L'argument dataset.yaml utilise la norme path, train, et val champs permettant de localiser les répertoires d'images. Deux champs supplémentaires, train_json et val_json, spécifiez les fichiers d'annotation COCO qui COCOTrainer dit. Le nc et names Ces champs définissent le nombre de classes et leurs noms, en respectant l'ordre trié de categories dans le fichier JSON.

path: /path/to/images # root directory with train/ and val/ subfolders
train: train
val: val

# COCO JSON annotation files
train_json: /path/to/annotations/instances_train.json
val_json: /path/to/annotations/instances_val.json

nc: 80
names:
    0: person
    1: bicycle
    # ... remaining class names

Structure de répertoires prévue :

my_dataset/
  images/
    train/
      img_001.jpg
      ...
    val/
      img_100.jpg
      ...
  annotations/
    instances_train.json
    instances_val.json
  dataset.yaml

Formation à l'utilisation de COCO

Une fois le jeu de données, la classe d'entraînement et la configuration YAML en place, l'entraînement se déroule selon la procédure standard model.train() appel. La seule différence par rapport à une exécution d'entraînement normale est le trainer=COCOTrainer argument, qui indique à Ultralytics d'utiliser le chargeur de jeu de données personnalisé au lieu de celui par défaut.

from ultralytics import YOLO

model = YOLO("yolo26n.pt")
model.train(data="dataset.yaml", epochs=100, imgsz=640, trainer=COCOTrainer)

L'ensemble du processus d'entraînement se déroule comme prévu, y compris la validation, l'enregistrement des points de contrôle et la journalisation des métriques.

Mise en œuvre complète

Pour plus de commodité, vous trouverez ci-dessous l'implémentation complète sous la forme d'un script unique prêt à copier-coller. Elle comprend l'ensemble de données personnalisé, le modèle d'entraînement personnalisé et l'appel d'entraînement. Enregistrez ce fichier à côté de votre dataset.yaml et exécutez-le directement.

import json
from collections import defaultdict
from pathlib import Path

import numpy as np

from ultralytics import YOLO
from ultralytics.data.dataset import DATASET_CACHE_VERSION, YOLODataset
from ultralytics.data.utils import get_hash, load_dataset_cache_file, save_dataset_cache_file
from ultralytics.models.yolo.detect import DetectionTrainer
from ultralytics.utils import TQDM, colorstr


class COCODataset(YOLODataset):
    """Dataset that reads COCO JSON annotations directly without conversion to .txt files."""

    def __init__(self, *args, json_file="", **kwargs):
        self.json_file = json_file
        super().__init__(*args, data={"channels": 3}, **kwargs)

    def get_img_files(self, img_path):
        return []

    def cache_labels(self, path=Path("./labels.cache")):
        x = {"labels": []}
        with open(self.json_file) as f:
            coco = json.load(f)

        images = {img["id"]: img for img in coco["images"]}
        categories = {cat["id"]: i for i, cat in enumerate(sorted(coco["categories"], key=lambda c: c["id"]))}

        img_to_anns = defaultdict(list)
        for ann in coco["annotations"]:
            img_to_anns[ann["image_id"]].append(ann)

        for img_info in TQDM(coco["images"], desc="reading annotations"):
            h, w = img_info["height"], img_info["width"]
            im_file = Path(self.img_path) / img_info["file_name"]
            if not im_file.exists():
                continue

            self.im_files.append(str(im_file))
            bboxes = []
            for ann in img_to_anns.get(img_info["id"], []):
                if ann.get("iscrowd", False):
                    continue
                box = np.array(ann["bbox"], dtype=np.float32)
                box[:2] += box[2:] / 2
                box[[0, 2]] /= w
                box[[1, 3]] /= h
                if box[2] <= 0 or box[3] <= 0:
                    continue
                cls = categories[ann["category_id"]]
                bboxes.append([cls, *box.tolist()])

            lb = np.array(bboxes, dtype=np.float32) if bboxes else np.zeros((0, 5), dtype=np.float32)
            x["labels"].append(
                {
                    "im_file": str(im_file),
                    "shape": (h, w),
                    "cls": lb[:, 0:1],
                    "bboxes": lb[:, 1:],
                    "segments": [],
                    "normalized": True,
                    "bbox_format": "xywh",
                }
            )
        x["hash"] = get_hash([self.json_file, str(self.img_path)])
        save_dataset_cache_file(self.prefix, path, x, DATASET_CACHE_VERSION)
        return x

    def get_labels(self):
        cache_path = Path(self.json_file).with_suffix(".cache")
        try:
            cache = load_dataset_cache_file(cache_path)
            assert cache["version"] == DATASET_CACHE_VERSION
            assert cache["hash"] == get_hash([self.json_file, str(self.img_path)])
            self.im_files = [lb["im_file"] for lb in cache["labels"]]
        except (FileNotFoundError, AssertionError, AttributeError, KeyError, ModuleNotFoundError):
            cache = self.cache_labels(cache_path)
        cache.pop("hash", None)
        cache.pop("version", None)
        return cache["labels"]


class COCOTrainer(DetectionTrainer):
    """Trainer that uses COCODataset for direct COCO JSON training."""

    def build_dataset(self, img_path, mode="train", batch=None):
        json_file = self.data["train_json"] if mode == "train" else self.data.get("val_json", self.data["train_json"])
        return COCODataset(
            img_path=img_path,
            json_file=json_file,
            imgsz=self.args.imgsz,
            batch_size=batch,
            augment=mode == "train",
            hyp=self.args,
            rect=self.args.rect or mode == "val",
            cache=self.args.cache or None,
            single_cls=self.args.single_cls or False,
            stride=int(self.model.stride.max()) if hasattr(self, "model") and self.model else 32,
            pad=0.0 if mode == "train" else 0.5,
            prefix=colorstr(f"{mode}: "),
            task=self.args.task,
            classes=self.args.classes,
            fraction=self.args.fraction if mode == "train" else 1.0,
        )


model = YOLO("yolo26n.pt")
model.train(data="dataset.yaml", epochs=100, imgsz=640, trainer=COCOTrainer)

Pour obtenir des recommandations sur les hyperparamètres, consultez le guide « Conseils pour l'entraînement des modèles ».

FAQ

Quelle est la différence entre cette fonction et convert_coco() ?

convert_coco() écrit .txt enregistrer les fichiers d'étiquettes sur le disque lors d'une conversion unique. Cette approche analyse le fichier JSON au début de chaque session d'entraînement et convertit les annotations en mémoire. Utilisez convert_coco() lorsque l'on privilégie les étiquettes YOLO permanent ; utilisez cette approche pour que le fichier COCO reste la seule source de référence, sans générer de fichiers supplémentaires.

YOLO peut-il s'entraîner sur du JSON COCO sans code personnalisé ?

Ce n'est pas le cas avec le Ultralytics actuel Ultralytics , qui s'appuie sur YOLO .txt étiquettes par défaut. Ce guide fournit le code personnalisé minimal nécessaire : une classe de jeu de données et une classe d'entraîneur. Une fois ces éléments définis, l'entraînement ne nécessite qu'un model.train() appel.

Cela prend-il en charge la segmentation et l'estimation de pose ?

Ce guide traite détection d'objets. Pour ajouter segmentation d'instance prise en charge, inclure le segmentation données polygonales issues COCO dans le segments champ de chaque dictionnaire d'étiquettes. Pour estimation de pose, inclure keypoints. Le GroundingDataset code source fournit une implémentation de référence pour la gestion des segments.

Les augmentations fonctionnent-elles avec ce jeu de données personnalisé ?

Oui. COCODataset s'étend YOLODataset, donc toutes les méthodes intégrées augmentations de données — mosaïque, mixup, copier-coller, et d'autres — s'exécutent sans modification.

Comment les identifiants de catégorie sont-ils mis en correspondance avec les indices de classe ?

Les catégories sont triées par id et mappés à des indices séquentiels commençant à 0. Cela gère les ID basés sur 1 (COCO standard), les ID basés sur 0 et les ID non contigus. Le names dictionnaire dans dataset.yaml devrait respecter le même ordre de tri que le COCO categories tableau.

Y a-t-il une perte de performances par rapport aux étiquettes pré-converties ?

Le fichier COCO est analysé une seule fois lors du premier cycle d'entraînement. Les étiquettes analysées sont enregistrées dans un .cache fichier, de sorte que les exécutions suivantes se chargent instantanément sans nouvelle analyse. La vitesse d'entraînement est identique à celle de YOLO standard, car les annotations sont conservées en mémoire. Le cache est automatiquement reconstitué si le fichier JSON est modifié.



📅 Créé il y a 7 jours ✏️ Mis à jour il y a 7 jours
raimbekovmglenn-jocher

Commentaires