Перейти к содержанию

Как обучить YOLO COCO без преобразования

Зачем проводить обучение непосредственно на данных COCO

Аннотации в COCO формат можно использовать непосредственно для Ultralytics YOLO обучение без перехода на .txt сначала файлы. Для этого создают подкласс YOLODataset для оперативного разбора COCO и интеграции их в конвейер обучения с помощью настраиваемого тренера.

При таком подходе файл COCO остается единственным достоверным источником информации — нет convert_coco() вызов, без реорганизации каталога, без промежуточных файлов меток. YOLO26 а также поддерживаются все остальные моделиYOLO Ultralytics . Для моделей сегментации и определения позы требуются дополнительные поля меток (см. Часто задаваемые вопросы).

Или вы ищете вариант однократного преобразования?

Смотрите Руководство YOLO COCO YOLO для стандарта convert_coco() рабочий процесс.

Обзор архитектуры

Требуется два класса:

  1. COCOJSONDataset — считывает COCO и преобразует ограничивающие рамки в YOLO в памяти во время обучения
  2. COCOJSONTrainer — переопределения build_dataset() использовать COCOJSONDataset вместо значения по умолчанию YOLODataset

Реализация выполняется по той же схеме, что и встроенная GroundingDataset, который также напрямую считывает аннотации JSON. Переопределены три метода: get_img_files(), cache_labels()и get_labels().

Создание класса набора данных COCO

Параметр COCOJSONDataset класс наследуется от YOLODataset и переопределяет логику загрузки меток. Вместо чтения .txt файлы из каталога labels, открывает файл COCO , проходит по аннотациям, сгруппированным по изображениям, и преобразует каждую ограничительную рамку из COCO формата COCO [x_min, y_min, width, height] в формат YOLO по центру [x_center, y_center, width, height]. Аннотации пользователей (iscrowd: 1) и блоки с нулевой площадью пропускаются автоматически.

Параметр get_img_files() метод возвращает пустой список, поскольку пути к изображениям извлекаются из JSON file_name поле внутри cache_labels(). Идентификаторы категорий сортируются и преобразуются в индексы классов с нулевым индексом, поэтому правильно работают как схемы с индексацией, начинающейся с 1 (стандартная COCO), так и схемы с несмежными идентификаторами.

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

Проанализированные метки сохраняются в .cache файл рядом с JSON (например, instances_train.cache). При последующих прогонах модели кэш загружается напрямую, минуя этап разбора JSON. Если файл JSON изменяется, проверка хеша завершается с ошибкой, и кэш автоматически перестраивается.

Подключение набора данных к конвейеру обучения

Единственное изменение, которое необходимо внести в тренажер, — это переопределение build_dataset(). По умолчанию DetectionTrainer создает YOLODataset который выполняет сканирование на наличие .txt файлы меток. Заменив его на COCOJSONDataset... вместо этого тренер использует данные из файла COCO .

Путь к файлу JSON берется из настраиваемого train_json / val_json field in the data config (see Step 3). During training, 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 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,
        )

Настройка файла dataset.yaml COCO

Параметр dataset.yaml использует стандарт path, trainи val поля для указания каталогов с изображениями. Два дополнительных поля, train_json и val_json, укажите файлы COCO , которые COCOJSONTrainer гласит. 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

После подготовки набора данных, класса тренера и конфигурации YAML обучение проходит по стандартной схеме model.train() звонок. Единственное отличие от обычной тренировочной пробежки заключается в том, что trainer=COCOJSONTrainer аргумент, который указывает Ultralytics пользовательский модуль загрузки наборов данных вместо модуля по умолчанию.

from ultralytics import YOLO

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

Весь процесс обучения проходит в соответствии с ожиданиями, включая валидацию, сохранение контрольных точек и регистрацию метрик.

Полное внедрение

Для удобства ниже приводится полный код в виде одного скрипта, который можно просто скопировать и вставить. Он включает в себя набор данных, собственный тренажер и вызов функции обучения. Сохраните его рядом с вашим 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 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)

Рекомендации по гиперпараметрам см. в руководстве «Советы по обучению модели ».

Часто задаваемые вопросы

В чём разница между этой функцией и convert_coco()?

convert_coco() пишет .txt сохранить файлы с аннотациями на диск в ходе однократного преобразования. При таком подходе JSON-файл анализируется в начале каждого цикла обучения, а аннотации преобразуются в памяти. Используйте convert_coco() когда предпочтительны постоянные метки YOLO; используйте этот подход, чтобы сохранить файл COCO в качестве единственного достоверного источника данных без создания дополнительных файлов.

Может ли YOLO на данных COCO без написания собственного кода?

Но не в рамках текущего Ultralytics , где предполагается использование YOLO .txt по умолчанию. В данном руководстве представлен минимальный объем необходимого пользовательского кода — один класс набора данных и один класс тренера. После их определения для обучения требуется лишь стандартный model.train() звонок.

Поддерживает ли это сегментацию и оценку положения?

В данном руководстве рассматриваются обнаружения объектов. Чтобы добавить сегментация экземпляров поддержку, включая segmentation данные о многоугольниках из COCO в segments поле каждого словаря меток. Для оценка позы, в том числе keypoints. GroundingDataset исходный код предоставляет эталонную реализацию для обработки сегментов.

Работают ли расширения с этим настраиваемым набором данных?

Да. COCOJSONDataset продлевает YOLODataset, поэтому все встроенные расширение набора данных — mosaic, mixup, copy-pasteи другие — работают без изменений.

Как идентификаторы категорий сопоставляются с индексами классов?

Категории отсортированы по id и сопоставляются последовательным индексам, начиная с 0. Это позволяет обрабатывать идентификаторы с нулевым индексом (стандарт COCO), идентификаторы с нулевым индексом и несмежные идентификаторы. names словарь в dataset.yaml должен следовать тому же отсортированному порядку, что и COCO categories массив.

Есть ли снижение производительности по сравнению с предварительно преобразованными метками?

Файл COCO анализируется один раз при первом прогоне обучения. Проанализированные метки сохраняются в .cache файл, поэтому последующие запуски загружаются мгновенно без повторного анализа. Скорость обучения такая же, как при стандартном YOLO , поскольку аннотации хранятся в памяти. Кэш автоматически обновляется при изменении файла JSON.



📅 Создано 0 дней назад ✏️ Обновлено 0 дней назад
glenn-jocherraimbekovm

Комментарии