Vai al contenuto

Come addestrare YOLO COCO senza conversione

Perché esercitarsi direttamente su COCO

Annotazioni in COCO il formato può essere utilizzato direttamente per Ultralytics YOLO allenarsi senza passare a .txt prima i file. Ciò avviene creando una sottoclasse YOLODataset per analizzare COCO in tempo reale e integrarli nella pipeline di addestramento tramite un trainer personalizzato.

Questo approccio mantiene il COCO come unica fonte attendibile — nessun convert_coco() chiamata, nessuna riorganizzazione della directory, nessun file di etichetta intermedio. YOLO26 e sono supportati tutti gli altri modelliYOLO Ultralytics . I modelli di segmentazione e di posa richiedono campi di etichetta aggiuntivi (vedi FAQ).

Preferisci invece una conversione una tantum?

Vedere il file Guida YOLO COCO YOLO per lo standard convert_coco() flusso di lavoro.

Panoramica dell'architettura

Sono necessarie due classi:

  1. COCOJSONDataset — legge COCO e li converte bounding box YOLO in memoria durante l'addestramento
  2. COCOJSONTrainer — sovrascrizioni build_dataset() da usare COCOJSONDataset anziché quello predefinito YOLODataset

L'implementazione segue lo stesso schema della funzione integrata GroundingDataset, che legge direttamente anche le annotazioni JSON. Sono stati sovrascritti tre metodi: get_img_files(), cache_labels(), e get_labels().

Creazione della classe del set di dati COCO

Il COCOJSONDataset la classe eredita da YOLODataset e sovrascrive la logica di caricamento delle etichette. Invece di leggere .txt i file presenti nella directory "labels", apre il file COCO , esegue un'iterazione sulle annotazioni raggruppate per immagine e converte ogni riquadro di delimitazione dal formato COCO [x_min, y_min, width, height] nel formato YOLO al centro [x_center, y_center, width, height]. Annotazioni della comunità (iscrowd: 1) e le caselle con area nulla vengono saltate automaticamente.

Il get_img_files() Il metodo restituisce una lista vuota perché i percorsi delle immagini vengono ricavati dal JSON file_name all'interno del campo cache_labels(). Gli ID delle categorie vengono ordinati e rimappati in indici di classe con indicizzazione a partire da zero, in modo che funzionino correttamente sia gli schemi di ID con indicizzazione a partire da 1 ( COCO standard) sia quelli non contigui.

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

Le etichette analizzate vengono salvate in un .cache file accanto al JSON (ad es. instances_train.cache). Nelle successive sessioni di addestramento, la cache viene caricata direttamente, saltando l'analisi del JSON. Se il file JSON viene modificato, il controllo dell'hash fallisce e la cache viene ricostruita automaticamente.

Collegamento del set di dati alla pipeline di addestramento

L'unica modifica necessaria nel trainer consiste nel sovrascrivere build_dataset(). L'impostazione predefinita DetectionTrainer crea un YOLODataset che ricerca .txt file di etichette. Sostituendolo con COCOJSONDataset, il trainer legge invece dal file COCO .

Il percorso del file JSON viene ricavato da un file personalizzato train_json / val_json field in the data config (see Step 3). During training, mode="train" decide di train_json; durante la convalida, mode="val" decide di val_json. Se val_json se non è impostato, si ricorre a train_json.

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


class COCOJSONTrainer(DetectionTrainer):
    """Trainer that uses COCOJSONDataset 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 COCOJSONDataset(
            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,
        )

Configurazione del file dataset.yaml COCO

Il dataset.yaml utilizza lo standard path, train, e val campi per individuare le directory delle immagini. Due campi aggiuntivi, train_json e val_json, specificare i file COCO che COCOJSONTrainer si legge. Il nc e names I campi definiscono il numero delle classi e i loro nomi, seguendo l'ordine in cui sono elencati categories nel 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

Struttura di directory prevista:

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

Formazione sulla gestione di COCO

Una volta preparati il set di dati, il modello di addestramento e il file di configurazione YAML, l'addestramento procede secondo la procedura standard model.train() chiamata. L'unica differenza rispetto a una normale sessione di allenamento è la trainer=COCOJSONTrainer argomento che indica Ultralytics utilizzare il caricatore di set di dati personalizzato anziché quello predefinito.

from ultralytics import YOLO

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

L'intero flusso di addestramento funziona come previsto, comprese la convalida, il salvataggio dei checkpoint e la registrazione delle metriche.

Piena attuazione

Per comodità, di seguito è riportata l'implementazione completa sotto forma di unico script da copiare e incollare. Include il set di dati personalizzato, il trainer personalizzato e la chiamata di addestramento. Salva questo file insieme al tuo dataset.yaml e eseguirlo direttamente.

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 COCOJSONDataset(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 COCOJSONTrainer(DetectionTrainer):
    """Trainer that uses COCOJSONDataset 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 COCOJSONDataset(
            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=COCOJSONTrainer)

Per suggerimenti sugli iperparametri, consultare la guida " Suggerimenti per l'addestramento dei modelli ".

FAQ

Qual è la differenza tra questa funzione e convert_coco()?

convert_coco() scrive .txt salvare i file delle etichette su disco come conversione una tantum. Questo approccio analizza il JSON all'inizio di ogni sessione di addestramento e converte le annotazioni in memoria. Utilizza convert_coco() quando si preferisce utilizzare etichette permanenti YOLO; utilizzare questo approccio per mantenere il file COCO come unica fonte attendibile senza generare file aggiuntivi.

È possibile YOLO sui COCO senza codice personalizzato?

Non con l'attuale Ultralytics , che prevede l'uso di YOLO .txt etichette predefinite. Questa guida fornisce il codice personalizzato minimo necessario: una classe per il set di dati e una classe per il modello di addestramento. Una volta definiti, l'addestramento richiede solo un model.train() chiamata.

Supporta la segmentazione e la stima della posa?

Questa guida tratta il rilevamento di oggetti. Per aggiungere segmentazione delle istanze supporto, includere il segmentation dati poligonali provenienti dalle COCO nel segments campo di ciascun dizionario di etichette. Per stima della posa, tra cui keypoints. GroundingDataset codice sorgente fornisce un'implementazione di riferimento per la gestione dei segmenti.

Le estensioni funzionano con questo set di dati personalizzato?

Sì. COCOJSONDataset si estende YOLODataset, quindi tutte le funzioni integrate ampliamenti dei dati — mosaic, mixup, copia-incolla, e altri — funzionano senza modifiche.

In che modo gli ID delle categorie vengono associati agli indici delle classi?

Le categorie sono ordinate per id e mappati su indici sequenziali a partire da 0. In questo modo si gestiscono gli ID a partire da 1 (standard COCO), gli ID a partire da 0 e gli ID non contigui. Il names dizionario in dataset.yaml dovrebbe seguire lo stesso ordine di ordinamento del COCO categories matrice.

Si riscontra un calo di prestazioni rispetto alle etichette pre-convertite?

Il file COCO viene analizzato una volta durante il primo ciclo di addestramento. Le etichette analizzate vengono salvate in un .cache file, quindi le successive esecuzioni vengono caricate istantaneamente senza bisogno di una nuova analisi. La velocità di addestramento è identica a YOLO standard, poiché le annotazioni sono conservate in memoria. La cache viene ricostruita automaticamente se il file JSON viene modificato.



📅 Creato 0 giorni fa ✏️ Aggiornato 0 giorni fa
glenn-jocherraimbekovm

Commenti