Come addestrare YOLO su COCO JSON senza convertire
Perché addestrare direttamente su COCO JSON
Le annotazioni in formato COCO JSON possono essere utilizzate direttamente per l'addestramento di Ultralytics YOLO senza convertire prima i file in .txt. Questo si ottiene creando una sottoclasse di YOLODataset per analizzare COCO JSON al volo e collegandola alla pipeline di addestramento tramite un trainer personalizzato.
Questo approccio mantiene il COCO JSON come unica fonte di verità: nessuna chiamata convert_coco(), nessuna riorganizzazione delle directory, nessun file di etichetta intermedio. YOLO26 e tutti gli altri modelli di rilevamento di Ultralytics YOLO sono supportati. I modelli di segmentazione e posa richiedono campi di etichetta aggiuntivi (vedi FAQ).
Consulta la guida alla conversione da COCO a YOLO per il flusso di lavoro standard convert_coco().
Panoramica dell'architettura
Sono necessarie due classi:
COCODataset: legge il COCO JSON e converte i bounding box nel formato YOLO in memoria durante l'addestramentoCOCOTrainer: sovrascrivebuild_dataset()per utilizzareCOCODatasetal posto delYOLODatasetpredefinito
L'implementazione segue lo stesso schema del GroundingDataset integrato, che legge anch'esso le annotazioni JSON direttamente. Vengono sovrascritti tre metodi: get_img_files(), cache_labels() e get_labels().
Creazione della classe Dataset per COCO JSON
La classe COCODataset eredita da YOLODataset e sovrascrive la logica di caricamento delle etichette. Invece di leggere i file .txt da una directory di etichette, apre il file COCO JSON, scorre le annotazioni raggruppate per immagine e converte ogni bounding box dal formato pixel COCO [x_min, y_min, width, height] al formato normalizzato centro YOLO [x_center, y_center, width, height]. Le annotazioni di gruppo (iscrowd: 1) e le box con area zero vengono saltate automaticamente.
Il metodo get_img_files() restituisce un elenco vuoto perché i percorsi delle immagini vengono risolti dal campo file_name del JSON all'interno di cache_labels(). Gli ID delle categorie vengono ordinati e rimappati in indici di classe basati sullo zero, quindi funzionano correttamente sia gli schemi basati su 1 (COCO standard) che 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 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"]Le etichette analizzate vengono salvate in un file .cache accanto al JSON (es. instances_train.cache). Nelle sessioni di addestramento successive, la cache viene caricata direttamente, saltando l'analisi del JSON. Se il file JSON cambia, il controllo hash fallisce e la cache viene ricostruita automaticamente.
Collegamento del Dataset alla pipeline di addestramento
L'unica modifica necessaria nel trainer è sovrascrivere build_dataset(). Il DetectionTrainer predefinito crea un YOLODataset che scansiona i file di etichette .txt. Sostituendolo con COCODataset, il trainer legge invece dal COCO JSON.
Il percorso del file JSON viene estratto da un campo personalizzato train_json / val_json nella configurazione dei dati (vedi Passaggio 3). Durante l'addestramento, mode="train" viene risolto in train_json; durante la convalida, mode="val" viene risolto in val_json. Se val_json non è impostato, ripiega su 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,
)Configurazione di dataset.yaml per COCO JSON
Il dataset.yaml utilizza i campi standard path, train e val per individuare le directory delle immagini. Due campi aggiuntivi, train_json e val_json, specificano i file di annotazione COCO che COCOTrainer legge. I campi nc e names definiscono il numero di classi e i loro nomi, corrispondendo all'ordine ordinato delle 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 namesStruttura della directory prevista:
my_dataset/
images/
train/
img_001.jpg
...
val/
img_100.jpg
...
annotations/
instances_train.json
instances_val.json
dataset.yamlEsecuzione dell'addestramento su COCO JSON
Con la classe del dataset, la classe del trainer e la configurazione YAML pronte, l'addestramento funziona tramite la chiamata standard model.train(). L'unica differenza rispetto a una normale sessione di addestramento è l'argomento trainer=COCOTrainer, che indica a Ultralytics di utilizzare il caricatore di dataset personalizzato invece di quello predefinito.
from ultralytics import YOLO
model = YOLO("yolo26n.pt")
model.train(data="dataset.yaml", epochs=100, imgsz=640, trainer=COCOTrainer)L'intera pipeline di addestramento viene eseguita come previsto, inclusa la convalida, il salvataggio dei checkpoint e la registrazione delle metriche.
Implementazione completa
Per comodità, l'implementazione completa viene fornita di seguito come script copia-incolla singolo. Include il dataset personalizzato, il trainer personalizzato e la chiamata di addestramento. Salvalo insieme al tuo dataset.yaml ed eseguilo 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 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)Per consigli sugli iperparametri, consulta la guida Consigli per l'addestramento del modello.
FAQ
Qual è la differenza tra questo e convert_coco()?
convert_coco() scrive i file di etichette .txt su disco come una conversione una tantum. Questo approccio analizza il JSON all'inizio di ogni sessione di addestramento e converte le annotazioni in memoria. Usa convert_coco() quando si preferiscono etichette in formato YOLO permanente; usa questo approccio per mantenere il COCO JSON come unica fonte di verità senza generare file aggiuntivi.
YOLO può addestrarsi su COCO JSON senza codice personalizzato?
Non con l'attuale pipeline Ultralytics, che si aspetta etichette YOLO .txt per impostazione predefinita. Questa guida fornisce il codice personalizzato minimo necessario: una classe di dataset e una classe di trainer. Una volta definite, l'addestramento richiede solo una chiamata model.train() standard.
Supporta la segmentazione e la stima della posa?
This guide covers object detection. To add instance segmentation support, include the segmentation polygon data from COCO annotations in the segments field of each label dictionary. For pose estimation, include keypoints. The GroundingDataset source code provides a reference implementation for handling segments.
Le aumentazioni funzionano con questo dataset personalizzato?
Sì. COCODataset estende YOLODataset, quindi tutte le aumentazioni dei dati integrate ( mosaic, mixup, copy-paste e altre) funzionano senza modifiche.
Come vengono mappati gli ID delle categorie agli indici di classe?
Le categorie sono ordinate per id e mappate su indici sequenziali che partono da 0. Questo gestisce ID basati su 1 (COCO standard), ID basati su 0 e ID non contigui. Il dizionario names in dataset.yaml dovrebbe seguire lo stesso ordine ordinato dell'array categories in COCO.
C'è un sovraccarico di prestazioni rispetto alle etichette pre-convertite?
Il COCO JSON viene analizzato una volta durante la prima sessione di addestramento. Le etichette analizzate vengono salvate in un file .cache, quindi le sessioni successive si caricano istantaneamente senza ri-analisi. La velocità di addestramento è identica a quella dell'addestramento YOLO standard poiché le annotazioni sono conservate in memoria. La cache viene ricostruita automaticamente se il file JSON cambia.