Как обучать YOLO на COCO JSON без конвертации
Зачем обучать напрямую на COCO JSON
Аннотации в формате COCO JSON можно использовать напрямую для обучения Ultralytics YOLO, не конвертируя их сначала в файлы .txt. Это делается путем создания подкласса YOLODataset для разбора COCO JSON «на лету» и внедрения его в конвейер обучения через кастомный тренер.
Такой подход позволяет сохранить COCO JSON в качестве единственного источника истины — никаких вызовов convert_coco(), никакой реорганизации папок и промежуточных файлов разметки. YOLO26 и все остальные модели детектирования Ultralytics YOLO поддерживаются. Модели сегментации и оценки позы требуют дополнительных полей меток (см. FAQ).
Ознакомься с руководством по конвертации из COCO в YOLO для стандартного рабочего процесса convert_coco().
Обзор архитектуры
Необходимо два класса:
COCODataset— считывает COCO JSON и преобразует ограничивающие рамки в формат YOLO в оперативной памяти во время обученияCOCOTrainer— переопределяетbuild_dataset(), чтобы использоватьCOCODatasetвместо стандартногоYOLODataset
Реализация следует тому же шаблону, что и встроенный GroundingDataset, который также считывает JSON-аннотации напрямую. Переопределяются три метода: get_img_files(), cache_labels() и get_labels().
Создание класса датасета для COCO JSON
Класс COCODataset наследуется от YOLODataset и переопределяет логику загрузки меток. Вместо чтения файлов .txt из директории с метками, он открывает файл COCO JSON, перебирает аннотации, сгруппированные по изображениям, и преобразует каждую ограничивающую рамку из формата пикселей COCO [x_min, y_min, width, height] в нормализованный формат YOLO [x_center, y_center, width, height]. Аннотации толпы (iscrowd: 1) и рамки с нулевой площадью пропускаются автоматически.
Метод get_img_files() возвращает пустой список, так как пути к изображениям разрешаются из поля file_name в JSON внутри cache_labels(). Идентификаторы категорий сортируются и переотображаются в индексы классов, начинающиеся с нуля, поэтому работают как 1-базированные (стандарт COCO), так и несмежные схемы ID.
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"]Разобранные метки сохраняются в файл .cache рядом с JSON (например, instances_train.cache). При последующих запусках обучения кэш загружается напрямую, минуя разбор JSON. Если файл JSON изменяется, проверка хеша не проходит, и кэш автоматически перестраивается.
Подключение датасета к конвейеру обучения
Единственное изменение, необходимое в тренере, — это переопределение build_dataset(). Стандартный DetectionTrainer создает YOLODataset, который сканирует файлы меток .txt. Заменив его на COCODataset, тренер начинает считывать данные из COCO JSON.
Путь к файлу JSON берется из кастомного поля train_json / val_json в конфигурации данных (см. Шаг 3). Во время обучения mode="train" разрешается в train_json; во время валидации mode="val" разрешается в val_json. Если val_json не задан, он возвращается к использованию 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,
)Настройка dataset.yaml для COCO JSON
Файл dataset.yaml использует стандартные поля path, train и val для поиска директорий с изображениями. Два дополнительных поля, train_json и val_json, указывают файлы аннотаций COCO, которые считывает COCOTrainer. Поля nc и names определяют количество классов и их названия, соответствующие отсортированному порядку categories в 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Ожидаемая структура директорий:
my_dataset/
images/
train/
img_001.jpg
...
val/
img_100.jpg
...
annotations/
instances_train.json
instances_val.json
dataset.yamlЗапуск обучения на COCO JSON
Когда класс датасета, класс тренера и YAML-конфиг готовы, обучение запускается через стандартный вызов model.train(). Единственное отличие от обычного обучения — это аргумент trainer=COCOTrainer, который сообщает Ultralytics о необходимости использовать кастомный загрузчик датасета вместо стандартного.
from ultralytics import YOLO
model = YOLO("yolo26n.pt")
model.train(data="dataset.yaml", epochs=100, imgsz=640, trainer=COCOTrainer)Весь конвейер обучения работает как ожидалось, включая валидацию, сохранение чекпоинтов и логирование метрик.
Полная реализация
Для удобства полная реализация представлена ниже в виде готового скрипта. Она включает кастомный датасет, кастомный тренер и вызов обучения. Сохрани этот код рядом с твоим dataset.yaml и запусти его напрямую.
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)Для рекомендаций по гиперпараметрам ознакомься с руководством Советы по обучению моделей.
Часто задаваемые вопросы (FAQ)
В чем разница между этим методом и convert_coco()?
convert_coco() записывает файлы меток .txt на диск в рамках однократной конвертации. Этот подход разбирает JSON в начале каждого запуска обучения и преобразует аннотации в оперативной памяти. Используй convert_coco(), если предпочитаешь постоянные метки в формате YOLO; используй этот подход, чтобы сохранить COCO JSON как единственный источник истины без создания дополнительных файлов.
Может ли YOLO обучаться на COCO JSON без кастомного кода?
Нет, текущий конвейер Ultralytics по умолчанию ожидает метки YOLO в формате .txt. Это руководство предоставляет минимально необходимый кастомный код — один класс датасета и один класс тренера. После их определения для обучения требуется только стандартный вызов model.train().
Поддерживает ли это сегментацию и оценку позы?
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.
Работают ли аугментации с этим кастомным датасетом?
Да. COCODataset расширяет YOLODataset, поэтому все встроенные аугментации данных — mosaic, mixup, copy-paste и другие — работают без изменений.
Как ID категорий отображаются на индексы классов?
Категории сортируются по id и отображаются на последовательные индексы, начиная с 0. Это обрабатывает 1-базированные ID (стандарт COCO), 0-базированные ID и несмежные ID. Словарь names в dataset.yaml должен следовать тому же отсортированному порядку, что и массив categories в COCO.
Есть ли снижение производительности по сравнению с заранее конвертированными метками?
COCO JSON разбирается один раз при первом запуске обучения. Разобранные метки сохраняются в файл .cache, поэтому последующие запуски загружаются мгновенно без повторного разбора. Скорость обучения идентична стандартному обучению YOLO, так как аннотации хранятся в оперативной памяти. Кэш перестраивается автоматически, если файл JSON изменяется.