mmdetection之MaxIoUAssigner

MaxIoUAssigner的作用是根据bboxes和gt_bboxes的最大IoU对bboxes分配标签。所有的Assigner都继承自BaseAssigner,实现assign函数。

一、计算IoU

下面代码会计算gt_bboxes和bboxes之间的IoU，需要注意的是，把gt_bboxes放在bboxes前面时，overlaps每一行表示的是每一个gt_bbox与所有bbox的IoU，每一列表示的是每一个bbox与所有gt_bbox的IoU。

 overlaps = self.iou_calculator(gt_bboxes, bboxes)

二、根据iof设置忽略样本

if (self.ignore_iof_thr > 0 and gt_bboxes_ignore is not None
        and gt_bboxes_ignore.numel() > 0 and bboxes.numel() > 0):
    if self.ignore_wrt_candidates:
        ignore_overlaps = self.iou_calculator(
            bboxes, gt_bboxes_ignore, mode='iof')
        ignore_max_overlaps, _ = ignore_overlaps.max(dim=1)
    else:
        ignore_overlaps = self.iou_calculator(
            gt_bboxes_ignore, bboxes, mode='iof')
        ignore_max_overlaps, _ = ignore_overlaps.max(dim=0)
    overlaps[:, ignore_max_overlaps > self.ignore_iof_thr] = -1

$$IOF=\frac{AreaofOverlap}{AreaofGroundTruth}$$

ignore_wrt_candidates控制着ground truth是bboxes还是gt_bboxes_ignore

三、调用assign_wrt_overlaps为bboxes分配label

1、首先需要有一个记录bboxes标记（正样本、负样本、忽略样本）的数组。

assigned_gt_inds = overlaps.new_full((num_bboxes,), -1,
                                     dtype=torch.long)

形状为一个1D（num_bboxes,），num_bboxes为bbox的数量，通过下面代码获取：

num_gts, num_bboxes = overlaps.size(0), overlaps.size(1)

初始值为-1，代表全部为忽略样本

2、计算每一个bbox与所有gt_bbox的IoU的最大值

max_overlaps, argmax_overlaps = overlaps.max(dim=0)

dim=0可以获取每一列的最大值，返回值为一个一维数组max_overlaps记录最大IoU，一个一维数组argmax_overlaps 记录此最大IoU的列坐标

3、当没有gt_bbox时，将所有bbox划分为负样本

if num_gts == 0 or num_bboxes == 0:
    max_overlaps = overlaps.new_zeros((num_bboxes,))
    if num_gts == 0:
        assigned_gt_inds[:] = 0
    if gt_labels is None:
        assigned_labels = None
    else:
        assigned_labels = overlaps.new_full((num_bboxes,),
                                            -1,
                                            dtype=torch.long)
    return AssignResult(
        num_gts,
        assigned_gt_inds,
        max_overlaps,
        labels=assigned_labels)

当bbox的数量为0时，相当于没有样本，assigned_gt_inds所有值为-1（忽略样本），通常不会发生这种情况
当gt_bbox数量为0时，将所有bbox标记为负样本

4、将最大IoU在[0，neg_iou_thr）之间的bbox标记为负样本

if isinstance(self.neg_iou_thr, float):
    # 将0<iou<neg_iou_thr的样本设为负样本，label为0
    assigned_gt_inds[
        (max_overlaps >= 0) & (max_overlaps < self.neg_iou_thr)] = 0
elif isinstance(self.neg_iou_thr, tuple):
    # 将neg_iou_thr[0]<iou<neg_iou_thr[1]的样本设为负样本，label为0
    assert len(self.neg_iou_thr) == 2
    assigned_gt_inds[(max_overlaps >= self.neg_iou_thr[0]) & (
            max_overlaps <= self.neg_iou_thr[1])] = 0

当neg_iou_thr为浮点数时，将最大IoU在[0，neg_iou_thr）之间的bbox标记为负样本
当neg_iou_thr为元组时，将最大IoU在[neg_iou_thr[0]，neg_iou_thr[1]）之间的bbox标记为负样本

5、将最大IoU大于等于pos_iou_thr的bbox标记为正样本

pos_inds = max_overlaps >= self.pos_iou_thr
assigned_gt_inds[pos_inds] = argmax_overlaps[pos_inds] + 1

这里对坐标+1是因为0已经分配给了负样本

6、开启低质量匹配模式

if self.match_low_quality:
    for i in range(num_gts):
        if gt_max_overlaps[i] >= self.min_pos_iou:
            if self.gt_max_assign_all:
                max_iou_inds = overlaps[i, :] == gt_max_overlaps[i]
                assigned_gt_inds[max_iou_inds] = i + 1
            else:
                assigned_gt_inds[gt_argmax_overlaps[i]] = i + 1

高质量匹配模式会为每一个bbox匹配到相应的gt_bbox,但是会存在有的gt_bbox匹配不到bbox的情况
此时，开启低质量模式时会为每一个gt_bbox匹配一个IoU最大的bbox，如果IoU>=min_pos_iou就将此bboxes与此gt_bboxes进行匹配
低质量模式时会出现覆盖问题，例如有2个gt_bbox,1个bbox，iou为[0.9,0.8]，高质量匹配模式下会将bbox与gt1进行匹配，开启低质量模式时后，gt_bbox1会匹配到bbox1，gt_bbox2也会匹配到bbox1，此时会出现覆盖问题，将bbox与gt_bbox2进行匹配，此时IoU不如原来的高，因此匹配的质量会比较低。
当使用bbox匹配gt_bbox时，最大的IoU可能会对应好几个gt_bbox，不能将bbox与所有的gt_bbox进行匹配，我们将bbox与第一个gt_bbox进行匹配。
但是，使用gt_bbox匹配bbox时，最大的IoU对应好几个bbox，我们可以将这些bbox都与这个gt_bbox进行匹配。
此时使用gt_max_assign_all来决定是否将所有的bbox都与这个gt_bbox进行匹配。

7、为正样本分配标签

if gt_labels is not None:
    assigned_label = assigned_gt_inds.new_full((num_bboxes,), -1)
    # 返回大于0元素的坐标
    pos_inds = torch.nonzero(assigned_gt_inds > 0, as_tuple=False)
    if pos_inds.numel() > 0:
        # -1,对应的gt的坐标
        assigned_label[pos_inds] = gt_labels[
            assigned_gt_inds[pos_inds] - 1]
else:
    assigned_label = None

上面，我们将bbox划分为了正样本（>0）,负样本（0），忽略样本（-1），现在我们要为正样本分配label。
assigned_gt_inds[pos_inds]表示正样本匹配到的gt_bbox，-1表示gt_bbox的坐标，gt_labels[assigned_gt_inds[pos_inds] - 1]会索引gt_bbox对应的label。

8、返回结果

return AssignResult(
    num_gts, assigned_gt_inds, max_overlaps, labels=assigned_labels)