跳转至内容

如何在不进行转换的情况下，YOLO COCO 训练YOLO

为什么直接在COCO 上进行训练

标注 在 COCO 该格式可直接用于 Ultralytics YOLO 在不转换的情况下进行训练 .txt 首先处理文件。这是通过继承类来实现的 YOLODataset 实时解析COCO 数据，并通过自定义训练器将其集成到训练管道中。

这种方法将COCO 作为唯一可信数据源——不 convert_coco() 调用，不进行目录重组，不生成中间标签文件。 YOLO26 并支持所有其他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 格式 [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 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,
        )

配置COCO yaml dataset.yaml 文件

字段 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 YOLO ？

在当前的Ultralytics 中并非如此，该管道预期使用YOLO .txt 默认情况下使用标签。本指南提供了所需的最小自定义代码——一个数据集类和一个训练器类。定义完成后，训练只需一个标准的 model.train() 呼叫。

这支持分割和姿势估计 吗？

本指南涵盖 对象检测. 要添加 实例分割 支持，包括 segmentation 来自COCO 的 segments 每个标签字典的字段。对于 姿势估计, 包括 keypoints。 GroundingDataset 源代码 提供了处理分段的参考实现。

增强学习方法适用于这个自定义数据集吗？

是的。 COCOJSONDataset 扩展 YOLODataset，因此所有内置的 数据增强 — mosaic, mixup, 复制-粘贴……等——均可直接运行，无需修改。

类别 ID 是如何映射到类索引的？

分类按以下方式排序： id 并映射到从 0 开始的顺序索引。这支持 1 为起始点的 ID（标准COCO）、0 为起始点的 ID 以及非连续的 ID。该 names 词典在 dataset.yaml 应遵循与COCO相同的排序顺序 categories 数组。

与预转换标签相比，是否会产生性能开销？

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



📅 创建于 0 天前 ✏️ 更新于 0 天前
glenn-jocherraimbekovm

评论