MaskrcnnBenchmark 源码解析-数据结构(structures)

发表于   |   更新于 | 分类于 | 评论数： | 阅读次数：
本文字数： 17k | 阅读时长 ≈ 15 分钟

源码文件

在 ./maskrcnn_benchmark/structures/ 文件夹中, 定义了许多检测模式使用的数据结构, 如 BoxList, ImageList 等, 下面我们就将对这些数据结构以后适用于它们的操作进行讲解, 涉及到的源码文件如下所示:

bounding_box.py 文件解析

该文件的主要内容就是定义了 class BoxList(object) 类, 该类用于表示一系列的 bounding boxes. 这些 boxes 会以 N×4 大小的 Tensor 来表示. 为了唯一的确定 boxes 在图片中的准确位置, 该类还保存了图片的维度, 另外, 也可以添加额外的信息到特定的 bounding box 中, 如标签信息.

文件概览

该文件篇幅较长, 文件概览如下所示:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
# ./maskrcnn_benchmark/structures/bounding_box.py

class BoxList(object):

    def __init__(self, bbox, image_size, mode="xyxy"):
        # 初始化函数
        # ...

    def add_field(self, field, field_data):
        self.extra_fields[field] = field_data

    def get_field(self, field):
        return self.extra_fields[field]

    def fields(self):
        return list(self.extra_fields.keys())

    def _copy_extra_fields(self, bbox):
        for k, v in bbox.extra_fields.items():
            self.extra_fields[k] = v

    def convert(self, mode):
        # ...

    def _split_into_xyxy(self):
        # ...

    def resize(self, size, *args, **kwargs):
        # ...

    def transpose(self, method):
        # ...

    def crop(self, box):
        # ...

    def to(self, device):
        # ...

    def __getitem__(self, item):
        # ...

    def __len__(self):
        return self.bbox.shape[0]

    def clip_to_image(self, remove_empty=True):
        # ...

    def area(self):
        # ...

    def copy_with_fields(self, fields):
        # ...

    def __repr__(self):
        # ...

class BoxList 解析

下面我们对该类进行解析, 由于这个类的定义比较长, 且这个文件没有外部函数, 所有的函数都是类的成员函数, 因此, 我们将这个类拆分来讲解, 以方便我们查阅.

初始化函数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
# ./maskrcnn_benchmark/structures/bounding_box.py


import torch

# 转换标志
FLIP_LEFT_RIGHT = 0
FLIP_TOP_BOTTOM = 1

class BoxList(object):
    # 该类用于表示一系列的 bounding boxes. 这些 boxes 会以 N×4 大小的 Tensor 来表示.
    # 为了唯一的确定 boxes 在图片中的准确位置, 该类还保存了图片的维度,
    # 另外, 也可以添加额外的信息到特定的 bounding box 中, 如标签信息.

    def __init__(self, bbox, image_size, mode="xyxy"):
        # bbox (tensor): n×4, 代表n个box, 如: [[0,0,10,10],[0,0,5,5]]
        # image_size: (width, height)

        # 根据 bbox 的数据类型获取对应的 device
        device = bbox.device if isinstance(bbox, torch.Tensor) else torch.device("cpu")

        # 将 bbox 转换成 tensor 类型
        bbox = torch.as_tensor(bbox, dtype=torch.float32, device=device)

        # bbox 的维度数量必须为2, 并且第二维必须为 4, 即 shape=(n, 4), 代表 n 个 box
        if bbox.ndimension() != 2:
            raise ValueError(
                "bbox should have 2 dimensions, got {}".format(bbox.ndimension())
            )
        if bbox.size(-1) !=4:
            raise ValueError(
                "last dimension of bbox should have a"
                "size of 4, got {}".format(bbox.size(-1))
            )

        # 只支持以下两种模式
        if mode not in ("xyxy", "xywh"):
            raise ValueError("mode should be 'xyxy' or 'xywh'")

        # 为成员变量赋值
        self.bbox = bbox
        self.size = image_size
        self.mode = mode
        self.extra_fields = {} # 以字典结构存储额外信息

extra_fields 操作函数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
# ./maskrcnn_benchmark/structures/bounding_box.py

class BoxList(object):
    #...

    # 添加新的键值或覆盖旧的键值
    def add_field(self, field, field_data):
            self.extra_fields[field] = field_data

    # 获取指定键对应的值
    def get_field(self, field):
        return self.extra_fields[field]

    # 判断额外信息中是否存在该键
    def has_field(self, field):
        return field in self.extra_fields

    # 以列表的形式返回所有的键的名称
    def fields(self):
        return list(self.extra_fields.keys())

    # 将另一个 BoxList 类型的额外信息(字典)复制到到的额外信息(extra_fields)中.
    def _copy_extra_fields(self, bbox):
        for k, v in bbox.extra_fields.items():
            self.extra_fields[k] = v

convert() 模式转换函数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
# ./maskrcnn_benchmark/structures/bounding_box.py

class BoxList(object):
    #...

    # 将当前 bbox 的表示形式转换成参数指定的模式
    def convert(self, mode):
        # 只支持以下两种模式
        if mode not in ("xyxy", "xywh"):
            raise ValueError("mode should be 'xyxy' or 'xywh'")
        if mode == self.mode:
            return self

        # 调用成员函数, 将坐标表示转化成 (x1,y1,x2,y2) 的形式
        xmin, ymin, xmax, ymax = self._split_into_xyxy()

        if mode == "xyxy":
            # 如果模式为 "xyxy", 则直接将 xmin, ymin, xmax, ymax 合并成 n×4 的 bbox
            bbox = torch.cat((xmin, ymin, xmax, ymax), dim=-1)
            # 这里创建了一个新的 BoxList 实例
            bbox = BoxList(bbox, self.size, mode)
        else:
            # 否则的话, 就将 xmin, ymin, xmax, ymax 转化成 (x,y,w,h) 后再连接在一起
            TO_REMOVE = 1
            torch.cat((xmin, ymin, xmax - xmin + TO_REMOVE, ymax - ymin + TO_REMOVE), dim=-1)
            # 创建新的 BoxList 实例
            bbox = BoxList(bbox, self.size, mode=mode)

        # 复制当前实例的 extra_fields 信息到这个新创建的实例当中, 并将这个新实例返回
        bbox._copy_extra_fields(self)
        return bbox

    # 获取 bbox 的 (x1,y1,x2,y2)形式的坐标表示, .split 为 torch 内置函数, 用于将 tensor 分成多块
    def _split_into_xyxy(self):
        if self.mode == "xyxy":
            # x, y 的 shape 为 n × 1, 代表着 n 个 box 的 x, y 坐标
            xmin, ymin, xmax, ymax = self.split(1, dim=-1)
            return xmin, ymin, xmax, ymax
        elif self.mode == "xywh":
            TO_REMOVE = 1
            # 4 个 tensor 的 shape 均为 n×1
            xmin, ymin, w, h = self.bbox.split(1, dim=-1)
            return(
                xmin,
                ymin,
                xmin + (w - TO_REMOVE).clamp(min=0),
                ymin + (w - TO_REMOVE).clamp(min=0),
            )
        else:
            raise RuntimeError("Should not be here")

resize() 放缩函数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
# ./maskrcnn_benchmark/structures/bounding_box.py

class BoxList(object):
    #...

    # 将所有的 boxes 按照给定的 size 和图片的尺寸进行放缩, 创建一个副本存储放缩后的 boxes 并返回
    def resize(self, *args, **kwargs):
        # size: 指定放缩后的大小 (width, height)

        # 计算宽和高的放缩比例(new_size/old_size)
        ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(size, self.size))

        # 宽高放缩比例相同
        if ratios[0] == ratios[1]:
            ratio = ratios[0]

            # 令所有的 bbox 都乘以放缩比例, 不论 bbox 是以 xyxy 形式还是以 xywh 表示
            # 乘以系数就可以正确的将 bbox 的坐标转化到放缩后图片的对应坐标
            scaled_box = self.box * ratio

            bbox = BoxList(scaled_box, size, mode=self.mode)

            # 复制/转化其他信息
            for k, v in self.extra_fields.items():
                if not isinstance(v, torch.Tensor):
                    v = v.resize(size, *args, **kwargs)
                bbox.add_field(k, v)
            return bbox

        # 宽高的放缩比例不同, 因此, 需要拆分后分别放缩然后在连接在一起
        ratio_width, ratio_height = ratios

        # 获取 bbox 的左上角和右下角的坐标
        xmin, ymin, xmax, ymax = self._split_into_xyxy()

        # 分别对宽 (xmax, xmin) 和高 (ymax, ymin) 进行放缩
        scaled_xmin = xmin * ratio_width
        scaled_xmax = xmax * ratio_width
        scaled_ymin = ymin * ratio_height
        scaled_ymax = ymax * ratio_height

        # 将左上角和右下角坐标连接起来, 组合放缩后的 bbox 表示
        scaled_box = torch.cat(
            (scaled_xmin, scaled_ymin, scaled_xmax, scaled_ymax), dim=-1
        )

        bbox = BoxList(scaled_box, size, mode="xyxy")

        # 复制或转化其他信息
        for k, v in self.extra_fields.items():
            if k not isinstance(v, torch.Tensor):
                v = v.resize(size, *args, **kwargs)

            bbox.add_field(k, v)

        # 将 bbox 转换成指定的模式(因为前面强制转换成 xyxy 模式了, 所这里要转回去)
        retrun bbox.convert(self.mode)

transpose() 图片翻转或旋转

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
# ./maskrcnn_benchmark/structures/bounding_box.py
class BoxList(object):
    #...

    def transpose(self, method):
        # 对 bbox 进行转换(翻转或者旋转90度)
        # methon (int) 此处只能为 0 或 1, 目前仅仅支持两个转换方法
        # FLIP_LEFT_RIGHT = 0, FLIP_TOP_BOTTOM = 1
        # :param method: One of :py:attr:`PIL.Image.FLIP_LEFT_RIGHT`,
        # :py:attr:`PIL.Image.FLIP_TOP_BOTTOM`, :py:attr:`PIL.Image.ROTATE_90`,
        # :py:attr:`PIL.Image.ROTATE_180`, :py:attr:`PIL.Image.ROTATE_270`,
        # :py:attr:`PIL.Image.TRANSPOSE` or :py:attr:`PIL.Image.TRANSVERSE`.

        # 目前仅仅支持 FLIP_LEFT_RIGHT 和 FLIP_TOP_BOTTOM 两种方式
        if method not in (FLIP_LEFT_RIGHT, FLIP_TOP_BOTTOM):
            raise NotImplementedError(
                "Only FLIP_LEFT_RIGHT and FLIP_TOP_BOTTOM implemented"
            )

        # 获取图片的宽和高
        image_width, image_height = self.size

        # 获取左上角和右下角的坐标
        xmin, ymin, xmax, ymax = self._split_into_xyxy()
        if method == FLIP_LEFT_RIGHT: # method=0
            TO_REMOVE = 1
            transposed_xmin = image_width - xmax - TO_REMOVE
            transposed_xmax = image_width - xmin - TO_REMOVE
            transposed_ymin = ymin
            transposed_ymax = ymax
        elif method == FLIP_TOP_BOTTOM:
            transposed_xmin = xmin
            transposed_xmax = xmax
            transposed_ymin = image_height - ymax
            transposed_ymax = image_height - ymin

        # 将转换后的坐标组合起来形成新的 boxes
        transposed_boxes = torch.cat(
            (transposed_xmin, transposed_ymin, transposed_xmax, transposed_ymax), dim=-1
        )

        # 根据转换后的 boxes 坐标创建一个新的 BoxList 实例, 同时将 extra_fields 信息复制
        bbox = BoxList(transposed_boxes, self.size, mode="xyxy")
        for k, v in self.extra_fields.items():
            if not isinstance(v, torch.Tensor):
                v = v.transpose(method)
            bbox.add_field(k, v)

        # 将 bbox 的 mode 转换后返回
        return bbox.convert(self.mode)

crop() 裁剪函数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
# ./maskrcnn_benchmark/structures/bounding_box.py
class BoxList(object):
    #...

    def crop(self, box):
        # box 是一个4元组, 指定了希望剪裁的区域的左上角和右下角

        # 获取当前所有 boxes 的最左, 最上, 最下, 最右的坐标
        xmin, ymin, xmax, ymax = self._split_into_xyxy()

        # 获取欲剪裁的 box 的宽和高
        w, h = box[2] - box[0], box[3] - box[1]

        # 根据 box 指定的区域, 对所有的 proposals boxes 进行剪裁
        # 即改变其坐标的位置, 如果发现有超出规定尺寸的情况, 则将其截断
        cropped_xmin = (xmin - box[0]).clamp(min=0, max=w)
        cropped_ymin = (ymin - box[1]).clamp(min=0, max=h)
        cropped_xmax = (xmax - box[0]).clamp(min=0, max=w)
        cropped_ymax = (ymax - box[1]).clamp(min=0, max=h)

        # 将新的剪裁后的 box 坐标连接起来创建一个新的 BoxList 实例
        bbox = BoxList(cropped_box, (w, h), mode="xyxy")
        # 复制其他信息
        for k, v in self.extra_fields.items():
            if not isinstance(v, torch.Tensor):
                v = v.crop(box)
            bbox.add_field(k, v)

        # 将 bbox 的模式转换成传入时的模式
        return bbox.convert(self.mode)

to() 设备转移函数

将 BoxList 实例转化到指定的设备上(创建新实例返回).

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# ./maskrcnn_benchmark/structures/bounding_box.py
class BoxList(object):
    #...

    def to(self, device):
        # device: "cuda:x" or "cpu"
        # 将当前的 bbox 移动到指定的 device 上, 并且重新创建一个新的 BoxList 实例
        bbox = BoxList(self.bbox.to(device), self.size, self.mode)
        # 深度复制
        for k, v in self.extra_fields.items():
            if hasattr(v, "to"):
                v = v.to(device)
            bbox.add_field(k, v)
        return bbox

clip_to_image()

该函数将 bbox 的坐标限制在 image 的尺寸内, 以便可以将边框显示在 image uh.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# ./maskrcnn_benchmark/structures/bounding_box.py
class BoxList(object):
    #...

    def clip_to_image(self, remove_empty=True):
        TO_REMOVE = 1
        self.bbox[:, 0].clamp_(min=0, max=self.size[0] - TO_REMOVE)
        self.bbox[:, 1].clamp_(min=0, max=self.size[1] - TO_REMOVE)
        self.bbox[:, 2].clamp_(min=0, max=self.size[0] - TO_REMOVE)
        self.bbox[:, 3].clamp_(min=0, max=self.size[1] - TO_REMOVE)

        if remove_empty:
            box = self.bbox

            # 该语句会返回一个 n×1 的列表, 对应着 n 个 box, 如果 box 的坐标满足
            # 下面的语句条件, 则对应位为 1, 否则 为 0,
            keep = (box[:, 3] > box[:, 1]) & (box[:, 2] > box[:, 0])

            # 返回那些对应位为 1 的 box
            return self[keep]

        return self

area() 获取区域面积函数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# ./maskrcnn_benchmark/structures/bounding_box.py
class BoxList(object):
    #...

    def area(self):
        box = self.bbox
        if self.mode = "xyxy":
            TO_REMOVE = 1
            # 一个像素点的面积我们认为是 1, 而不是 0
            area = (box[:, 2] - box[:, 0] + TO_REMOVE) * (box[:, 3] - box[:, 1] + TO_REMOVE)
        elif self.mode == "xywh":
            # 直接令 w * h
            area = box[:, 2] * box[:, 3]
        else:
            raise RuntimeError("Should not be here")

copy_with_fields() 深度复制函数

连带 BoxList 的 extra_fields 信息进行复制, 返回一个新的 BoxList 实例

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# ./maskrcnn_benchmark/structures/bounding_box.py
class BoxList(object):
    #...

    def copy_with_fields(self, fields):
        bbox = BoxList(self.bbox, self.size, self.mode)
        if not isinstance(fields, (list, tuple)):
            # 将 fields 包裹在列表里面, 注意 [fields] 和 list(fields) 的区别
            fields = [fields]

        # 遍历 fields 中的所有元素, 并将其添加到当前的 BoxList 实例 bbox 中
        for field in fields:
            bbox.add_field(field, self.get_field(field))

        return  bbox

其他函数

getitem(), len(), repr()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
# ./maskrcnn_benchmark/structures/bounding_box.py
class BoxList(object):
        #...

        def __getitem__(self, item):
            # item 必须是列表类型
            # 假设 bbox 是一个 BoxList 实例, 那么我们可以利用下面语句得到该实例的子集
            # sub_bbox = bbox[[0,3,4,8]]
            # one_bbox = bbox[[2]]

            # 创建新的子集实例
            bbox = BoxList(self.bbox[item], self.size, self.mode)
            # 复制其他信息
            for k, v in self.extra_fields.items():
                bbox.add_field(k, v[item])

            return bbox

        def __len__(self):
            # 获取当前 BoxList 中含有的 box 的数量
            return self.bbox.shape[0]

        def __repr__(self):
            # 改变 print(BoxList_a) 的打印信息, 使之显示更多的有用信息, 示例如下:
            # BoxList(num_boxes=2, image_width=10, image_height=10, mode=xyxy)
            s = self.__class__.__name__ + "("
            s += "num_boxes={}, ".format(len(self))
            s += "image_width={}, ".format(self.size[0])
            s += "image_height={}, ".format(self.size[1])
            s += "mode={})".format(self.mode)

            return s

boxlist_ops.py

该文件定义了一些与 BoxList 类型数据有关的操作, 根据函数的相互调用关系, 可以分为以下四个主要的函数:

  • boxlist_nms():
  • remove_small_boxes():
  • boxlist_iou():
  • cat_boxlist():

下面我们按照顺序对分别上面的函数进行解析

boxlist_nms() 函数

该函数会对一个 BoxList 类型数据中的 box 执行非极大值抑制算法, 这些 box 的 socre 值可以通过 score_field 获取, 代码的核心逻辑实际上是调用了 ./maskrcnn_benchmark/layers/nms.py 文件中的 _box_nms() 函数, 关于该函数的详细解析可以看 nms 解析

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
# ./maskrcnn_benchmark/structures/boxlist_ops.py

def boxlist_nms(boxlist, nms_thresh, max_proposals=-1, score_field="score"):
    # boxlist (BoxList)
    # nms_thresh (float)
    # max_proposals (int): top-k
    # score_field (str): 键

    if nms_thresh <= 0:
        return boxlist

    mode = boxlist.mode # 缓存当前的模式
    boxlist = boxlist.convert("xyxy") # 转换成指定模式
    boxes = boxlist.bbox # 获取 n*4 的 bbox 列表
    score = boxlist.get_field(score_field) # 获取对应的 socre 列表

    # 调用 _box_nms 执行非极大值抑制抑制
    keep = _box_nms(boxes, score, nms_thresh)
    if max_proposals > 0: # 只保留 top-k
        keep = keep[: max_proposals]

    # keep为下标列表, 指示了需要保存哪些box, 这里已经重写了 __getitem__ 方法, 因此会根据下标返回 BoxList
    boxlist = boxlist[keep]

    return boxlist.convert[mode]

remove_small_boxes() 函数

该函数执行后, BoxList 中只会保留那些尺寸大于一定值的区域 box.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# ./maskrcnn_benchmark/structures/boxlist_ops.py

def remove_small_boxes(boxlist, min_size):
    # boxlist (BoxList)
    # min_size (int)

    # 获取 xywh 形式的bbox列表
    xywh_boxes = boxlist.convert("xywh").bbox

    # torch 的 unbind 函数, 用于 "移除" tensor 的维度,
    _, _, ws, hs = xywh_boxes.unbind(dim=1)

    keep = (
        (ws >= min_size) & (hs >= min_size)
    ).()

boxlist_iou() 函数

cat_boxlist() 函数

该函数会将一个组成元素为 BoxList 的列表合并成一个 BoxList 对象. 注意, 列表中的 BoxList 对象必须具有相同 image size.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
# ./maskrcnn_benchmark/structures/boxlist.py

def cat_boxlist(bboxes):
    # bboxes (list[BoxList])

    # 确保类型为列表或元组, 且其中元素类型为 BoxList
    assert isinstance(bboxes, (list, tuple))
    assert all(isinstance(bbox, BoxList) for bbox in bboxes)

    # 确保所有的 BoxList 的 size , mode, 以及 extra_fields 字典的 keys 是相同的
    size = bboxes[0].size
    assert all(bbox.size == size for bbox in bboxes)

    mode = bboxes[0].mode
    assert all(bbox.mode == mode for bbox in bboxes)

    fields = set(bboxes[0].fields()) # 获取字典的所有 key 值
    assert all(set(bbox.fields()) == fields for bbox in bboxes)

    # 调用本文件的 _cat() 方法, 将 bboxes 里面的 BoxList 数据连接成一个 BoxList, 具体解析看下方
    cat_boxes = BoxList(_cat([bbox.bbox for bbox in boxes], dim=0), size, mode)

    # 将各个 BoxList 的 fields 补充上
    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

def _cat(tensors, dim=0):
    # 调用 torch.cat 将数据连接
    assert isinstance(tenosrs, (list, tuple))
    if len(tensors) == 1:
        return tensors[0]
    return torch.cat(tensors, dim)

image_list.py

segmentation_mask.py