iou3d_calculator.py 12.5 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
2
3
import torch

4
from mmdet3d.registry import TASK_UTILS
zhangshilong's avatar
zhangshilong committed
5
6
from mmdet3d.structures.bbox_3d import get_box_type
from mmdet.structures.bbox import bbox_overlaps
zhangwenwei's avatar
zhangwenwei committed
7
8


9
@TASK_UTILS.register_module()
zhangwenwei's avatar
zhangwenwei committed
10
class BboxOverlapsNearest3D(object):
zhangwenwei's avatar
zhangwenwei committed
11
    """Nearest 3D IoU Calculator.
zhangwenwei's avatar
zhangwenwei committed
12
13
14

    Note:
        This IoU calculator first finds the nearest 2D boxes in bird eye view
15
        (BEV), and then calculates the 2D IoU using :meth:`bbox_overlaps`.
zhangwenwei's avatar
zhangwenwei committed
16
17

    Args:
wangtai's avatar
wangtai committed
18
        coordinate (str): 'camera', 'lidar', or 'depth' coordinate system.
zhangwenwei's avatar
zhangwenwei committed
19
20
21
22
23
    """

    def __init__(self, coordinate='lidar'):
        assert coordinate in ['camera', 'lidar', 'depth']
        self.coordinate = coordinate
zhangwenwei's avatar
zhangwenwei committed
24
25

    def __call__(self, bboxes1, bboxes2, mode='iou', is_aligned=False):
zhangwenwei's avatar
zhangwenwei committed
26
        """Calculate nearest 3D IoU.
zhangwenwei's avatar
zhangwenwei committed
27
28
29
30
31
32
33

        Note:
            If ``is_aligned`` is ``False``, then it calculates the ious between
            each bbox of bboxes1 and bboxes2, otherwise it calculates the ious
            between each aligned pair of bboxes1 and bboxes2.

        Args:
34
            bboxes1 (torch.Tensor): shape (N, 7+N)
35
                [x, y, z, x_size, y_size, z_size, ry, v].
36
            bboxes2 (torch.Tensor): shape (M, 7+N)
37
                [x, y, z, x_size, y_size, z_size, ry, v].
zhangwenwei's avatar
zhangwenwei committed
38
39
            mode (str): "iou" (intersection over union) or iof
                (intersection over foreground).
wangtai's avatar
wangtai committed
40
            is_aligned (bool): Whether the calculation is aligned.
zhangwenwei's avatar
zhangwenwei committed
41
42

        Return:
43
44
            torch.Tensor: If ``is_aligned`` is ``True``, return ious between
                bboxes1 and bboxes2 with shape (M, N). If ``is_aligned`` is
zhangwenwei's avatar
zhangwenwei committed
45
46
                ``False``, return shape is M.
        """
zhangwenwei's avatar
zhangwenwei committed
47
48
        return bbox_overlaps_nearest_3d(bboxes1, bboxes2, mode, is_aligned,
                                        self.coordinate)
zhangwenwei's avatar
zhangwenwei committed
49
50

    def __repr__(self):
wangtai's avatar
wangtai committed
51
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
52
        repr_str = self.__class__.__name__
zhangwenwei's avatar
zhangwenwei committed
53
        repr_str += f'(coordinate={self.coordinate}'
zhangwenwei's avatar
zhangwenwei committed
54
55
56
        return repr_str


57
@TASK_UTILS.register_module()
zhangwenwei's avatar
zhangwenwei committed
58
class BboxOverlaps3D(object):
zhangwenwei's avatar
zhangwenwei committed
59
    """3D IoU Calculator.
zhangwenwei's avatar
zhangwenwei committed
60

61
    Args:
zhangwenwei's avatar
zhangwenwei committed
62
63
        coordinate (str): The coordinate system, valid options are
            'camera', 'lidar', and 'depth'.
64
65
66
    """

    def __init__(self, coordinate):
zhangwenwei's avatar
zhangwenwei committed
67
        assert coordinate in ['camera', 'lidar', 'depth']
68
69
70
        self.coordinate = coordinate

    def __call__(self, bboxes1, bboxes2, mode='iou'):
zhangwenwei's avatar
zhangwenwei committed
71
        """Calculate 3D IoU using cuda implementation.
zhangwenwei's avatar
zhangwenwei committed
72
73
74
75
76
77
78

        Note:
            This function calculate the IoU of 3D boxes based on their volumes.
            IoU calculator ``:class:BboxOverlaps3D`` uses this function to
            calculate the actual 3D IoUs of boxes.

        Args:
79
80
81
82
            bboxes1 (torch.Tensor): with shape (N, 7+C),
                (x, y, z, x_size, y_size, z_size, ry, v*).
            bboxes2 (torch.Tensor): with shape (M, 7+C),
                (x, y, z, x_size, y_size, z_size, ry, v*).
zhangwenwei's avatar
zhangwenwei committed
83
84
85
86
            mode (str): "iou" (intersection over union) or
                iof (intersection over foreground).

        Return:
87
            torch.Tensor: Bbox overlaps results of bboxes1 and bboxes2
zhangwenwei's avatar
zhangwenwei committed
88
89
                with shape (M, N) (aligned mode is not supported currently).
        """
90
        return bbox_overlaps_3d(bboxes1, bboxes2, mode, self.coordinate)
zhangwenwei's avatar
zhangwenwei committed
91
92

    def __repr__(self):
zhangwenwei's avatar
zhangwenwei committed
93
        """str: return a string that describes the module"""
zhangwenwei's avatar
zhangwenwei committed
94
        repr_str = self.__class__.__name__
zhangwenwei's avatar
zhangwenwei committed
95
        repr_str += f'(coordinate={self.coordinate}'
zhangwenwei's avatar
zhangwenwei committed
96
97
98
        return repr_str


zhangwenwei's avatar
zhangwenwei committed
99
100
101
102
103
def bbox_overlaps_nearest_3d(bboxes1,
                             bboxes2,
                             mode='iou',
                             is_aligned=False,
                             coordinate='lidar'):
zhangwenwei's avatar
zhangwenwei committed
104
    """Calculate nearest 3D IoU.
zhangwenwei's avatar
zhangwenwei committed
105

zhangwenwei's avatar
zhangwenwei committed
106
107
    Note:
        This function first finds the nearest 2D boxes in bird eye view
wangtai's avatar
wangtai committed
108
        (BEV), and then calculates the 2D IoU using :meth:`bbox_overlaps`.
109
        This IoU calculator :class:`BboxOverlapsNearest3D` uses this
zhangwenwei's avatar
zhangwenwei committed
110
111
112
113
114
115
        function to calculate IoUs of boxes.

        If ``is_aligned`` is ``False``, then it calculates the ious between
        each bbox of bboxes1 and bboxes2, otherwise the ious between each
        aligned pair of bboxes1 and bboxes2.

zhangwenwei's avatar
zhangwenwei committed
116
    Args:
117
118
119
120
        bboxes1 (torch.Tensor): with shape (N, 7+C),
            (x, y, z, x_size, y_size, z_size, ry, v*).
        bboxes2 (torch.Tensor): with shape (M, 7+C),
            (x, y, z, x_size, y_size, z_size, ry, v*).
liyinhao's avatar
liyinhao committed
121
        mode (str): "iou" (intersection over union) or iof
zhangwenwei's avatar
zhangwenwei committed
122
            (intersection over foreground).
zhangwenwei's avatar
zhangwenwei committed
123
        is_aligned (bool): Whether the calculation is aligned
zhangwenwei's avatar
zhangwenwei committed
124
125

    Return:
126
127
        torch.Tensor: If ``is_aligned`` is ``True``, return ious between
            bboxes1 and bboxes2 with shape (M, N). If ``is_aligned`` is
zhangwenwei's avatar
zhangwenwei committed
128
            ``False``, return shape is M.
zhangwenwei's avatar
zhangwenwei committed
129
    """
zhangwenwei's avatar
zhangwenwei committed
130
    assert bboxes1.size(-1) == bboxes2.size(-1) >= 7
zhangwenwei's avatar
zhangwenwei committed
131

wuyuefeng's avatar
Demo  
wuyuefeng committed
132
    box_type, _ = get_box_type(coordinate)
zhangwenwei's avatar
zhangwenwei committed
133
134
135

    bboxes1 = box_type(bboxes1, box_dim=bboxes1.shape[-1])
    bboxes2 = box_type(bboxes2, box_dim=bboxes2.shape[-1])
zhangwenwei's avatar
zhangwenwei committed
136
137
138
139

    # Change the bboxes to bev
    # box conversion and iou calculation in torch version on CUDA
    # is 10x faster than that in numpy version
zhangwenwei's avatar
zhangwenwei committed
140
141
142
    bboxes1_bev = bboxes1.nearest_bev
    bboxes2_bev = bboxes2.nearest_bev

zhangwenwei's avatar
zhangwenwei committed
143
144
145
146
147
    ret = bbox_overlaps(
        bboxes1_bev, bboxes2_bev, mode=mode, is_aligned=is_aligned)
    return ret


148
def bbox_overlaps_3d(bboxes1, bboxes2, mode='iou', coordinate='camera'):
zhangwenwei's avatar
zhangwenwei committed
149
    """Calculate 3D IoU using cuda implementation.
zhangwenwei's avatar
zhangwenwei committed
150

zhangwenwei's avatar
zhangwenwei committed
151
    Note:
152
        This function calculates the IoU of 3D boxes based on their volumes.
zhangwenwei's avatar
zhangwenwei committed
153
        IoU calculator :class:`BboxOverlaps3D` uses this function to
zhangwenwei's avatar
zhangwenwei committed
154
155
        calculate the actual IoUs of boxes.

zhangwenwei's avatar
zhangwenwei committed
156
    Args:
157
158
159
160
        bboxes1 (torch.Tensor): with shape (N, 7+C),
            (x, y, z, x_size, y_size, z_size, ry, v*).
        bboxes2 (torch.Tensor): with shape (M, 7+C),
            (x, y, z, x_size, y_size, z_size, ry, v*).
liyinhao's avatar
liyinhao committed
161
        mode (str): "iou" (intersection over union) or
zhangwenwei's avatar
zhangwenwei committed
162
            iof (intersection over foreground).
liyinhao's avatar
liyinhao committed
163
        coordinate (str): 'camera' or 'lidar' coordinate system.
zhangwenwei's avatar
zhangwenwei committed
164
165

    Return:
166
        torch.Tensor: Bbox overlaps results of bboxes1 and bboxes2
zhangwenwei's avatar
zhangwenwei committed
167
            with shape (M, N) (aligned mode is not supported currently).
zhangwenwei's avatar
zhangwenwei committed
168
    """
zhangwenwei's avatar
zhangwenwei committed
169
    assert bboxes1.size(-1) == bboxes2.size(-1) >= 7
wuyuefeng's avatar
Demo  
wuyuefeng committed
170
171

    box_type, _ = get_box_type(coordinate)
zhangwenwei's avatar
zhangwenwei committed
172
173
174
175
176

    bboxes1 = box_type(bboxes1, box_dim=bboxes1.shape[-1])
    bboxes2 = box_type(bboxes2, box_dim=bboxes2.shape[-1])

    return bboxes1.overlaps(bboxes1, bboxes2, mode=mode)
177
178


179
@TASK_UTILS.register_module()
180
181
182
183
184
185
186
187
class AxisAlignedBboxOverlaps3D(object):
    """Axis-aligned 3D Overlaps (IoU) Calculator."""

    def __call__(self, bboxes1, bboxes2, mode='iou', is_aligned=False):
        """Calculate IoU between 2D bboxes.

        Args:
            bboxes1 (Tensor): shape (B, m, 6) in <x1, y1, z1, x2, y2, z2>
188
                format or empty.
189
            bboxes2 (Tensor): shape (B, n, 6) in <x1, y1, z1, x2, y2, z2>
190
                format or empty.
191
                B indicates the batch dim, in shape (B1, B2, ..., Bn).
192
                If ``is_aligned`` is ``True``, then m and n must be equal.
193
194
195
            mode (str): "iou" (intersection over union) or "giou" (generalized
                intersection over union).
            is_aligned (bool, optional): If True, then m and n must be equal.
196
                Defaults to False.
197
        Returns:
198
            Tensor: shape (m, n) if ``is_aligned`` is False else shape (m,)
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
        """
        assert bboxes1.size(-1) == bboxes2.size(-1) == 6
        return axis_aligned_bbox_overlaps_3d(bboxes1, bboxes2, mode,
                                             is_aligned)

    def __repr__(self):
        """str: a string describing the module"""
        repr_str = self.__class__.__name__ + '()'
        return repr_str


def axis_aligned_bbox_overlaps_3d(bboxes1,
                                  bboxes2,
                                  mode='iou',
                                  is_aligned=False,
                                  eps=1e-6):
    """Calculate overlap between two set of axis aligned 3D bboxes. If
216
    ``is_aligned`` is ``False``, then calculate the overlaps between each bbox
217
218
219
220
221
    of bboxes1 and bboxes2, otherwise the overlaps between each aligned pair of
    bboxes1 and bboxes2.

    Args:
        bboxes1 (Tensor): shape (B, m, 6) in <x1, y1, z1, x2, y2, z2>
222
            format or empty.
223
        bboxes2 (Tensor): shape (B, n, 6) in <x1, y1, z1, x2, y2, z2>
224
            format or empty.
225
            B indicates the batch dim, in shape (B1, B2, ..., Bn).
226
            If ``is_aligned`` is ``True``, then m and n must be equal.
227
228
229
        mode (str): "iou" (intersection over union) or "giou" (generalized
            intersection over union).
        is_aligned (bool, optional): If True, then m and n must be equal.
230
            Defaults to False.
231
        eps (float, optional): A value added to the denominator for numerical
232
            stability. Defaults to 1e-6.
233
234

    Returns:
235
        Tensor: shape (m, n) if ``is_aligned`` is False else shape (m,)
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260

    Example:
        >>> bboxes1 = torch.FloatTensor([
        >>>     [0, 0, 0, 10, 10, 10],
        >>>     [10, 10, 10, 20, 20, 20],
        >>>     [32, 32, 32, 38, 40, 42],
        >>> ])
        >>> bboxes2 = torch.FloatTensor([
        >>>     [0, 0, 0, 10, 20, 20],
        >>>     [0, 10, 10, 10, 19, 20],
        >>>     [10, 10, 10, 20, 20, 20],
        >>> ])
        >>> overlaps = axis_aligned_bbox_overlaps_3d(bboxes1, bboxes2)
        >>> assert overlaps.shape == (3, 3)
        >>> overlaps = bbox_overlaps(bboxes1, bboxes2, is_aligned=True)
        >>> assert overlaps.shape == (3, )
    Example:
        >>> empty = torch.empty(0, 6)
        >>> nonempty = torch.FloatTensor([[0, 0, 0, 10, 9, 10]])
        >>> assert tuple(bbox_overlaps(empty, nonempty).shape) == (0, 1)
        >>> assert tuple(bbox_overlaps(nonempty, empty).shape) == (1, 0)
        >>> assert tuple(bbox_overlaps(empty, empty).shape) == (0, 0)
    """

    assert mode in ['iou', 'giou'], f'Unsupported mode {mode}'
261
    # Either the boxes are empty or the length of boxes's last dimension is 6
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
    assert (bboxes1.size(-1) == 6 or bboxes1.size(0) == 0)
    assert (bboxes2.size(-1) == 6 or bboxes2.size(0) == 0)

    # Batch dim must be the same
    # Batch dim: (B1, B2, ... Bn)
    assert bboxes1.shape[:-2] == bboxes2.shape[:-2]
    batch_shape = bboxes1.shape[:-2]

    rows = bboxes1.size(-2)
    cols = bboxes2.size(-2)
    if is_aligned:
        assert rows == cols

    if rows * cols == 0:
        if is_aligned:
            return bboxes1.new(batch_shape + (rows, ))
        else:
            return bboxes1.new(batch_shape + (rows, cols))

    area1 = (bboxes1[..., 3] -
             bboxes1[..., 0]) * (bboxes1[..., 4] - bboxes1[..., 1]) * (
                 bboxes1[..., 5] - bboxes1[..., 2])
    area2 = (bboxes2[..., 3] -
             bboxes2[..., 0]) * (bboxes2[..., 4] - bboxes2[..., 1]) * (
                 bboxes2[..., 5] - bboxes2[..., 2])

    if is_aligned:
        lt = torch.max(bboxes1[..., :3], bboxes2[..., :3])  # [B, rows, 3]
        rb = torch.min(bboxes1[..., 3:], bboxes2[..., 3:])  # [B, rows, 3]

        wh = (rb - lt).clamp(min=0)  # [B, rows, 2]
        overlap = wh[..., 0] * wh[..., 1] * wh[..., 2]

        if mode in ['iou', 'giou']:
            union = area1 + area2 - overlap
        else:
            union = area1
        if mode == 'giou':
            enclosed_lt = torch.min(bboxes1[..., :3], bboxes2[..., :3])
            enclosed_rb = torch.max(bboxes1[..., 3:], bboxes2[..., 3:])
    else:
        lt = torch.max(bboxes1[..., :, None, :3],
                       bboxes2[..., None, :, :3])  # [B, rows, cols, 3]
        rb = torch.min(bboxes1[..., :, None, 3:],
                       bboxes2[..., None, :, 3:])  # [B, rows, cols, 3]

        wh = (rb - lt).clamp(min=0)  # [B, rows, cols, 3]
        overlap = wh[..., 0] * wh[..., 1] * wh[..., 2]

        if mode in ['iou', 'giou']:
            union = area1[..., None] + area2[..., None, :] - overlap
        if mode == 'giou':
            enclosed_lt = torch.min(bboxes1[..., :, None, :3],
                                    bboxes2[..., None, :, :3])
            enclosed_rb = torch.max(bboxes1[..., :, None, 3:],
                                    bboxes2[..., None, :, 3:])

    eps = union.new_tensor([eps])
    union = torch.max(union, eps)
    ious = overlap / union
    if mode in ['iou']:
        return ious
    # calculate gious
    enclose_wh = (enclosed_rb - enclosed_lt).clamp(min=0)
    enclose_area = enclose_wh[..., 0] * enclose_wh[..., 1] * enclose_wh[..., 2]
    enclose_area = torch.max(enclose_area, eps)
    gious = ious - (enclose_area - union) / enclose_area
    return gious