Ir al contenido

Cómo entrenar YOLO con COCO JSON sin convertir

¿Por qué entrenar directamente con COCO JSON?

Anotaciones en COCO JSON el formato se puede usar directamente para Ultralytics YOLO entrenamiento sin convertir a .txt archivos primero. Esto se logra mediante la subclase de YOLODataset para analizar COCO JSON sobre la marcha e integrarlo en el pipeline de entrenamiento a través de un entrenador personalizado.

Este enfoque mantiene el COCO JSON como la única fuente de verdad — sin convert_coco() llamada, sin reorganización de directorios, sin archivos de etiquetas intermedios. YOLO26 y todos los demás modelos de detección de Ultralytics YOLO son compatibles. Los modelos de segmentación y pose requieren campos de etiqueta adicionales (ver Preguntas frecuentes).

¿Busca una conversión única en su lugar?

Consulte la Guía de conversión de COCO a YOLO para el flujo de trabajo estándar convert_coco() flujo de trabajo.

Descripción General de la Arquitectura

Se necesitan dos clases:

  1. COCODataset — lee COCO JSON y convierte cajas delimitadoras a formato YOLO en memoria durante el entrenamiento
  2. COCOTrainer — anula build_dataset() para usar COCODataset en lugar del predeterminado YOLODataset

La implementación sigue el mismo patrón que el integrado GroundingDataset, que también lee anotaciones JSON directamente. Se sobrescriben tres métodos: get_img_files(), cache_labels(), y get_labels().

Construyendo la clase de conjunto de datos COCO JSON

El COCODataset la clase hereda de YOLODataset y sobrescribe la lógica de carga de etiquetas. En lugar de leer .txt archivos de un directorio de etiquetas, abre el archivo JSON de COCO, itera sobre las anotaciones agrupadas por imagen y convierte cada cuadro delimitador del formato de píxeles de COCO [x_min, y_min, width, height] al formato de centro normalizado de YOLO [x_center, y_center, width, height]. Las anotaciones de multitud (iscrowd: 1) y los cuadros de área cero se omiten automáticamente.

El get_img_files() el método devuelve una lista vacía porque las rutas de las imágenes se resuelven a partir del JSON file_name campo dentro de cache_labels(). Los IDs de categoría se ordenan y se reasignan a índices de clase basados en cero, por lo que tanto los esquemas de ID basados en 1 (COCO estándar) como los no contiguos funcionan correctamente.

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"]

Las etiquetas analizadas se guardan en un .cache archivo junto al JSON (por ejemplo, instances_train.cache). En ejecuciones de entrenamiento posteriores, la caché se carga directamente, omitiendo el análisis JSON. Si el archivo JSON cambia, la verificación de hash falla y la caché se reconstruye automáticamente.

Conectando el conjunto de datos al pipeline de entrenamiento

El único cambio necesario en el entrenador es sobrescribir build_dataset(). El predeterminado DetectionTrainer construye un YOLODataset que escanea en busca de .txt archivos de etiquetas. Al reemplazarlo con COCODataset, el entrenador lee del JSON de COCO en su lugar.

La ruta del archivo JSON se obtiene de un train_json / val_json campo personalizado en la configuración de datos (ver Paso 3). Durante el entrenamiento, mode="train" se resuelve a train_json; durante la validación, mode="val" se resuelve a val_json. Si val_json no está configurado, se recurre a 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,
        )

Configurando dataset.yaml para COCO JSON

El dataset.yaml utiliza los campos estándar de path, train, y val para localizar directorios de imágenes. Dos campos adicionales, train_json y val_json, especifican los archivos de anotación COCO que COCOTrainer lee. Los campos nc y names definen el número de clases y sus nombres, coincidiendo con el orden alfabético de categories en el 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

Estructura de directorio esperada:

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

Entrenamiento con COCO JSON

Con la clase de conjunto de datos, la clase de entrenador y la configuración YAML implementadas, el entrenamiento se realiza mediante la llamada estándar a model.train() . La única diferencia con una ejecución de entrenamiento normal es el argumento trainer=COCOTrainer , que indica a Ultralytics que utilice el cargador de conjunto de datos personalizado en lugar del predeterminado.

from ultralytics import YOLO

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

El pipeline completo de entrenamiento se ejecuta según lo esperado, incluyendo la validación, el guardado de puntos de control y el registro de métricas.

Implementación completa

Para mayor comodidad, la implementación completa se proporciona a continuación como un script único de copiar y pegar. Incluye el conjunto de datos personalizado, el entrenador personalizado y la llamada de entrenamiento. Guárdelo junto con su archivo dataset.yaml y ejecútelo directamente.

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)

Para recomendaciones de hiperparámetros, consulte la guía de Consejos para el Entrenamiento de Modelos.

Preguntas frecuentes

¿Cuál es la diferencia entre esto y convert_coco()?

convert_coco() escribe .txt archivos de etiquetas en disco como una conversión única. Este enfoque analiza el JSON al inicio de cada ejecución de entrenamiento y convierte las anotaciones en memoria. Utilice convert_coco() cuando se prefieren etiquetas permanentes en formato YOLO; use este enfoque para mantener el JSON de COCO como la única fuente de verdad sin generar archivos adicionales.

¿Puede YOLO entrenar con COCO JSON sin código personalizado?

No con el pipeline actual de Ultralytics, que espera etiquetas YOLO .txt por defecto. Esta guía proporciona el código personalizado mínimo necesario: una clase de conjunto de datos y una clase de entrenador. Una vez definidos, el entrenamiento solo requiere una llamada estándar a model.train() .

¿Soporta esto segmentación y estimación de pose?

Esta guía cubre detección de objetos. Para añadir soporte para segmentación de instancias , incluya el segmentation datos de polígono de las anotaciones COCO en el segments campo de cada diccionario de etiquetas. Para estimación de pose, incluya keypoints. El GroundingDataset código fuente proporciona una implementación de referencia para el manejo de segments.

¿Funcionan las aumentaciones con este conjunto de datos personalizado?

Sí. COCODataset extiende YOLODataset, por lo que todas las aumentos de datosmosaico, mixup, copiar-pegar, y otros — se ejecutan sin modificaciones.

¿Cómo se mapean los IDs de categoría a los índices de clase?

Las categorías se ordenan por id y se mapean a índices secuenciales que comienzan desde 0. Esto maneja IDs basados en 1 (COCO estándar), IDs basados en 0 e IDs no contiguos. El names diccionario en dataset.yaml debe seguir el mismo orden de clasificación que el categories array COCO.

¿Existe una sobrecarga de rendimiento en comparación con las etiquetas preconvertidas?

El JSON COCO se analiza una vez en la primera ejecución de entrenamiento. Las etiquetas analizadas se guardan en un .cache archivo, por lo que las ejecuciones posteriores se cargan instantáneamente sin necesidad de volver a analizar. La velocidad de entrenamiento es idéntica a la del entrenamiento YOLO estándar, ya que las anotaciones se mantienen en memoria. La caché se reconstruye automáticamente si el archivo JSON cambia.



📅 Creado hace 7 días ✏️ Actualizado hace 7 días
raimbekovmglenn-jocher

Comentarios