mmdetection源码笔记（三）：创建数据集模型之datasets/coco.py的解读（上）

最新推荐文章于 2025-06-11 15:10:54 发布

原创

最新推荐文章于 2025-06-11 15:10:54 发布 · 5.3k 阅读

20 ·

CC 4.0 BY-SA版权

文章标签：

#mmdetection #目标检测

本文主要解读mmdetection框架中CocoDataset类的源码，它是CustomDataset的子类，用于目标检测任务。内容包括加载图片信息、获取annotation详情、过滤图片以及解析图片annotation等关键方法的详细解释。同时，提到了父类CustomDataset的其他重要函数，并链接了相关系列文章供深入学习。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

引言

同样，数据集也是需要build模型的。类CocoDataset是CustomDataset类的子类，而CustomDataset是Dataset的子类。(之前的创建模型，都是torch.nn.module的子类，数据集的创建就不是咯，注意一下)
关于CustomDataset的定义和其类方法的代码解读，可以看下面这篇文章：

mmdetection源码笔记（三）：创建数据集模型之datasets/custom.py的解读（下）

类CocoDataset依然作为形参，添加到@DATASETS.register_module中，作用就是将其保存到注册表的module中。共有四个方法：

load_annotations()：加载标注文件中的annotation字典，返回图片信息，比如：info{"filename":"284193,faa9000f2678b5e.jpg"}。
get_ann_info(self,idx)：获得annotation的信息，其实是调用了_parse_ann_info()；它的形参是指定的图片id，返回值是个字典：bboxes,bboxes_ignore, labels, masks, mask_polys, poly_lens.
_filter_imgs(self, min_size)：过滤图片，去除没有annotation标注文件的图片，以及图片尺寸小于min-size的图片。
_parse_ann_info(self, ann_info, with_mask=True)：解析一张图片的annotation的信息，主要是bbox和mask信息，返回值为：bboxes,bboxes_ignore, labels, masks, mask_polys, poly_lens.（如果没用mask分支，就没用后面的三个返回值了）

在父类中custom.py有其初始化，还有另外的几个重要的函数，比如prepare_train_img()、prepare_test_img()等。因为其代码行数太长，所以在上面的链接里写一篇讲解。
以下是coco.py的代码，如有错误的地方，还请指出，后面博主也会继续修改，增加对各个代码段的理解。

coco.py代码注释

import numpy as np
from pycocotools.coco import COCO
from .custom import CustomDataset
from .registry import DATASETS
@DATASETS.register_module
class CocoDataset(CustomDataset):
""" coco api
self.coco = COCO(ann_file)
create class members
		# 一维数组，值为对应原coco数据集的annotation/images/categories信息
        coco.anns = anns  
        coco.imgs = imgs
        coco.cats = cats
        # 两个默认value为list的字典 比如：imgToAnns{"1":[ann1,ann2,ann3,....]},'1'为image_id = 1,ann1时其为 1 的annotation。
        coco.imgToAnns = imgToAnns
        coco.catToImgs = catToImgs
        以上五个members，在实例化COCO时，被创建
"""
    CLASSES = ('person', 'mask')
    def load_annotations(self, ann_file):
        self.coco = COCO(ann_file)    
        
        #函数 getImgIds()、getCatIds()、getAnnIds(),返回值为integer array of img/cat/ann ids，形参为过滤条件
        self.cat_ids = self.coco.getCatIds()  # [1,2] , integer array of cat ids
        self.cat2label = {
   
       #  dict
            cat_id: i + 1     #  cat_id : cat_id + 1 ？？
            for i, cat_id in enumerate(self