当前位置：首页 > news >正文

COCO 数据集 80 类标注解析：从 JSON 结构到 PyTorch DataLoader 实现

news 2026/7/7 12:16:34

COCO 数据集 80 类标注解析：从 JSON 结构到 PyTorch DataLoader 实现

计算机视觉领域的研究者和开发者们，一定对 COCO 数据集不陌生。这个由微软团队创建的大规模数据集，已经成为目标检测、实例分割等任务的事实标准。但你是否曾深入探究过它的标注文件结构？又是否遇到过将 COCO 格式转换为模型输入时的各种"坑"？本文将带你从 JSON 文件解析开始，一步步构建完整的 PyTorch 数据加载流程。

1. COCO 数据集标注结构解析

COCO 标注文件采用 JSON 格式存储，主要包含五个关键部分：

{ "info": {}, "licenses": [], "images": [], "annotations": [], "categories": [] }

1.1 核心字段详解

images 数组存储所有图像的基础信息，每个元素包含：

id: 图像唯一标识符
width/height: 图像尺寸
file_name: 图像文件名
coco_url: 在线访问地址

annotations 数组是最复杂的部分，每个标注对象包含：

{ "id": 1768, "image_id": 376765, # 对应images中的id "category_id": 7, # 类别ID "segmentation": [...], # 分割多边形坐标 "area": 702.105, # 目标区域面积 "bbox": [473.07, 395.93, 38.65, 28.67], # [x,y,width,height] "iscrowd": 0 # 是否为一组对象 }

categories 数组定义了 80 个对象类别，结构如下：

{ "id": 1, "name": "person", "supercategory": "animal" # 父类别 }

1.2 边界框格式解析

COCO 采用[x, y, width, height]格式表示边界框，其中：

(x,y)是矩形框左上角坐标
width和height是框的宽高

这种表示法与常见的(x_min, y_min, x_max, y_max)或(center_x, center_y, w, h)格式不同，需要在数据加载时进行转换。

2. PyTorch Dataset 类实现

2.1 基础实现框架

from pycocotools.coco import COCO from torch.utils.data import Dataset import torchvision.transforms as T class CocoDetection(Dataset): def __init__(self, img_dir, ann_file, transforms=None): self.img_dir = img_dir self.coco = COCO(ann_file) self.img_ids = list(sorted(self.coco.imgs.keys())) self.transforms = transforms or T.Compose([ T.ToTensor() ]) def __getitem__(self, idx): img_id = self.img_ids[idx] img_info = self.coco.loadImns(img_id)[0] # 加载图像 img_path = os.path.join(self.img_dir, img_info['file_name']) img = Image.open(img_path).convert('RGB') # 获取标注 ann_ids = self.coco.getAnnIds(imgIds=img_id) anns = self.coco.loadAnns(ann_ids) # 转换标注格式 boxes = [] labels = [] for ann in anns: boxes.append(ann['bbox']) labels.append(ann['category_id']) # 转换为Tensor boxes = torch.as_tensor(boxes, dtype=torch.float32) labels = torch.as_tensor(labels, dtype=torch.int64) # 应用变换 if self.transforms: img = self.transforms(img) return img, {'boxes': boxes, 'labels': labels} def __len__(self): return len(self.img_ids)

2.2 边界框格式转换

将 COCO 格式转换为 YOLO 使用的中心点坐标格式：

def coco_to_yolo(box, img_width, img_height): """ 转换 [x,y,w,h] 到 [cx,cy,w,h] 并归一化 """ x, y, w, h = box cx = (x + w/2) / img_width cy = (y + h/2) / img_height nw = w / img_width nh = h / img_height return [cx, cy, nw, nh]

3. 数据增强与预处理

3.1 常用增强策略

from torchvision.transforms import functional as F import random class RandomHorizontalFlip(object): def __init__(self, prob=0.5): self.prob = prob def __call__(self, image, target): if random.random() < self.prob: height, width = image.shape[-2:] image = F.hflip(image) bbox = target["boxes"] # 翻转x坐标 bbox[:, 0] = width - bbox[:, 0] - bbox[:, 2] target["boxes"] = bbox return image, target

3.2 完整变换管道

def get_transform(train): transforms = [] transforms.append(T.ToTensor()) if train: transforms.append(RandomHorizontalFlip(0.5)) transforms.append(T.ColorJitter( brightness=0.2, contrast=0.2, saturation=0.2)) return T.Compose(transforms)

4. DataLoader 配置与优化

4.1 批处理函数实现

由于图像中的目标数量不同，需要自定义 collate_fn：

def collate_fn(batch): images = [] targets = [] for img, target in batch: images.append(img) targets.append(target) return torch.stack(images), targets

4.2 完整数据加载流程

# 初始化数据集 train_dataset = CocoDetection( img_dir='train2017', ann_file='annotations/instances_train2017.json', transforms=get_transform(True) ) # 创建DataLoader train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=8, shuffle=True, num_workers=4, collate_fn=collate_fn, pin_memory=True ) # 使用示例 for images, targets in train_loader: # 训练代码... pass

5. 性能优化技巧

5.1 并行加载优化

# 在Linux系统上使用更高效的多进程加载 if __name__ == '__main__': train_loader = DataLoader( dataset, batch_size=16, num_workers=8, pin_memory=True, prefetch_factor=2, collate_fn=collate_fn )

5.2 缓存机制实现

class CachedCocoDataset(CocoDetection): def __init__(self, img_dir, ann_file, transforms=None, cache_dir='.cache'): super().__init__(img_dir, ann_file, transforms) os.makedirs(cache_dir, exist_ok=True) self.cache_dir = cache_dir def __getitem__(self, idx): img_id = self.img_ids[idx] cache_path = os.path.join(self.cache_dir, f'{img_id}.pt') if os.path.exists(cache_path): return torch.load(cache_path) # 原始加载逻辑 img, target = super().__getitem__(idx) # 保存到缓存 torch.save((img, target), cache_path) return img, target

6. 常见问题解决方案

6.1 类别ID不连续问题

COCO 的类别ID从1开始且不连续，需要建立映射表：

# COCO类别ID到连续ID的映射 COCO_ID_MAP = { 1: 0, 2: 1, 3: 2, 4: 3, 5: 4, 6: 5, 7: 6, 8: 7, 9: 8, 10: 9, 11: 10, 13: 11, 14: 12, 15: 13, 16: 14, 17: 15, 18: 16, 19: 17, 20: 18, 21: 19, 22: 20, 23: 21, 24: 22, 25: 23, 27: 24, 28: 25, 31: 26, 32: 27, 33: 28, 34: 29, 35: 30, 36: 31, 37: 32, 38: 33, 39: 34, 40: 35, 41: 36, 42: 37, 43: 38, 44: 39, 46: 40, 47: 41, 48: 42, 49: 43, 50: 44, 51: 45, 52: 46, 53: 47, 54: 48, 55: 49, 56: 50, 57: 51, 58: 52, 59: 53, 60: 54, 61: 55, 62: 56, 63: 57, 64: 58, 65: 59, 67: 60, 70: 61, 72: 62, 73: 63, 74: 64, 75: 65, 76: 66, 77: 67, 78: 68, 79: 69, 80: 70, 81: 71, 82: 72, 84: 73, 85: 74, 86: 75, 87: 76, 88: 77, 89: 78, 90: 79 }

6.2 处理 crowd 标注

对于iscrowd=1的标注，通常需要特殊处理：

# 在__getitem__方法中过滤crowd标注 ann_ids = self.coco.getAnnIds(imgIds=img_id, iscrowd=False)

7. 完整代码示例

以下是一个整合了所有功能的完整实现：

import os import torch from PIL import Image from pycocotools.coco import COCO from torch.utils.data import Dataset import torchvision.transforms as T class CocoDataset(Dataset): def __init__(self, root, annotation, transforms=None): self.root = root self.transforms = transforms self.coco = COCO(annotation) self.ids = list(sorted(self.coco.imgs.keys())) # COCO类别ID映射 self.coco_ids = [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90 ] self.class_ids = {coco_id: i for i, coco_id in enumerate(self.coco_ids)} def __getitem__(self, index): coco = self.coco img_id = self.ids[index] # 加载图像 img_info = coco.loadImgs(img_id)[0] path = img_info['file_name'] img = Image.open(os.path.join(self.root, path)).convert('RGB') # 获取标注 ann_ids = coco.getAnnIds(imgIds=img_id, iscrowd=False) anns = coco.loadAnns(ann_ids) # 解析标注 boxes = [] labels = [] for ann in anns: # 跳过未在coco_ids中的类别 if ann['category_id'] not in self.coco_ids: continue x, y, w, h = ann['bbox'] # 处理越界情况 if w <= 0 or h <= 0: continue boxes.append([x, y, x + w, y + h]) # 转换为xyxy格式 labels.append(self.class_ids[ann['category_id']]) # 转换为Tensor boxes = torch.as_tensor(boxes, dtype=torch.float32) labels = torch.as_tensor(labels, dtype=torch.int64) # 创建目标字典 target = {} target['boxes'] = boxes target['labels'] = labels target['image_id'] = torch.tensor([img_id]) # 应用变换 if self.transforms is not None: img, target = self.transforms(img, target) return img, target def __len__(self): return len(self.ids)