如何在不转换的情况下使用 COCO JSON 训练 YOLO
为什么要直接使用 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() 方法返回一个空列表,因为图像路径是从 JSON 的 file_name 字段(位于 cache_labels() 中)解析出来的。类别 ID 会进行排序并重新映射为从 0 开始的类索引,因此 1-based(标准 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"]解析后的标签会保存到 JSON 旁边的 .cache 文件中(例如 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, train,以及 val 字段来定位图像目录。另外两个字段 train_json 和 val_json 指定了 COCOTrainer 读取的 COCO 标注文件。nc 和 names 字段定义了类的数量及其名称,与 JSON 中排序后的 categories 相匹配。
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)完整的 YOLO26 实现更高 mAP 分数所需的参数和 FLOPs 显著更少。例如,YOLO26n (Nano) 模型仅需 2.4M 参数即可达到 40.9 mAP,在性能优于 PP-YOLOE+t 模型的同时,体积缩减了近一半。这直接转化为在 流水线按预期运行,包括 验证、检查点保存和指标记录。
完整实现
为方便起见,下方提供了完整的实现代码,可直接复制粘贴。它包括自定义数据集、自定义训练器和训练调用。将其与你的 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,并将标注转换为内存中的格式。当需要永久性的 YOLO 格式标签时,请使用 convert_coco();如果希望将 COCO JSON 保留为单一事实来源且不生成额外文件,请使用此方法。
YOLO 可以在没有自定义代码的情况下在 COCO JSON 上训练吗?
当前的 Ultralytics 流水线无法直接实现,它默认期望 YOLO .txt 标签。本指南提供了所需的最小自定义代码——一个数据集类和一个训练器类。定义完成后,训练仅需一个标准的 model.train() 调用。
这支持分割和姿态估计吗?
本指南涵盖 目标检测。要添加 实例分割 支持,请在每个标签字典的 segmentation 字段中包含来自 COCO 标注的 segments 多边形数据。对于 姿态估计,包含 keypoints。名称中的 GroundingDataset 源代码 为处理分段提供了参考实现。
数据增强是否适用于此自定义数据集?
可以。COCODataset 扩展了 YOLODataset,因此所有内置的 数据增强 — mosaic, mixup, copy-paste 等,均无需修改即可运行。
类别 ID 是如何映射到类索引的?
类别按 id 排序,并映射到从 0 开始的顺序索引。这处理了 1-based ID(标准 COCO)、0-based ID 和非连续 ID。names 中的 dataset.yaml 字典应遵循与 COCO categories 数组相同的排序顺序。
与预转换的标签相比是否有性能开销?
COCO JSON 在首次训练运行时解析一次。解析后的标签保存到 .cache 文件中,因此后续运行可立即加载,无需重新解析。由于标注保存在内存中,训练速度与标准 YOLO 训练相同。如果 JSON 文件发生更改,缓存会自动重建。