"example/60_gemm_multiABD/CMakeLists.txt" did not exist on "a8780c322cf4993b826b0dcc0842a18c73c3f429"
base_box3d.py 25.6 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
2
import warnings
zhangwenwei's avatar
zhangwenwei committed
3
from abc import abstractmethod
4
from typing import Iterator, Optional, Sequence, Tuple, Union
5

6
7
import numpy as np
import torch
8
from mmcv.ops import box_iou_rotated, points_in_boxes_all, points_in_boxes_part
9
from torch import Tensor
10

11
from mmdet3d.structures.points import BasePoints
12
from .utils import limit_period
zhangwenwei's avatar
zhangwenwei committed
13

14

15
class BaseInstance3DBoxes:
zhangwenwei's avatar
zhangwenwei committed
16
    """Base class for 3D Boxes.
17

zhangwenwei's avatar
zhangwenwei committed
18
    Note:
19
20
        The box is bottom centered, i.e. the relative position of origin in the
        box is (0.5, 0.5, 0).
zhangwenwei's avatar
zhangwenwei committed
21

zhangwenwei's avatar
zhangwenwei committed
22
    Args:
23
24
25
26
27
28
29
        tensor (Tensor or np.ndarray or Sequence[Sequence[float]]): The boxes
            data with shape (N, box_dim).
        box_dim (int): Number of the dimension of a box. Each row is
            (x, y, z, x_size, y_size, z_size, yaw). Defaults to 7.
        with_yaw (bool): Whether the box is with yaw rotation. If False, the
            value of yaw will be set to 0 as minmax boxes. Defaults to True.
        origin (Tuple[float]): Relative position of the box origin.
30
            Defaults to (0.5, 0.5, 0). This will guide the box be converted to
wuyuefeng's avatar
wuyuefeng committed
31
            (0.5, 0.5, 0) mode.
Wenwei Zhang's avatar
Wenwei Zhang committed
32
33

    Attributes:
34
35
36
        tensor (Tensor): Float matrix with shape (N, box_dim).
        box_dim (int): Integer indicating the dimension of a box. Each row is
            (x, y, z, x_size, y_size, z_size, yaw, ...).
Wenwei Zhang's avatar
Wenwei Zhang committed
37
38
        with_yaw (bool): If True, the value of yaw will be set to 0 as minmax
            boxes.
39
40
    """

41
42
    YAW_AXIS: int = 0

43
44
45
46
47
48
49
50
    def __init__(
        self,
        tensor: Union[Tensor, np.ndarray, Sequence[Sequence[float]]],
        box_dim: int = 7,
        with_yaw: bool = True,
        origin: Tuple[float, float, float] = (0.5, 0.5, 0)
    ) -> None:
        if isinstance(tensor, Tensor):
51
52
53
54
55
            device = tensor.device
        else:
            device = torch.device('cpu')
        tensor = torch.as_tensor(tensor, dtype=torch.float32, device=device)
        if tensor.numel() == 0:
56
57
58
59
60
61
62
            # Use reshape, so we don't end up creating a new tensor that does
            # not depend on the inputs (and consequently confuses jit)
            tensor = tensor.reshape((-1, box_dim))
        assert tensor.dim() == 2 and tensor.size(-1) == box_dim, \
            ('The box dimension must be 2 and the length of the last '
             f'dimension must be {box_dim}, but got boxes with shape '
             f'{tensor.shape}.')
wuyuefeng's avatar
wuyuefeng committed
63

wuyuefeng's avatar
wuyuefeng committed
64
        if tensor.shape[-1] == 6:
65
66
            # If the dimension of boxes is 6, we expand box_dim by padding 0 as
            # a fake yaw and set with_yaw to False
wuyuefeng's avatar
wuyuefeng committed
67
68
69
70
            assert box_dim == 6
            fake_rot = tensor.new_zeros(tensor.shape[0], 1)
            tensor = torch.cat((tensor, fake_rot), dim=-1)
            self.box_dim = box_dim + 1
wuyuefeng's avatar
wuyuefeng committed
71
            self.with_yaw = False
wuyuefeng's avatar
wuyuefeng committed
72
73
        else:
            self.box_dim = box_dim
wuyuefeng's avatar
wuyuefeng committed
74
            self.with_yaw = with_yaw
75
        self.tensor = tensor.clone()
76

wuyuefeng's avatar
wuyuefeng committed
77
78
        if origin != (0.5, 0.5, 0):
            dst = self.tensor.new_tensor((0.5, 0.5, 0))
zhangwenwei's avatar
zhangwenwei committed
79
80
81
            src = self.tensor.new_tensor(origin)
            self.tensor[:, :3] += self.tensor[:, 3:6] * (dst - src)

82
83
84
85
86
    @property
    def shape(self) -> torch.Size:
        """torch.Size: Shape of boxes."""
        return self.tensor.shape

zhangwenwei's avatar
zhangwenwei committed
87
    @property
88
89
    def volume(self) -> Tensor:
        """Tensor: A vector with volume of each box in shape (N, )."""
90
91
        return self.tensor[:, 3] * self.tensor[:, 4] * self.tensor[:, 5]

zhangwenwei's avatar
zhangwenwei committed
92
    @property
93
94
    def dims(self) -> Tensor:
        """Tensor: Size dimensions of each box in shape (N, 3)."""
zhangwenwei's avatar
zhangwenwei committed
95
96
        return self.tensor[:, 3:6]

zhangwenwei's avatar
zhangwenwei committed
97
    @property
98
99
    def yaw(self) -> Tensor:
        """Tensor: A vector with yaw of each box in shape (N, )."""
zhangwenwei's avatar
zhangwenwei committed
100
101
        return self.tensor[:, 6]

102
    @property
103
104
    def height(self) -> Tensor:
        """Tensor: A vector with height of each box in shape (N, )."""
105
106
        return self.tensor[:, 5]

107
    @property
108
109
    def top_height(self) -> Tensor:
        """Tensor: A vector with top height of each box in shape (N, )."""
110
111
112
        return self.bottom_height + self.height

    @property
113
114
    def bottom_height(self) -> Tensor:
        """Tensor: A vector with bottom height of each box in shape (N, )."""
115
116
        return self.tensor[:, 2]

zhangwenwei's avatar
zhangwenwei committed
117
    @property
118
    def center(self) -> Tensor:
zhangwenwei's avatar
zhangwenwei committed
119
120
121
        """Calculate the center of all the boxes.

        Note:
122
123
            In MMDetection3D's convention, the bottom center is usually taken
            as the default center.
zhangwenwei's avatar
zhangwenwei committed
124

125
126
127
128
129
            The relative position of the centers in different kinds of boxes
            are different, e.g., the relative center of a boxes is
            (0.5, 1.0, 0.5) in camera and (0.5, 0.5, 0) in lidar. It is
            recommended to use ``bottom_center`` or ``gravity_center`` for
            clearer usage.
zhangwenwei's avatar
zhangwenwei committed
130
131

        Returns:
132
            Tensor: A tensor with center of each box in shape (N, 3).
zhangwenwei's avatar
zhangwenwei committed
133
134
135
        """
        return self.bottom_center

zhangwenwei's avatar
zhangwenwei committed
136
    @property
137
138
    def bottom_center(self) -> Tensor:
        """Tensor: A tensor with center of each box in shape (N, 3)."""
zhangwenwei's avatar
zhangwenwei committed
139
        return self.tensor[:, :3]
140

zhangwenwei's avatar
zhangwenwei committed
141
    @property
142
143
144
145
146
147
148
    def gravity_center(self) -> Tensor:
        """Tensor: A tensor with center of each box in shape (N, 3)."""
        bottom_center = self.bottom_center
        gravity_center = torch.zeros_like(bottom_center)
        gravity_center[:, :2] = bottom_center[:, :2]
        gravity_center[:, 2] = bottom_center[:, 2] + self.tensor[:, 5] * 0.5
        return gravity_center
149

zhangwenwei's avatar
zhangwenwei committed
150
    @property
151
152
    def corners(self) -> Tensor:
        """Tensor: A tensor with 8 corners of each box in shape (N, 8, 3)."""
153
154
        pass

155
    @property
156
157
158
    def bev(self) -> Tensor:
        """Tensor: 2D BEV box of each box with rotation in XYWHR format, in
        shape (N, 5)."""
159
160
161
        return self.tensor[:, [0, 1, 3, 4, 6]]

    @property
162
163
    def nearest_bev(self) -> Tensor:
        """Tensor: A tensor of 2D BEV box of each box without rotation."""
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
        # Obtain BEV boxes with rotation in XYWHR format
        bev_rotated_boxes = self.bev
        # convert the rotation to a valid range
        rotations = bev_rotated_boxes[:, -1]
        normed_rotations = torch.abs(limit_period(rotations, 0.5, np.pi))

        # find the center of boxes
        conditions = (normed_rotations > np.pi / 4)[..., None]
        bboxes_xywh = torch.where(conditions, bev_rotated_boxes[:,
                                                                [0, 1, 3, 2]],
                                  bev_rotated_boxes[:, :4])

        centers = bboxes_xywh[:, :2]
        dims = bboxes_xywh[:, 2:]
        bev_boxes = torch.cat([centers - dims / 2, centers + dims / 2], dim=-1)
        return bev_boxes

181
182
183
    def in_range_bev(
            self, box_range: Union[Tensor, np.ndarray,
                                   Sequence[float]]) -> Tensor:
184
185
186
        """Check whether the boxes are in the given range.

        Args:
187
188
            box_range (Tensor or np.ndarray or Sequence[float]): The range of
                box in order of (x_min, y_min, x_max, y_max).
189
190

        Note:
191
192
193
            The original implementation of SECOND checks whether boxes in a
            range by checking whether the points are in a convex polygon, we
            reduce the burden for simpler cases.
194
195

        Returns:
196
197
            Tensor: A binary vector indicating whether each box is inside the
            reference range.
198
199
200
201
202
203
204
        """
        in_range_flags = ((self.bev[:, 0] > box_range[0])
                          & (self.bev[:, 1] > box_range[1])
                          & (self.bev[:, 0] < box_range[2])
                          & (self.bev[:, 1] < box_range[3]))
        return in_range_flags

205
    @abstractmethod
206
207
208
209
210
211
    def rotate(
        self,
        angle: Union[Tensor, np.ndarray, float],
        points: Optional[Union[Tensor, np.ndarray, BasePoints]] = None
    ) -> Union[Tuple[Tensor, Tensor], Tuple[np.ndarray, np.ndarray], Tuple[
            BasePoints, Tensor], None]:
212
213
        """Rotate boxes with points (optional) with the given angle or rotation
        matrix.
214
215

        Args:
216
217
218
            angle (Tensor or np.ndarray or float): Rotation angle or rotation
                matrix.
            points (Tensor or np.ndarray or :obj:`BasePoints`, optional):
219
                Points to rotate. Defaults to None.
220
221
222
223
224

        Returns:
            tuple or None: When ``points`` is None, the function returns None,
            otherwise it returns the rotated points and the rotation matrix
            ``rot_mat_T``.
225
226
227
228
        """
        pass

    @abstractmethod
229
230
231
232
233
    def flip(
        self,
        bev_direction: str = 'horizontal',
        points: Optional[Union[Tensor, np.ndarray, BasePoints]] = None
    ) -> Union[Tensor, np.ndarray, BasePoints, None]:
234
235
236
        """Flip the boxes in BEV along given BEV direction.

        Args:
237
238
239
240
241
242
243
244
245
            bev_direction (str): Direction by which to flip. Can be chosen from
                'horizontal' and 'vertical'. Defaults to 'horizontal'.
            points (Tensor or np.ndarray or :obj:`BasePoints`, optional):
                Points to flip. Defaults to None.

        Returns:
            Tensor or np.ndarray or :obj:`BasePoints` or None: When ``points``
            is None, the function returns None, otherwise it returns the
            flipped points.
246
        """
247
248
        pass

249
    def translate(self, trans_vector: Union[Tensor, np.ndarray]) -> None:
250
        """Translate boxes with the given translation vector.
251
252

        Args:
253
254
            trans_vector (Tensor or np.ndarray): Translation vector of size
                1x3.
255
        """
256
        if not isinstance(trans_vector, Tensor):
zhangwenwei's avatar
zhangwenwei committed
257
258
            trans_vector = self.tensor.new_tensor(trans_vector)
        self.tensor[:, :3] += trans_vector
259

260
261
262
    def in_range_3d(
            self, box_range: Union[Tensor, np.ndarray,
                                   Sequence[float]]) -> Tensor:
zhangwenwei's avatar
zhangwenwei committed
263
        """Check whether the boxes are in the given range.
264
265

        Args:
266
267
            box_range (Tensor or np.ndarray or Sequence[float]): The range of
                box (x_min, y_min, z_min, x_max, y_max, z_max).
268

zhangwenwei's avatar
zhangwenwei committed
269
        Note:
270
271
272
            In the original implementation of SECOND, checking whether a box in
            the range checks whether the points are in a convex polygon, we try
            to reduce the burden for simpler cases.
zhangwenwei's avatar
zhangwenwei committed
273

274
        Returns:
275
276
            Tensor: A binary vector indicating whether each point is inside the
            reference range.
277
        """
278
279
280
281
282
283
284
        gravity_center = self.gravity_center
        in_range_flags = ((gravity_center[:, 0] > box_range[0])
                          & (gravity_center[:, 1] > box_range[1])
                          & (gravity_center[:, 2] > box_range[2])
                          & (gravity_center[:, 0] < box_range[3])
                          & (gravity_center[:, 1] < box_range[4])
                          & (gravity_center[:, 2] < box_range[5]))
zhangwenwei's avatar
zhangwenwei committed
285
        return in_range_flags
286

287
    @abstractmethod
288
289
290
291
    def convert_to(self,
                   dst: int,
                   rt_mat: Optional[Union[Tensor, np.ndarray]] = None,
                   correct_yaw: bool = False) -> 'BaseInstance3DBoxes':
Wenwei Zhang's avatar
Wenwei Zhang committed
292
        """Convert self to ``dst`` mode.
293
294

        Args:
295
296
            dst (int): The target Box mode.
            rt_mat (Tensor or np.ndarray, optional): The rotation and
297
                translation matrix between different coordinates.
298
299
300
301
302
303
                Defaults to None. The conversion from ``src`` coordinates to
                ``dst`` coordinates usually comes along the change of sensors,
                e.g., from camera to LiDAR. This requires a transformation
                matrix.
            correct_yaw (bool): Whether to convert the yaw angle to the target
                coordinate. Defaults to False.
304
305

        Returns:
306
307
            :obj:`BaseInstance3DBoxes`: The converted box of the same type in
            the ``dst`` mode.
308
309
310
        """
        pass

311
    def scale(self, scale_factor: float) -> None:
zhangwenwei's avatar
zhangwenwei committed
312
        """Scale the box with horizontal and vertical scaling factors.
zhangwenwei's avatar
zhangwenwei committed
313
314

        Args:
liyinhao's avatar
liyinhao committed
315
            scale_factors (float): Scale factors to scale the boxes.
zhangwenwei's avatar
zhangwenwei committed
316
        """
zhangwenwei's avatar
zhangwenwei committed
317
        self.tensor[:, :6] *= scale_factor
318
        self.tensor[:, 7:] *= scale_factor  # velocity
zhangwenwei's avatar
zhangwenwei committed
319

320
    def limit_yaw(self, offset: float = 0.5, period: float = np.pi) -> None:
zhangwenwei's avatar
zhangwenwei committed
321
        """Limit the yaw to a given period and offset.
zhangwenwei's avatar
zhangwenwei committed
322
323

        Args:
324
325
            offset (float): The offset of the yaw. Defaults to 0.5.
            period (float): The expected period. Defaults to np.pi.
zhangwenwei's avatar
zhangwenwei committed
326
327
        """
        self.tensor[:, 6] = limit_period(self.tensor[:, 6], offset, period)
zhangwenwei's avatar
zhangwenwei committed
328

329
    def nonempty(self, threshold: float = 0.0) -> Tensor:
330
331
        """Find boxes that are non-empty.

332
333
        A box is considered empty if either of its side is no larger than
        threshold.
334

zhangwenwei's avatar
zhangwenwei committed
335
        Args:
336
            threshold (float): The threshold of minimal sizes. Defaults to 0.0.
zhangwenwei's avatar
zhangwenwei committed
337

338
        Returns:
339
340
            Tensor: A binary vector which represents whether each box is empty
            (False) or non-empty (True).
341
342
343
344
345
346
347
348
349
        """
        box = self.tensor
        size_x = box[..., 3]
        size_y = box[..., 4]
        size_z = box[..., 5]
        keep = ((size_x > threshold)
                & (size_y > threshold) & (size_z > threshold))
        return keep

350
351
352
    def __getitem__(
            self, item: Union[int, slice, np.ndarray,
                              Tensor]) -> 'BaseInstance3DBoxes':
353
        """
354
355
356
        Args:
            item (int or slice or np.ndarray or Tensor): Index of boxes.

357
358
        Note:
            The following usage are allowed:
359
360
361
362
363
364
365
366

            1. `new_boxes = boxes[3]`: Return a `Boxes` that contains only one
               box.
            2. `new_boxes = boxes[2:10]`: Return a slice of boxes.
            3. `new_boxes = boxes[vector]`: Where vector is a
               torch.BoolTensor with `length = len(boxes)`. Nonzero elements in
               the vector will be selected.

367
            Note that the returned Boxes might share storage with this Boxes,
368
            subject to PyTorch's indexing semantics.
369
370

        Returns:
371
            :obj:`BaseInstance3DBoxes`: A new object of
372
            :class:`BaseInstance3DBoxes` after indexing.
373
374
375
        """
        original_type = type(self)
        if isinstance(item, int):
wuyuefeng's avatar
wuyuefeng committed
376
377
378
379
            return original_type(
                self.tensor[item].view(1, -1),
                box_dim=self.box_dim,
                with_yaw=self.with_yaw)
380
381
382
        b = self.tensor[item]
        assert b.dim() == 2, \
            f'Indexing on Boxes with {item} failed to return a matrix!'
wuyuefeng's avatar
wuyuefeng committed
383
        return original_type(b, box_dim=self.box_dim, with_yaw=self.with_yaw)
384

385
    def __len__(self) -> int:
wangtai's avatar
wangtai committed
386
        """int: Number of boxes in the current object."""
387
388
        return self.tensor.shape[0]

389
390
    def __repr__(self) -> str:
        """str: Return a string that describes the object."""
391
392
393
        return self.__class__.__name__ + '(\n    ' + str(self.tensor) + ')'

    @classmethod
394
395
    def cat(cls, boxes_list: Sequence['BaseInstance3DBoxes']
            ) -> 'BaseInstance3DBoxes':
396
        """Concatenate a list of Boxes into a single Boxes.
397

liyinhao's avatar
liyinhao committed
398
        Args:
399
            boxes_list (Sequence[:obj:`BaseInstance3DBoxes`]): List of boxes.
zhangwenwei's avatar
zhangwenwei committed
400

401
        Returns:
402
            :obj:`BaseInstance3DBoxes`: The concatenated boxes.
403
404
405
406
407
408
409
410
        """
        assert isinstance(boxes_list, (list, tuple))
        if len(boxes_list) == 0:
            return cls(torch.empty(0))
        assert all(isinstance(box, cls) for box in boxes_list)

        # use torch.cat (v.s. layers.cat)
        # so the returned boxes never share storage with input
zhangwenwei's avatar
zhangwenwei committed
411
412
        cat_boxes = cls(
            torch.cat([b.tensor for b in boxes_list], dim=0),
413
            box_dim=boxes_list[0].box_dim,
zhangwenwei's avatar
zhangwenwei committed
414
            with_yaw=boxes_list[0].with_yaw)
415
416
        return cat_boxes

417
418
419
420
    def numpy(self) -> np.ndarray:
        """Reload ``numpy`` from self.tensor."""
        return self.tensor.numpy()

421
422
    def to(self, device: Union[str, torch.device], *args,
           **kwargs) -> 'BaseInstance3DBoxes':
zhangwenwei's avatar
zhangwenwei committed
423
        """Convert current boxes to a specific device.
zhangwenwei's avatar
zhangwenwei committed
424
425

        Args:
426
            device (str or :obj:`torch.device`): The name of the device.
zhangwenwei's avatar
zhangwenwei committed
427
428

        Returns:
429
            :obj:`BaseInstance3DBoxes`: A new boxes object on the specific
430
            device.
zhangwenwei's avatar
zhangwenwei committed
431
        """
432
        original_type = type(self)
wuyuefeng's avatar
wuyuefeng committed
433
        return original_type(
JingweiZhang12's avatar
JingweiZhang12 committed
434
            self.tensor.to(device, *args, **kwargs),
wuyuefeng's avatar
wuyuefeng committed
435
436
            box_dim=self.box_dim,
            with_yaw=self.with_yaw)
437

438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
    def cpu(self) -> 'BaseInstance3DBoxes':
        """Convert current boxes to cpu device.

        Returns:
            :obj:`BaseInstance3DBoxes`: A new boxes object on the cpu device.
        """
        original_type = type(self)
        return original_type(
            self.tensor.cpu(), box_dim=self.box_dim, with_yaw=self.with_yaw)

    def cuda(self, *args, **kwargs) -> 'BaseInstance3DBoxes':
        """Convert current boxes to cuda device.

        Returns:
            :obj:`BaseInstance3DBoxes`: A new boxes object on the cuda device.
        """
        original_type = type(self)
        return original_type(
            self.tensor.cuda(*args, **kwargs),
            box_dim=self.box_dim,
            with_yaw=self.with_yaw)

460
461
    def clone(self) -> 'BaseInstance3DBoxes':
        """Clone the boxes.
462
463

        Returns:
464
465
            :obj:`BaseInstance3DBoxes`: Box object with the same properties as
            self.
466
467
        """
        original_type = type(self)
wuyuefeng's avatar
wuyuefeng committed
468
469
        return original_type(
            self.tensor.clone(), box_dim=self.box_dim, with_yaw=self.with_yaw)
470

471
472
473
474
475
476
477
478
479
480
481
    def detach(self) -> 'BaseInstance3DBoxes':
        """Detach the boxes.

        Returns:
            :obj:`BaseInstance3DBoxes`: Box object with the same properties as
            self.
        """
        original_type = type(self)
        return original_type(
            self.tensor.detach(), box_dim=self.box_dim, with_yaw=self.with_yaw)

482
    @property
483
484
    def device(self) -> torch.device:
        """torch.device: The device of the boxes are on."""
485
486
        return self.tensor.device

487
488
    def __iter__(self) -> Iterator[Tensor]:
        """Yield a box as a Tensor at a time.
wuyuefeng's avatar
wuyuefeng committed
489
490

        Returns:
491
            Iterator[Tensor]: A box of shape (box_dim, ).
492
493
        """
        yield from self.tensor
494
495

    @classmethod
496
497
    def height_overlaps(cls, boxes1: 'BaseInstance3DBoxes',
                        boxes2: 'BaseInstance3DBoxes') -> Tensor:
zhangwenwei's avatar
zhangwenwei committed
498
        """Calculate height overlaps of two boxes.
499
500

        Note:
501
502
            This function calculates the height overlaps between ``boxes1`` and
            ``boxes2``, ``boxes1`` and ``boxes2`` should be in the same type.
503
504

        Args:
505
506
            boxes1 (:obj:`BaseInstance3DBoxes`): Boxes 1 contain N boxes.
            boxes2 (:obj:`BaseInstance3DBoxes`): Boxes 2 contain M boxes.
507
508

        Returns:
509
            Tensor: Calculated height overlap of the boxes.
510
        """
511
512
        assert isinstance(boxes1, BaseInstance3DBoxes)
        assert isinstance(boxes2, BaseInstance3DBoxes)
513
514
515
        assert type(boxes1) == type(boxes2), \
            '"boxes1" and "boxes2" should be in the same type, ' \
            f'but got {type(boxes1)} and {type(boxes2)}.'
516
517
518
519
520
521
522
523
524
525
526
527
528

        boxes1_top_height = boxes1.top_height.view(-1, 1)
        boxes1_bottom_height = boxes1.bottom_height.view(-1, 1)
        boxes2_top_height = boxes2.top_height.view(1, -1)
        boxes2_bottom_height = boxes2.bottom_height.view(1, -1)

        heighest_of_bottom = torch.max(boxes1_bottom_height,
                                       boxes2_bottom_height)
        lowest_of_top = torch.min(boxes1_top_height, boxes2_top_height)
        overlaps_h = torch.clamp(lowest_of_top - heighest_of_bottom, min=0)
        return overlaps_h

    @classmethod
529
530
531
532
    def overlaps(cls,
                 boxes1: 'BaseInstance3DBoxes',
                 boxes2: 'BaseInstance3DBoxes',
                 mode: str = 'iou') -> Tensor:
zhangwenwei's avatar
zhangwenwei committed
533
        """Calculate 3D overlaps of two boxes.
534
535

        Note:
Wenwei Zhang's avatar
Wenwei Zhang committed
536
537
            This function calculates the overlaps between ``boxes1`` and
            ``boxes2``, ``boxes1`` and ``boxes2`` should be in the same type.
538
539

        Args:
540
541
            boxes1 (:obj:`BaseInstance3DBoxes`): Boxes 1 contain N boxes.
            boxes2 (:obj:`BaseInstance3DBoxes`): Boxes 2 contain M boxes.
542
            mode (str): Mode of iou calculation. Defaults to 'iou'.
543
544

        Returns:
545
            Tensor: Calculated 3D overlap of the boxes.
546
547
548
        """
        assert isinstance(boxes1, BaseInstance3DBoxes)
        assert isinstance(boxes2, BaseInstance3DBoxes)
549
550
551
        assert type(boxes1) == type(boxes2), \
            '"boxes1" and "boxes2" should be in the same type, ' \
            f'but got {type(boxes1)} and {type(boxes2)}.'
552
553
554

        assert mode in ['iou', 'iof']

zhangwenwei's avatar
zhangwenwei committed
555
556
557
558
559
        rows = len(boxes1)
        cols = len(boxes2)
        if rows * cols == 0:
            return boxes1.tensor.new(rows, cols)

560
561
562
        # height overlap
        overlaps_h = cls.height_overlaps(boxes1, boxes2)

563
564
565
566
        # Restrict the min values of W and H to avoid memory overflow in
        # ``box_iou_rotated``.
        boxes1_bev, boxes2_bev = boxes1.bev, boxes2.bev
        boxes1_bev[:, 2:4] = boxes1_bev[:, 2:4].clamp(min=1e-4)
567
        boxes2_bev[:, 2:4] = boxes2_bev[:, 2:4].clamp(min=1e-4)
568

569
        # bev overlap
570
571
        iou2d = box_iou_rotated(boxes1_bev, boxes2_bev)
        areas1 = (boxes1_bev[:, 2] * boxes1_bev[:, 3]).unsqueeze(1).expand(
572
            rows, cols)
573
        areas2 = (boxes2_bev[:, 2] * boxes2_bev[:, 3]).unsqueeze(0).expand(
574
575
            rows, cols)
        overlaps_bev = iou2d * (areas1 + areas2) / (1 + iou2d)
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590

        # 3d overlaps
        overlaps_3d = overlaps_bev.to(boxes1.device) * overlaps_h

        volume1 = boxes1.volume.view(-1, 1)
        volume2 = boxes2.volume.view(1, -1)

        if mode == 'iou':
            # the clamp func is used to avoid division of 0
            iou3d = overlaps_3d / torch.clamp(
                volume1 + volume2 - overlaps_3d, min=1e-8)
        else:
            iou3d = overlaps_3d / torch.clamp(volume1, min=1e-8)

        return iou3d
wuyuefeng's avatar
wuyuefeng committed
591

592
593
594
    def new_box(
        self, data: Union[Tensor, np.ndarray, Sequence[Sequence[float]]]
    ) -> 'BaseInstance3DBoxes':
wuyuefeng's avatar
wuyuefeng committed
595
596
        """Create a new box object with data.

597
598
        The new box and its tensor has the similar properties as self and
        self.tensor, respectively.
wuyuefeng's avatar
wuyuefeng committed
599
600

        Args:
601
602
            data (Tensor or np.ndarray or Sequence[Sequence[float]]): Data to
                be copied.
wuyuefeng's avatar
wuyuefeng committed
603
604

        Returns:
605
606
            :obj:`BaseInstance3DBoxes`: A new bbox object with ``data``, the
            object's other properties are similar to ``self``.
wuyuefeng's avatar
wuyuefeng committed
607
        """
zhangwenwei's avatar
zhangwenwei committed
608
        new_tensor = self.tensor.new_tensor(data) \
609
            if not isinstance(data, Tensor) else data.to(self.device)
wuyuefeng's avatar
wuyuefeng committed
610
611
612
        original_type = type(self)
        return original_type(
            new_tensor, box_dim=self.box_dim, with_yaw=self.with_yaw)
613

614
615
616
617
    def points_in_boxes_part(
            self,
            points: Tensor,
            boxes_override: Optional[Tensor] = None) -> Tensor:
618
        """Find the box in which each point is.
619
620

        Args:
621
622
623
624
            points (Tensor): Points in shape (1, M, 3) or (M, 3), 3 dimensions
                are (x, y, z) in LiDAR or depth coordinate.
            boxes_override (Tensor, optional): Boxes to override `self.tensor`.
                Defaults to None.
625
626

        Note:
627
628
629
630
631
632
633
            If a point is enclosed by multiple boxes, the index of the first
            box will be returned.

        Returns:
            Tensor: The index of the first box that each point is in with shape
            (M, ). Default value is -1 (if the point is not enclosed by any
            box).
634
635
636
637
638
        """
        if boxes_override is not None:
            boxes = boxes_override
        else:
            boxes = self.tensor
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653

        points_clone = points.clone()[..., :3]
        if points_clone.dim() == 2:
            points_clone = points_clone.unsqueeze(0)
        else:
            assert points_clone.dim() == 3 and points_clone.shape[0] == 1

        boxes = boxes.to(points_clone.device).unsqueeze(0)
        box_idx = points_in_boxes_part(points_clone, boxes)

        return box_idx.squeeze(0)

    def points_in_boxes_all(self,
                            points: Tensor,
                            boxes_override: Optional[Tensor] = None) -> Tensor:
654
        """Find all boxes in which each point is.
655
656

        Args:
657
658
659
660
            points (Tensor): Points in shape (1, M, 3) or (M, 3), 3 dimensions
                are (x, y, z) in LiDAR or depth coordinate.
            boxes_override (Tensor, optional): Boxes to override `self.tensor`.
                Defaults to None.
661
662

        Returns:
663
664
665
666
            Tensor: A tensor indicating whether a point is in a box with shape
            (M, T). T is the number of boxes. Denote this tensor as A, it the
            m^th point is in the t^th box, then `A[m, t] == 1`, otherwise
            `A[m, t] == 0`.
667
668
669
670
671
672
673
674
675
676
677
678
679
        """
        if boxes_override is not None:
            boxes = boxes_override
        else:
            boxes = self.tensor

        points_clone = points.clone()[..., :3]
        if points_clone.dim() == 2:
            points_clone = points_clone.unsqueeze(0)
        else:
            assert points_clone.dim() == 3 and points_clone.shape[0] == 1

        boxes = boxes.to(points_clone.device).unsqueeze(0)
680
        box_idxs_of_pts = points_in_boxes_all(points_clone, boxes)
681
682

        return box_idxs_of_pts.squeeze(0)
683

684
685
686
687
688
    def points_in_boxes(self,
                        points: Tensor,
                        boxes_override: Optional[Tensor] = None) -> Tensor:
        warnings.warn('DeprecationWarning: points_in_boxes is a deprecated '
                      'method, please consider using points_in_boxes_part.')
689
690
        return self.points_in_boxes_part(points, boxes_override)

691
692
693
694
    def points_in_boxes_batch(
            self,
            points: Tensor,
            boxes_override: Optional[Tensor] = None) -> Tensor:
695
696
697
698
        warnings.warn('DeprecationWarning: points_in_boxes_batch is a '
                      'deprecated method, please consider using '
                      'points_in_boxes_all.')
        return self.points_in_boxes_all(points, boxes_override)