从YOLOv5到Detectron2：COCO数据集在不同CV框架下的加载与预处理实战

从YOLOv5到Detectron2：COCO数据集跨框架加载与预处理实战指南

在计算机视觉领域，COCO数据集已成为目标检测和实例分割任务的事实标准。但对于开发者而言，面对PyTorch生态中YOLOv5、MMDetection和Detectron2等不同框架时，数据加载和预处理往往成为项目落地的第一道门槛。本文将深入剖析三大主流框架下的COCO数据处理差异，提供可复用的工程实践方案。

1. COCO数据集核心特性与框架适配要点

COCO数据集包含33万张图像，涵盖80个日常物体类别，其标注信息之丰富远超同类数据集。但这也带来了框架适配的复杂性：

• 多任务支持：同时提供目标检测（bbox）、实例分割（mask）和关键点检测标注
• 层次结构：标注文件采用JSON格式，包含images、annotations、categories三层嵌套
• 小目标密集：平均每张图像包含7.2个目标，存在大量小尺寸物体

不同框架对COCO数据的处理方式存在显著差异：

提示：选择框架时需考虑团队技术栈和项目规模，小团队快速验证推荐YOLOv5，大型项目建议采用Detectron2的完整生态

2. YOLOv5的COCO适配实战

YOLOv5采用独特的TXT标注格式，与COCO原生JSON存在显著差异。以下是完整的转换流程：

2.1 标注格式转换

import json from pathlib import Path def coco2yolo(coco_json, output_dir): with open(coco_json) as f: data = json.load(f) # 创建类别映射 cat_map = {cat['id']: i for i, cat in enumerate(data['categories'])} for img in data['images']: img_id = img['id'] anns = [a for a in data['annotations'] if a['image_id'] == img_id] txt_path = Path(output_dir) / f"{Path(img['file_name']).stem}.txt" with open(txt_path, 'w') as f: for ann in anns: # 转换bbox格式：xywh -> xyxy -> 归一化 x, y, w, h = ann['bbox'] x_center = (x + w/2) / img['width'] y_center = (y + h/2) / img['height'] norm_w = w / img['width'] norm_h = h / img['height'] line = f"{cat_map[ann['category_id']]} {x_center} {y_center} {norm_w} {norm_h}\n" f.write(line)

关键转换步骤：

1. 将COCO的绝对坐标xywh转为YOLO格式的归一化xywh
2. 类别ID重新映射为连续整数
3. 每个图像生成对应的TXT标注文件

2.2 自定义数据加载

YOLOv5通过dataset.py实现数据加载，主要优化点包括：

class COCODataset(torch.utils.data.Dataset): def __init__(self, img_dir, label_dir, img_size=640): self.img_files = list(Path(img_dir).glob('*.jpg')) self.label_files = [Path(label_dir)/f"{f.stem}.txt" for f in self.img_files] self.img_size = img_size self.transform = A.Compose([ A.HorizontalFlip(p=0.5), A.RandomBrightnessContrast(p=0.2), A.Resize(img_size, img_size) ], bbox_params=A.BboxParams(format='yolo')) def __getitem__(self, idx): img = cv2.imread(str(self.img_files[idx])) labels = np.loadtxt(self.label_files[idx], dtype=np.float32) # Albumentations增强 transformed = self.transform(image=img, bboxes=labels[:, 1:]) img = transformed['image'] labels[:, 1:] = np.array(transformed['bboxes']) return torch.from_numpy(img).permute(2,0,1), torch.from_numpy(labels)

性能优化技巧：

• 使用albumentations替代原生torchvision变换
• 开启Dataloader的persistent_workers减少进程创建开销
• 对固定尺寸模型预处理时设置collate_fn=identity

3. Detectron2的COCO原生支持解析

Detectron2作为Facebook官方框架，对COCO数据集的支持最为完善。其核心优势在于：

3.1 内置数据注册系统

from detectron2.data import DatasetCatalog, MetadataCatalog # 注册COCO数据集 def get_coco_dicts(img_dir, json_file): from detectron2.data.datasets import load_coco_json return load_coco_json(json_file, img_dir) DatasetCatalog.register("my_coco_train", lambda: get_coco_dicts("train2017", "annotations_trainval2017.json")) MetadataCatalog.get("my_coco_train").set(thing_classes=["person", "bicycle", ...]) # 可视化验证 dataset_dicts = DatasetCatalog.get("my_coco_train") metadata = MetadataCatalog.get("my_coco_train")

3.2 高级数据增强管道

Detectron2通过Augmentation系统实现灵活的数据增强：

from detectron2.data import transforms as T aug_list = [ T.RandomFlip(horizontal=True, vertical=False), T.RandomApply(T.RandomRotation(angle=[-15, 15]), prob=0.3), T.RandomApply(T.RandomContrast(intensity_min=0.8, intensity_max=1.2), prob=0.5), T.ResizeShortestEdge(short_edge_length=(640, 672, 704, 736, 768), max_size=1333) ] # 构建mapper def mapper(dataset_dict): image = cv2.imread(dataset_dict["file_name"]) aug_input = T.AugInput(image) transforms = aug_list(aug_input) image = aug_input.image annos = [ utils.transform_instance_annotations(anno, transforms, image.shape[:2]) for anno in dataset_dict["annotations"] ] return { "image": torch.as_tensor(image.transpose(2,0,1)), "instances": utils.annotations_to_instances(annos, image.shape[:2]) }

Detectron2特有功能：

• 自动处理crowd区域标注
• 支持RLE格式的mask压缩存储
• 内置全景分割标注转换

4. MMDetection的多后端支持方案

OpenMMLab生态的MMDetection提供了最灵活的COCO处理方案，支持多种训练范式：

4.1 配置文件定义

# configs/_base_/datasets/coco_detection.py dataset_type = 'CocoDataset' data_root = 'data/coco/' train_pipeline = [ dict(type='LoadImageFromFile'), dict(type='LoadAnnotations', with_bbox=True), dict(type='Resize', img_scale=(1333, 800), keep_ratio=True), dict(type='RandomFlip', flip_ratio=0.5), dict(type='Normalize', mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375]), dict(type='Pad', size_divisor=32), dict(type='DefaultFormatBundle'), dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels']) ] data = dict( samples_per_gpu=2, workers_per_gpu=2, train=dict( type=dataset_type, ann_file=data_root + 'annotations/instances_train2017.json', img_prefix=data_root + 'train2017/', pipeline=train_pipeline), val=dict(...), test=dict(...) )

4.2 自定义数据混合

MMDetection支持多种数据集混合训练：

# 混合COCO和Objects365数据集 dataset_A_train = dict( type='CocoDataset', ann_file='data/coco/annotations/instances_train2017.json', img_prefix='data/coco/train2017/', pipeline=train_pipeline ) dataset_B_train = dict( type='CocoDataset', ann_file='data/objects365/annotations/train.json', img_prefix='data/objects365/train/', pipeline=train_pipeline ) data = dict( train=dict( type='ConcatDataset', datasets=[dataset_A_train, dataset_B_train]), ... )

性能对比测试（RTX 3090, batch_size=16）：

在实际项目中，当遇到COCO标注文件损坏或格式异常时，可先用pycocotools验证数据完整性：

from pycocotools.coco import COCO coco = COCO("annotations/instances_val2017.json") print(coco.getCatIds()) # 验证类别读取 print(len(coco.getImgIds())) # 验证图像数量