Как обучать YOLO на COCO JSON без конвертации#
Аннотации в формате COCO JSON можно использовать напрямую для обучения Ultralytics YOLO без предварительного преобразования в файлы .txt. Это работает за счет создания подкласса YOLODataset для разбора COCO JSON на лету и интеграции его в конвейер обучения через кастомный тренер.
Почему стоит обучать напрямую на COCO JSON#
Этот подход позволяет сохранить COCO JSON в качестве единственного источника достоверных данных — никаких вызовов convert_coco(), реорганизации директорий и промежуточных файлов меток. Поддерживаются YOLO26 и все другие модели детектирования Ultralytics YOLO. Модели сегментации и позы требуют дополнительных полей меток (см. Часто задаваемые вопросы).
Смотри руководство по конвертации 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):
"""Initialize the dataset with a COCO JSON annotation file."""
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)
# 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 в конфигурации данных (см. Настройка dataset.yaml). Во время обучения 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):
"""Build a COCODataset for the given split using the JSON file from the data config."""
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/my_dataset/images # root with train/ and val/ image subfolders
train: train
val: val
# COCO JSON annotation files (use absolute paths; these custom keys are not resolved against `path`)
train_json: /path/to/my_dataset/annotations/instances_train.json
val_json: /path/to/my_dataset/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):
"""Initialize the dataset with a COCO JSON annotation file."""
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)
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):
"""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"]
class COCOTrainer(DetectionTrainer):
"""Trainer that uses COCODataset for direct COCO JSON training."""
def build_dataset(self, img_path, mode="train", batch=None):
"""Build a COCODataset for the given split using the JSON file from the data config."""
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)Теперь у тебя есть минимальный набор данных и трейнер, которые обучают Ultralytics YOLO напрямую на COCO JSON, причем аннотации остаются единственным источником истины без промежуточных файлов .txt. Расширь метод cache_labels() с помощью segments или keypoints, чтобы охватить сегментацию и позы, и ознакомься с руководством Советы по обучению моделей для получения рекомендаций по настройке гиперпараметров.
FAQ#
В чем разница между этим методом и convert_coco()?#
convert_coco() записывает файлы меток .txt на диск в рамках однократного преобразования. Этот подход разбирает JSON в начале каждого запуска обучения и преобразует аннотации в памяти. Используй convert_coco(), если предпочитаешь постоянные метки в формате YOLO; применяй этот подход, чтобы держать COCO JSON в качестве единственного источника истины без генерации дополнительных файлов.
Может ли YOLO обучаться на COCO JSON без написания кастомного кода?#
Не с текущим конвейером Ultralytics, который по умолчанию ожидает метки YOLO .txt. В этом руководстве представлен необходимый минимальный кастомный код: один класс датасета и один класс тренера. После их определения для обучения требуется лишь стандартный вызов model.train().
Поддерживает ли этот метод сегментацию и оценку поз?#
В этом руководстве рассматривается обнаружение объектов. Чтобы добавить поддержку сегментации экземпляров, включи данные полигонов segmentation из аннотаций COCO в поле segments каждого словаря меток. Для оценки поз включи keypoints. Исходный код GroundingDataset предоставляет эталонную реализацию для работы с сегментами.
Работают ли аугментации с этим кастомным набором данных?#
Да. COCODataset расширяет YOLODataset, поэтому все встроенные аугментации данных — мозаика, mixup, copy-paste и другие — работают без изменений.
Как ID категорий отображаются на индексы классов?#
Категории сортируются по id и маппятся на последовательные индексы, начиная с 0. Это обрабатывает 1-индексированные ID (стандартный COCO), 0-индексированные ID и неконтигуированные ID. Словарь names в dataset.yaml должен следовать тому же отсортированному порядку, что и массив COCO categories.
Есть ли снижение производительности по сравнению с заранее конвертированными метками?#
Файл COCO JSON разбирается один раз при первом запуске обучения. Разобранные метки сохраняются в файл .cache, поэтому последующие запуски загружаются мгновенно без повторного разбора. Скорость обучения идентична стандартному обучению YOLO, так как аннотации хранятся в памяти. Кэш перестраивается автоматически при изменении файла JSON.