跳转至内容

如何在不转换的情况下在 COCO JSON 上训练 YOLO

为何直接在 COCO JSON 上训练

标注COCO JSON 格式可直接用于 Ultralytics YOLO 训练,无需先转换为 .txt 文件。这是通过子类化 YOLODataset 来实现的,它可以在训练过程中动态解析 COCO JSON 并通过自定义训练器将其集成到训练管道中。

这种方法将 COCO JSON 作为单一事实来源 — 无需 convert_coco() 调用,无需目录重组,也无需中间标签文件。 YOLO26 以及所有其他 Ultralytics YOLO 检测模型都受支持。分割和姿势估计模型需要额外的标签字段(参见 常见问题)。

是否正在寻找一次性转换方案?

请参阅 COCO 到 YOLO 转换指南 以了解标准 convert_coco() 工作流程。

架构概览

需要两个类:

  1. COCODataset — 读取 COCO JSON 并将 边界框 在训练期间在内存中转换为 YOLO 格式
  2. 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() 方法返回一个空列表,因为图像路径是从 JSON file_name 字段中解析的 cache_labels()。类别 ID 会被排序并重新映射到从零开始的类索引,因此基于 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,
        )

为 COCO JSON 配置 dataset.yaml

字段 dataset.yaml 使用标准 path, trainval 字段来定位图像目录。另外两个字段, train_jsonval_json,指定了 COCO 注释文件, COCOTrainer 读取。 ncnames 字段定义了类别数量及其名称,与 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)

有关 超参数 建议,请参阅 模型训练技巧 指南。

常见问题

这与 convert_coco() 有何不同?

convert_coco() 写入 .txt 标签文件到磁盘,作为一次性转换。此方法在每次训练运行开始时解析 JSON 并在内存中转换注释。使用 convert_coco() 当偏好永久性的 YOLO 格式标签时;使用此方法可将 COCO JSON 作为唯一的事实来源,而无需生成额外文件。

YOLO 是否可以在不编写自定义代码的情况下直接使用 COCO JSON 进行训练?

不适用于当前的 Ultralytics 管道,该管道默认期望 YOLO .txt 标签。本指南提供了所需的最小自定义代码 — 一个数据集类和一个训练器类。一旦定义,训练只需一个标准 model.train() 调用。

这是否支持分割和姿势估计?

本指南涵盖 对象检测。要添加 实例分割 支持,请在 segmentation 中包含 COCO 注释中的多边形数据。 segments 每个标签字典的字段。对于 姿势估计,包括 keypointsGroundingDataset 源代码 提供了处理 segment 的参考实现。

数据增强是否适用于此自定义数据集?

是的。 COCODataset 扩展了 YOLODataset,因此所有内置的 数据增强mosaic, mixup, 复制-粘贴等,都可以无需修改地运行。

类别 ID 如何映射到类索引?

类别按 id 排序,并映射到从 0 开始的连续索引。这可以处理基于 1 的 ID(标准 COCO)、基于 0 的 ID 和非连续 ID。 names 字典在 dataset.yaml 中应遵循与 COCO categories 数组相同的排序顺序。

与预转换的标签相比,是否存在性能开销?

COCO JSON 在首次训练运行时解析一次。解析后的标签保存到 .cache 文件,因此后续运行无需重新解析即可立即加载。由于标注保存在内存中,训练速度与标准 YOLO 训练相同。如果 JSON 文件发生更改,缓存会自动重建。



📅 创建于 7 天前 ✏️ 更新于 7 天前
raimbekovmglenn-jocher

评论