Detectron 源码解析-roidb数据结构

roidb数据结构

roidb的类型是list, 其中的每个元素的数据类型都是dict, roidb列表的长度为数据集的数量(即图片的数量), roidb中每个元素的详细情况如下表所示:

for entry in roidb	数据类型	详细说明
entry['id']	int	代表了当前image的img_id
entry['file_name']	string	表示当前图片的文件名(带有.jpg后缀)
entry['dataset']	string	指明所属的数据集?
entry['image']	string	当前image的文件路径
entry['flipped']	bool	当前图片是否进行翻转
entry['height']	int	当前图片的高度
entry['width']	int	当前图片的宽度
entry['has_visible_keypoints']	bool	是否含有关键点
entry['boxes']	float32, numpy数组(num_objs, 4)	num_objs为当前图片中的目标物体个数, 4代表bbox的坐标
entry['segms']	二维列表[[],[],…]	列表中每个元素都还是一个列表, 其中存储着每个物体的ploygon实例标签
entry['gt_classes']	int32, numpy数组(num_objs)	指明当前图片中每一个obj的真实类别
entry['seg_areas']	float32, numpy数组(num_objs)	代表当前图片中每一个obj的掩膜面积
entry['gt_overlaps']	float32, scipy.sparse.csr_matrix数据(num_objs, 81)	代表每一个obj与81个不同类别的overlap
entry['is_crowd']	bool, numpy数组(num_objs)	代表当前掩膜是否为群落
entry[‘box_to_gt_ind_map’]	int32, numpy数组(num_objs)	该列表存储着box的顺序下标值, 同样是一维数组, 直接拼接,将每一个roi映射到一个index上, index是在entry[‘gt_classes’]>0的rois列表的下标

combined_roidb_for_training() 方法

在目标检测类任务中, 有一个很重要的数据结构roidb, 它将作为基本的数据结构在数据队列中存在, Detectron 的数据载入类 RoIDdataLoader 也是将该数据结构作为成员变量使用的, 因此, 有必要对这个数据结构展开分析.

首先, 在运行训练脚本时, 就会调用到 detectron/utils/train.py 中的 train()函数, 而train()函数内部又会调用当前文件的add_model_training_inputs() 函数, 在这个函数内部, 就会调用到 detectron/datasets/roidb 文件中的 combined_roidb_for_training() 函数, 该函数的返回值正是roidb, 这是贯穿整个训练过程的训练数据, 故我们对此函数进行分析. 该函数代码解析如下:

# detectron/datasets/roidb.py

# 加载并连接一个或多个数据集的roidbs, along with optional object proposals
# 每个roidb entry都带有特定的元数据类型, 对其进行准备工作后进行训练
def combined_roidb_for_training(dataset_names, proposal_files):
    def get_roidb(dataset_name, proposal_file):
        # 注意 dataset_name 没有 's'

        # from detectron.datasets.json_dataset import JsonDataset
        # 可以看到, roidb 是利用JsonDataset类对象的get_roidb()方法获取的
        # 因此, 我们先在下面看一下这个类的实现细节
        ds = JsonDataset(dataset_name)
        roidb = ds.get_roidb(
            gt=True,
            proposal_file=proposal_file,
            crowd_filter_thresh=cfg.TRAIN.CROWD_FILTER_THRESH
        )
        if cfg.TRAIN.USE_FLIPPED:
            logger.info("Appending horizontally-flipped training examples...")
            extend_with_flipped_entries(roidb, ds)
        logger.info("Loaded dataset: {:s}".format(ds.name))
        return roidb
    if isinstance(dataset_names, basestring):
        #...
    #...

get_roidb() 方法

在上面的函数中我们可以发现, combined_roidb_for_training函数内部又定义了另一个函数get_roidb(), 而该函数主要是基于detectron/datasets/json_dataset.py中的JsonDataset类及该类的成员方法get_roidb实现的, 因此, 我们先跳到json_dataset.py文件中去看看这个类的内部实现是怎样的:

# detectron/datasets/json_dataset.py

class JsonDataset(object):
    # 这个类的设计主要是基于COCO的json格式数据集
    # 当我们需要训练自己的数据集时, 最好的方式就是将自己的数据集的格式改为
    # COCO数据集的json格式, 这样一来, 我们就无需重写数据载入代码了.
    def __init__(self, name):
        assert dataset_catalog.contains(name), \
            "Unknown dataset name: {}".format(name)
        assert...
        #...

        # 准备数据集的类别信息
        category_ids = self.COCO.getCatIds() # 1~80, 对应80个类
        # coco的loadCats函数, 必须指定需要加载的cat的id, 否则返回空列表
        # 若指定后, 则返回id对应的类别信息, 每个类别信息是一个字典, 包括'name','id','supercategory'三个字段
        # 获取每个类的名字, person, bicycle,bus等等, 返回的名字在列表中的位置与id在cat_ids列表中的位置一致
        categories = [c['name'] for c in sefl.COCO.loadCats(category_ids)]
        # 建立类别的name 与 id之间的对应关系, 其中cat_name为key,cat_id为值
        self.category_to_id_map = dict(zip(categories, category_ids)) # 注意, 没有'__background__'
        self.classes = ['__background__'] + categories # 将'__background__'添加到categories类别名字列表中
        self.num_classes = len(self.classes)
        # coco下标最大值为90,但实际上只有80个类, 有的地方跳过了, 因此id不是连续的,
        self.json_category_id_to_contiguous_id = {
            v: i + 1 # key为coco的非连续id, value为1~80的连续id, 均为整数
            for i, v in enumerate(self.COCO.getCatIds())
        }
        self.contiguous_category_id_to_json_id = {
            v: k # key为1~80的连续id, value为coco的非连续id, 均为整数
            for k, v in self.json_category_id_to_contiguous_id.items()
        }
        self._init_keypoints() # 调用类内的keypoints初始化方法.

    def get_roidb(
        self,
        gt=False,
        proposal_file=None,
        min_proposal_size=2,
        proposal_limit=-1,
        crowd_filter_thresh=0
    ):
        """
        返回json dataset对应的roidb数据, 提供以下四种选项:
        - 在roidb中包含gt boxes
        - 添加位于proposal file里面的特定proposals
        - 基于最短边长的proposals过滤器
        - 基于群落区域交集的proposals过滤器
        """

        assert gt is True or crowd_filter_thresh == 0, \
            "Crowd filter threshold must be 0 if gt " \
            "annotations are not included."
        # 这里调用了COCO的官方API, 关于COCO数据集的结构和标注格式解析, 可以查看我的另一篇文章
        # 没有指定筛选条件, 获取数据集标签中所有的图片id
        image_ids = self.COCO.getImgIds()
        image_ids.sort() # 将id按照从小到大的顺序排列
        # roidb为列表结构, 列表中的每一项是一个字典, 代表着对应imageid的标签内容.
        # 键值包括:coco_url, license, width, filename, height, flickr_url, id, date_captured
        roidb = copy.deepcopy(self.COCO.loadImgs(image_ids))
        for entry in roidb:
            # 调用了本类的私有函数 _prep_roidb_entry(), entry为字典.
            # 主要是为entry赋初值, 占位符等等, 包含box, segms,等各种字段, 详细信息可以看下面的函数解析
            # 注意, 这里的字段值都是预测值相关的值, 因此也会局域gt_overlap等字段
            self._prep_roidb_entry(entry)
        if gt:
            # 如果参数声明是gt信息, 则会调用_add_gt_annotations
            # 访问标注文件, 以便添加相关字段信息, 具体看下面的相关函数解析
            self.debug_timer.tic()
            for entry in roidb:
                # 注意, 是单独对每个entry调用该函数, 因此每次会载入指定imgid的相关标签
                # 关于_add_gt_annotations函数具体解析可以看后面的部分
                self._add_gt_annotations(entry)
            logger.debug(
                '_add_gt_annotations took {:.3f}s'.
                format(self.debug_timer.toc(average=False))
            )
        if proposal_file is not None:
            self.debug_timer.tic()
            # 加载proposals文件到roidb中, 关于此函数的详细解析可以看后文
            self._add_proposals_from_file(
                roidb, proposal_file, min_proposal_size, proposal_limit,
                crowd_filter_thresh
            )
            logger.debug(
                '_add_proposals_from_file took {:.3f}s'.
                format(self.debug_timer.toc(average=False))
            )
        # 类外部的函数, 用于计算与每个roidb相关的box的类别
        _add_class_assignments(roidb)
        return roidb

_prep_roidb_entry() 方法

数据准备函数 _prep_roidb_entry() 的实现解析

# detectron/datasets/json_dataset.py

class JsonDataset(object):
    def __init__(...):
        #...
    def get_roidb(...):
        #...
    # 该函数主要将空的元数据添加到roidb entry中
    def _prep_roidb_entry(self, entry):
        # entry的'dataset'关键字, 值为self.
        entry['dataset'] = self
        im_path = os.path.join(self.image_directory, self.image_prefix+entry['file_name'])
        assert os.path.exists(im_path), "Image \"{} \" not found".format(im_path)
        # entry的'image'关键字, 值为当前imageid对应的image路径
        entry['image'] = im_path
        entry['flipped'] = False # 禁止反转
        entry['has_visible_keypoints'] = False

        # 下面entry键的对应值均为空, 暂为占位键

        # entry的'boxes'关键字,值为n×4的numpy数组, n代表box的数量,这里暂时为0
        entry['boxes'] = np.empty((0,4), dtype=np.float32)
        entry['segms'] = [] # entry的'segms'关键字, 值为一个列表,暂时为空
        # entry的'gt_classes'关键字, 是个一维数组, 维度与box的数量n对应,暂时为0
        entry['gt_classes'] = np.empty((0), dtype=np.int32)
        # 代表掩膜的面积, 供n项, 与boxes数目相对
        entry['seg_areas'] = np.empty((0), dtype=np.float32)
        # TODO, 这里是一个矩阵压缩, 矩阵大小为n×c, c为类别数量, 没太搞懂要压缩成什么?
        entry['gt_overlaps'] = scipy.sparse.csr_matrix(
            np.empty((0, self.num_classes), dtype=np.float32)
        )
        # 同样是n行1列, n与boxes数目对应, 表示是否为`一群物体`
        entry['is_crowd'] = np.empty((0), dtype=np.bool)
        # shape大小与roi相关, 将每一个roi映射到一个index上
        # index是在entry['gt_classes']>0的rois列表的下标 TODO还是不太清楚映射关系
        entry['box_to_gt_ind_map'] = np.empty((0), dtype=np.int32)
        # 关键点信息, 默认情况下不设置
        if self.keypoints is not None:
            entry['gt_keypoints'] = np.empty(
                (0, 3, self.num_keypoints), dtype=np.int32
            )
        # 删除那些从json file中获取到的不需要的字段
        for k in ['date_captured', 'url', 'license', 'file_name']:
            if k in entry:
                del entry[k]

_add_gt_annotations() 方法

加载标注文件的函数 _add_gt_annotations()的实现解析

# detectron/datasets/json_dataset.py

class JsonDataset(object):
    def __init__(...):
        #...
    def get_roidb(...):
        #...
    def _prep_roidb_entry(self, entry):
        #...
    # 该函数将标注文件的元数据添加到roidb entry中
    def _add_gt_annotations(self, entry):
        # 获取指定imgid的annid列表 (对应多个box)
        ann_ids = self.COCO.getAnnIds(imgIds=entry['id'], iscrowd=None)
        # 根据annids的id列表, 获取这些id对应的标注信息, objs是一个列表
        # 列表中的每一个元素都是一个字典,字典的内容是标注文件中的内容,包含bbox,segmentation等字段
        objs = self.COCO.loadAnns(ann_ids)
        # 下面的代码会对bboxes进行清洗, 因为有些是无效的数据
        valid_objs=[] # 存储有效的objs
        valid_segms=[] # 存储有效的segms
        width = entry['width'] # 获取entry中的width字段, 代表图片的宽度
        height = entry['height'] # 获取entry中的height字段, 代表图片的高度
        for obj in objs:
            # crowd区域采用RLE编码
            # import detectron.utils.segms as segm_utils
            # 用于判断当前的segmentation是polygon编码还是rle编码, 前者是列表类型, 后者是字典类型
            # 返回True为polygon编码, 返回Fasle为rle编码
            if segm_utils.is_poly(obj['segmentation']):
                # poly编码必须含有>=3个点才能组成一个多边形, 因此需要>=6个坐标点
                # 类似于这样的检查操作只在PLOYGON中存在, 在面对RLE时无需检查, 可以直接接受后面的检查
                obj['segmentation'] = [
                    p for p in obj['segmentation'] if len(p) >=6
                ]
            if obj['area'] < cfg.TRAIN.GT_MIN_AREA:
                continue # 如果面积不达标, 则认为该标注无效, 不将其加入valid列表
            if 'ignore' in obj and obj['ignore'] == 1:
                continue
            # import detectron.utils.boxes as box_utils
            # 将[x1,y1,w,h]的边框格式转换成[x1,y1,x2,y2]的格式
            x1, y1, x2, y2 = box_utils.xywh_to_xyxy(obj['bbox'])
            # 将[x1,y1,x2,y2]的边框坐标限制在图片的[width,height]范围内, 防止越界
            x1, y1, x2, y2 = box_utils.clip_xyxy_to_image(
                x1, y1, x2, y2, height, width
            )

            if obj['area'] > 0 and x2 > x1 and y2 > y1: # 若数据有效, 则加入到列表当中
                obj['clean_bbox'] = [x1, y1, x2, y2]
                valid_objs.append(obj)
                valid_segms.append(obj['segmentation']) # 将数据的segms存在列表中(RLE/PLOYGON)
        num_valid_objs = len(valid_objs) # num_valid_objs持有objs的有效个数

        # 注意, 下面的数据内容都被初始化为0
        # boxes为 有效objs数×4 的numpy数组, 用来表示每个objs的边框坐标
        boxes = np.zeros((num_valid_objs,4), dtype=entry['seg_areas'].dtype)
        # 每个objs的真实类别
        gt_classes = np.zeros((num_valid_objs), dtype=entry['gt_classes'].dtype)
        gt_overlaps = np.zeros( # 形状为 有效objs数×num_class数 的numpy数组, 表示与每个类的IoU大小
            (num_valid_objs, self.num_classes),
            dtype=entry['gt_overlaps'].dtype
        )
        # 掩膜面积
        seg_areas = np.zeros((num_valid_objs), dtype=entry['seg_areas'].dtype)
        # 是否crowd
        is_crowd = np.zeros((num_valid_objs), dtype=entry['is_crowd'].dtype)
        # 这个是???
        box_to_gt_ind_map = np.zeros(
            (num_valid_objs), dtype=entry['box_to_gt_ind_map'].dtype
        )
        if self.keypoints is not None:
            gt_keypoints = np.zeros(
                (num_valid_objs, 3, self.num_keypoints),
                dtype=entry['gt_keypoints'].dtype
            )

        # 图片是否有可视的关键点?
        im_has_visible_keypoints = False
        for ix, obj in enumerate(valid_objs):# ix为下标, obj为下标对应元素
            # category_id为coco类别id,json_category_id_to_contiguous_id 为字典类型
            # 其中, key为coco的非连续id, value为1~80的连续id, 均为整数, 所以这里是将coco的非连续id转换成对应的连续id
            cls = self.json_category_id_to_contiguous_id[obj['category_id']]
            boxes[ix, :] = obj['clean_box'] # 将当前obj的box填入boxes列表
            gt_classes[ix] = cls # 将连续id填入gt_classes列表
            seg_areas[ix] = obj['area'] # 将area填入seg_areas列表
            is_crowd[ix] = obj['iscrowd']
            box_to_gt_ind_map[ix] = ix # 该列表存储着box的顺序下标值
            if self.keypoints is not None:
                # ...
            if obj['iscrowd']:
                # 如果当前物体是crowd的话, 则将所有类别的overlap都设置为-1,
                # 这样一来在训练的时候, 这些物体都会被排除在外!!
                gt_overlaps[ix, :] = -1.0
            else:
                gt_overlaps[ix, cls] = 1.0  # 仅仅将对应类的overlap设置为1, 其他为0
        # 将gt的boxes添加到entry中, 注意axis为0, 则会按照第0维拼接, 即最后是一个n×4的数组
        # 注意, entry['boxes']初始时候是空的, 因此这就相当于是只添加了真实的框
        entry['boxes'] = np.append(entry['boxes'], boxes, axis=0)
        # 由于segms是以列表形式存储, 所以利用列表的extend方法来将valid_segms添加到其中
        entry['segms'].extend(valid_segms)  
        # gt_classes的类型内一维numpy数组(维度为有效obj的数量), 因此这里不用指定axis的值, 直接按照一维数组拼接即可
        entry['gt_classes'] = np.append(entry['gt_classes'], gt_classes)
        # 同理, 一维numpy数组(维度为有效obj的数量), 无须指定axis的值
        entry['seg_areas'] = np.append(entry['seg_areas'], seg_areas)
        # gt_overlaps为 num_objs × num_classes的numpy数组, 表示每个obj与任意一个类的重叠度
        # 因为entry['gt_overlaps']的类型为scipy.sparse.csr.csr_matrix, 因此这里需要调用toarray方法将其转换成numpy数组, 然后再与gt_overlaps拼接,
        #由于entry['gt_overlaps']的维度为 0 × 81 , 因此拼接后的维度为num_objs × num_classes的numpy数组
        entry['gt_overlaps'] = np.append(
            entry['gt_overlaps'].toarray(), gt_overlaps, axis=0
        )
        # 再将其包装成scipy.sparse.csr.csr_matrix类型
        entry['gt_overlaps'] = scipy.sparse.csr_matrix(entry['gt_overlaps'])
        # 一维numpy数组, 可直接拼接
        entry['is_crowd'] = np.append(entry['is_crowd'], is_crowd)
        # 该列表存储着box的顺序下标值, 同样是一维数组, 直接拼接
        entry['box_to_gt_ind_map'] = np.append(
            entry['box_to_gt_ind_map'], box_to_gt_ind_map
        )
        if self.keypoints is not None:
            entry['gt_keypoints'] = np.append(
                entry['gt_keypoints'], gt_keypoints, axis=0
            )
            entry['has_visible_keypoints'] = im_has_visible_keypoints

_add_proposals_from_file()

# detectron/datasets/json_dataset.py

class JsonDataset(object):
    def __init__(...):
        #...
    def get_roidb(...):
        #...
    def _prep_roidb_entry(self, entry):
        #...
    def _add_gt_annotations(self, entry):
        #...
    #
    def _add_proposals_from_file(
        self, roidb, proposal_file, min_proposal_size, top_k, crowd_thresh
    ):

续解combined_roidb_for_training() 方法

接下来, 重新回到刚才detectron/datasets/roidb.py 文件 的 combined_roidb_for_training 函数中, 继续往下看:

# detectron/datasets/roidb.py

# 加载并连接一个或多个数据集的roidbs, along with optional object proposals
# 每个roidb entry都带有特定的元数据类型, 对其进行准备工作后进行训练
def combined_roidb_for_training(dataset_names, proposal_files):
    def get_roidb(dataset_name, proposal_file): # 注意没有 's'
        # from detectron.datasets.json_dataset import JsonDataset
        # 可以看到, roidb 是利用JsonDataset类对象的get_roidb()方法获取的
        # 注意gt参数是True, 所以表明加载的是训练集的真实数据及其标签
        ds = JsonDataset(dataset_name)
        roidb = ds.get_roidb(
            gt=True,
            proposal_file=proposal_file,
            crowd_filter_thresh=cfg.TRAIN.CROWD_FILTER_THRESH
        )
        # 如果图片翻转属性为真, 则对加载好以后的数据集进行翻转操作
        if cfg.TRAIN.USE_FLIPPED:
            logger.info("Appending horizontally-flipped training examples...")
            extend_with_flipped_entries(roidb, ds)
        logger.info("Loaded dataset: {:s}".format(ds.name))
        # 以上, 数据集加载操作完成, 将roidb数据结构返回
        return roidb
    if isinstance(dataset_names, basestring):
        dataset_names=(dataset_names, )
    if isinstance(proposal_files, basestring):
        proposal_files = (proposal_files, )
    if len(proposal_files) == 0:
        proposal_files = (None, ) * len(dataset_names)
    assert len(dataset_names) == len(proposal_files)
    roidbs = [get_roidb(*args) for args in zip(dataset_names, proposal_files)]
    roidb = roidbs[0]
    for r in roidbs[1:]:
        roidb.extend(r)
    roidb = filter_for_training(roidb)

    logger.info("Computing bounding-box regression targets...")
    # 为训练bounding-box 回归其添加必要的information
    add_bbox_regression_targets(roidb)
    logger.info("done")
    _compute_and_log_stats(roidb)

    return roidb