maskrcnn_benchmark 代码详解之 boxlist_ops.py

前言：

　　与Bounding Box有关的操作有很多，例如对边框列表进行非极大线性抑制、去除过小的边框、计算边框之间的Iou以及对两个边框列表进行合并等操作。在maskrcnn_benchmark中，这些操作都得以实现，其具体代码以及解释如下：

# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import torch

from .bounding_box import BoxList

from maskrcnn_benchmark.layers import nms as _box_nms


def boxlist_nms(boxlist, nms_thresh, max_proposals=-1, score_field="scores"):
    """
    Performs non-maximum suppression on a boxlist, with scores specified
    in a boxlist field via score_field.在边框列表上进行非极大线性抑制

    参数:
        boxlist(BoxList)：含有边框目标得分，原始图像大小，预测边框列表等信息的边框列表
        nms_thresh (float)：非极大线性抑制的阈值
        max_proposals (int): if > 0, then only the top max_proposals are kept
            after non-maximum suppression当大于０时，在非极大线性抑制后只保存前max_proposals个最大目标评分的的预测边框
        score_field (str)：保存目标得分列表的键值
    """
    # 如果非极大线性抑制的阈值小于０，则返回所有边框列表
    if nms_thresh <= 0:
        return boxlist
    # 得到预测边框的格式
    mode = boxlist.mode
    # 　将预测边框转换为左下右上两个点坐标的格式
    boxlist = boxlist.convert("xyxy")
    # 　得到boxlist里的预测边框
    boxes = boxlist.bbox
    # 　得到boxlist里的目标得分
    score = boxlist.get_field(score_field)
    # 进行非极大线性抑制操作，得到应该保留的相应的边框索引
    keep = _box_nms(boxes, score, nms_thresh)
    # 　当大于０时，在非极大线性抑制后只保存前max_proposals个最大目标评分的的预测边框
    if max_proposals > 0:
        keep = keep[: max_proposals]
    # 　保存经过非极大线性抑制操作，得到应该保留的相应的预测边框
    boxlist = boxlist[keep]
    # 将边框格式转换回去，然后返回边框列表
    return boxlist.convert(mode)


def remove_small_boxes(boxlist, min_size):
    """
    Only keep boxes with both sides >= min_size
    只保存边框的所有边都大于min_size的边框
    Arguments:
        boxlist (Boxlist)：需要处理的边框列表
        min_size (int)：所能接受的最小的边框边长
    """
    # TODO maybe add an API for querying the ws / hs把边框转换为中心点坐标＋宽高的格式
    xywh_boxes = boxlist.convert("xywh").bbox
    # 去除xywh_boxes中的第二个维度
    _, _, ws, hs = xywh_boxes.unbind(dim=1)
    # 得到应该保留的边框的索引
    keep = (
        (ws >= min_size) & (hs >= min_size)
    ).nonzero().squeeze(1)
    # 返回满足条件的边框
    return boxlist[keep]


# implementation from https://github.com/kuangliu/torchcv/blob/master/torchcv/utils/box.py
# with slight modifications
def boxlist_iou(boxlist1, boxlist2):
    """Compute the intersection over union of two set of boxes.计算两个边框列表里面的列表之间的重合度
    The box order must be (xmin, ymin, xmax, ymax).所有的边框数据必须为xyxy格式

    Arguments:两个需要计算重合度的边框列表
      box1: (BoxList) bounding boxes, sized [N,4].
      box2: (BoxList) bounding boxes, sized [M,4].

    Returns:
      (tensor) iou, sized [N,M].返回彼此所有边框互相之间的重合度

    Reference:
      https://github.com/chainer/chainercv/blob/master/chainercv/utils/bbox/bbox_iou.py
    """
    # 如果两个边框列表对应的图片大小不同，即对应的是不同的图片，则报错
    if boxlist1.size != boxlist2.size:
        raise RuntimeError(
                "boxlists should have same image size, got {}, {}".format(boxlist1, boxlist2))
    # 得到两个边框列表的长度
    N = len(boxlist1)
    M = len(boxlist2)
    # 得到两个边框列表中边框的面积列表
    area1 = boxlist1.area()
    area2 = boxlist2.area()
    # 得到两个边框列表
    box1, box2 = boxlist1.bbox, boxlist2.bbox
    # 分别计算左下角左边中较大的坐标，右上角坐标中较小的，即两个边框重合的坐标
    lt = torch.max(box1[:, None, :2], box2[:, :2])  # [N,M,2]
    rb = torch.min(box1[:, None, 2:], box2[:, 2:])  # [N,M,2]

    TO_REMOVE = 1
    # 计算重合部分的面积
    wh = (rb - lt + TO_REMOVE).clamp(min=0)  # [N,M,2]
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]
    # 计算重合度
    iou = inter / (area1[:, None] + area2 - inter)
    return iou


# TODO redundant, remove
def _cat(tensors, dim=0):
    """
    Efficient version of torch.cat that avoids a copy if there is only a single element in a list
    是对torch.cat的补充，使得其能避免列表中只存在单个元素，即避免模型没有采用FPN或者只是从单层特征层里提取边框信息
    """
    # 判断tensors是不是boxlist格式
    assert isinstance(tensors, (list, tuple))
    # 如果tensors只有一个元素，直接返回，不再拼接
    if len(tensors) == 1:
        return tensors[0]
    # 如果tensors有多个元素，拼接
    return torch.cat(tensors, dim)


def cat_boxlist(bboxes):
    """
    Concatenates a list of BoxList (having the same image size) into a
    single BoxList。将从不同的FPN特征层里得到的边框以及其对应的目标得分(objectness)并在一起
    他们对应的是同一张图片，同一数据格式
    Arguments:
        bboxes (list[BoxList])：不同的FPN特征层里得到的边框列表
    """
    # 判断bboxes格式是否相同
    assert isinstance(bboxes, (list, tuple))
    assert all(isinstance(bbox, BoxList) for bbox in bboxes)
    # 获得bboxes中所有边框列表对应的同一张图片的大小
    size = bboxes[0].size
    # 如果bboxes中所有边框列表对应不是同一张图片，则报错
    assert all(bbox.size == size for bbox in bboxes)
    # 获得bboxes中所有边框列表的数据格式，xyxy或者ｘｙwh
    mode = bboxes[0].mode
    # 如果bboxes中所有边框列表对应不是同一数据格式，则报错
    assert all(bbox.mode == mode for bbox in bboxes)
    # 获得第一层FPN特征层的bboxes中保存的其他数据的键名
    fields = set(bboxes[0].fields())
    # 如果FPN特征层的bboxes中保存的其他数据的键名不同，则报错
    assert all(set(bbox.fields()) == fields for bbox in bboxes)
    # 将不同的FPN特征层的bboxes的边框合并在一起，并保存为BoxList类型
    cat_boxes = BoxList(_cat([bbox.bbox for bbox in bboxes], dim=0), size, mode)
    # 将不同的FPN特征层的bboxes的边框的与其对应的目标得分(objectness)并在一起
    for field in fields:
        data = _cat([bbox.get_field(field) for bbox in bboxes], dim=0)
        cat_boxes.add_field(field, data)

    return cat_boxes