lidar_box3d.py 7.93 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
2
3
4
import numpy as np
import torch

zhangshilong's avatar
zhangshilong committed
5
from mmdet3d.structures.points import BasePoints
6
from .base_box3d import BaseInstance3DBoxes
7
from .utils import rotation_3d_in_axis
8
9
10


class LiDARInstance3DBoxes(BaseInstance3DBoxes):
zhangwenwei's avatar
zhangwenwei committed
11
    """3D boxes of instances in LIDAR coordinates.
zhangwenwei's avatar
zhangwenwei committed
12

zhangwenwei's avatar
zhangwenwei committed
13
    Coordinates in LiDAR:
14

zhangwenwei's avatar
zhangwenwei committed
15
    .. code-block:: none
zhangwenwei's avatar
zhangwenwei committed
16

17
18
19
20
21
                                up z    x front (yaw=0)
                                   ^   ^
                                   |  /
                                   | /
       (yaw=0.5*pi) left y <------ 0
zhangwenwei's avatar
zhangwenwei committed
22

wuyuefeng's avatar
wuyuefeng committed
23
    The relative coordinate of bottom center in a LiDAR box is (0.5, 0.5, 0),
zhangwenwei's avatar
zhangwenwei committed
24
    and the yaw is around the z axis, thus the rotation axis=2.
25
26
    The yaw is 0 at the positive direction of x axis, and increases from
    the positive direction of x to the positive direction of y.
27

28
29
30
    A refactor is ongoing to make the three coordinate systems
    easier to understand and convert between each other.

31
    Attributes:
liyinhao's avatar
liyinhao committed
32
        tensor (torch.Tensor): Float matrix of N x box_dim.
33
        box_dim (int): Integer indicating the dimension of a box.
zhangwenwei's avatar
zhangwenwei committed
34
            Each row is (x, y, z, x_size, y_size, z_size, yaw, ...).
liyinhao's avatar
liyinhao committed
35
        with_yaw (bool): If True, the value of yaw will be set to 0 as minmax
zhangwenwei's avatar
zhangwenwei committed
36
            boxes.
37
    """
38
    YAW_AXIS = 2
39

zhangwenwei's avatar
zhangwenwei committed
40
    @property
41
    def gravity_center(self):
42
        """torch.Tensor: A tensor with center of each box in shape (N, 3)."""
zhangwenwei's avatar
zhangwenwei committed
43
        bottom_center = self.bottom_center
44
45
        gravity_center = torch.zeros_like(bottom_center)
        gravity_center[:, :2] = bottom_center[:, :2]
zhangwenwei's avatar
zhangwenwei committed
46
        gravity_center[:, 2] = bottom_center[:, 2] + self.tensor[:, 5] * 0.5
47
48
        return gravity_center

zhangwenwei's avatar
zhangwenwei committed
49
50
    @property
    def corners(self):
Wenwei Zhang's avatar
Wenwei Zhang committed
51
52
        """torch.Tensor: Coordinates of corners of all the boxes
        in shape (N, 8, 3).
53

zhangwenwei's avatar
zhangwenwei committed
54
        Convert the boxes to corners in clockwise order, in form of
Wenwei Zhang's avatar
Wenwei Zhang committed
55
        ``(x0y0z0, x0y0z1, x0y1z1, x0y1z0, x1y0z0, x1y0z1, x1y1z1, x1y1z0)``
zhangwenwei's avatar
zhangwenwei committed
56
57

        .. code-block:: none
zhangwenwei's avatar
zhangwenwei committed
58

zhangwenwei's avatar
zhangwenwei committed
59
60
61
62
63
64
65
66
67
                                           up z
                            front x           ^
                                 /            |
                                /             |
                  (x1, y0, z1) + -----------  + (x1, y1, z1)
                              /|            / |
                             / |           /  |
               (x0, y0, z1) + ----------- +   + (x1, y1, z0)
                            |  /      .   |  /
68
                            | / origin    | /
zhangwenwei's avatar
zhangwenwei committed
69
70
            left y<-------- + ----------- + (x0, y1, z0)
                (x0, y0, z0)
71
        """
72
73
74
        if self.tensor.numel() == 0:
            return torch.empty([0, 8, 3], device=self.tensor.device)

zhangwenwei's avatar
zhangwenwei committed
75
        dims = self.dims
76
        corners_norm = torch.from_numpy(
zhangwenwei's avatar
zhangwenwei committed
77
            np.stack(np.unravel_index(np.arange(8), [2] * 3), axis=1)).to(
78
79
80
                device=dims.device, dtype=dims.dtype)

        corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]]
zhangwenwei's avatar
zhangwenwei committed
81
        # use relative origin [0.5, 0.5, 0]
zhangwenwei's avatar
zhangwenwei committed
82
83
        corners_norm = corners_norm - dims.new_tensor([0.5, 0.5, 0])
        corners = dims.view([-1, 1, 3]) * corners_norm.reshape([1, 8, 3])
84

zhangwenwei's avatar
zhangwenwei committed
85
        # rotate around z axis
86
87
        corners = rotation_3d_in_axis(
            corners, self.tensor[:, 6], axis=self.YAW_AXIS)
88
89
90
        corners += self.tensor[:, :3].view(-1, 1, 3)
        return corners

wuyuefeng's avatar
wuyuefeng committed
91
    def rotate(self, angle, points=None):
92
93
        """Rotate boxes with points (optional) with the given angle or rotation
        matrix.
94
95

        Args:
96
97
            angles (float | torch.Tensor | np.ndarray):
                Rotation angle or rotation matrix.
98
            points (torch.Tensor | np.ndarray | :obj:`BasePoints`, optional):
99
                Points to rotate. Defaults to None.
wuyuefeng's avatar
wuyuefeng committed
100
101

        Returns:
102
103
            tuple or None: When ``points`` is None, the function returns
                None, otherwise it returns the rotated points and the
wuyuefeng's avatar
wuyuefeng committed
104
                rotation matrix ``rot_mat_T``.
105
106
107
        """
        if not isinstance(angle, torch.Tensor):
            angle = self.tensor.new_tensor(angle)
108

109
110
111
112
        assert angle.shape == torch.Size([3, 3]) or angle.numel() == 1, \
            f'invalid rotation angle shape {angle.shape}'

        if angle.numel() == 1:
113
114
115
116
117
            self.tensor[:, 0:3], rot_mat_T = rotation_3d_in_axis(
                self.tensor[:, 0:3],
                angle,
                axis=self.YAW_AXIS,
                return_mat=True)
118
119
        else:
            rot_mat_T = angle
120
            rot_sin = rot_mat_T[0, 1]
121
122
            rot_cos = rot_mat_T[0, 0]
            angle = np.arctan2(rot_sin, rot_cos)
123
            self.tensor[:, 0:3] = self.tensor[:, 0:3] @ rot_mat_T
124
125
126

        self.tensor[:, 6] += angle

127
128
129
130
        if self.tensor.shape[1] == 9:
            # rotate velo vector
            self.tensor[:, 7:9] = self.tensor[:, 7:9] @ rot_mat_T[:2, :2]

wuyuefeng's avatar
wuyuefeng committed
131
132
133
134
        if points is not None:
            if isinstance(points, torch.Tensor):
                points[:, :3] = points[:, :3] @ rot_mat_T
            elif isinstance(points, np.ndarray):
135
                rot_mat_T = rot_mat_T.cpu().numpy()
wuyuefeng's avatar
wuyuefeng committed
136
                points[:, :3] = np.dot(points[:, :3], rot_mat_T)
137
            elif isinstance(points, BasePoints):
138
                points.rotate(rot_mat_T)
wuyuefeng's avatar
wuyuefeng committed
139
140
141
142
143
            else:
                raise ValueError
            return points, rot_mat_T

    def flip(self, bev_direction='horizontal', points=None):
zhangwenwei's avatar
zhangwenwei committed
144
        """Flip the boxes in BEV along given BEV direction.
145

wuyuefeng's avatar
wuyuefeng committed
146
147
148
149
        In LIDAR coordinates, it flips the y (horizontal) or x (vertical) axis.

        Args:
            bev_direction (str): Flip direction (horizontal or vertical).
150
            points (torch.Tensor | np.ndarray | :obj:`BasePoints`, optional):
151
                Points to flip. Defaults to None.
wuyuefeng's avatar
wuyuefeng committed
152
153
154

        Returns:
            torch.Tensor, numpy.ndarray or None: Flipped points.
155
        """
wuyuefeng's avatar
wuyuefeng committed
156
157
158
159
        assert bev_direction in ('horizontal', 'vertical')
        if bev_direction == 'horizontal':
            self.tensor[:, 1::7] = -self.tensor[:, 1::7]
            if self.with_yaw:
160
                self.tensor[:, 6] = -self.tensor[:, 6]
wuyuefeng's avatar
wuyuefeng committed
161
162
163
        elif bev_direction == 'vertical':
            self.tensor[:, 0::7] = -self.tensor[:, 0::7]
            if self.with_yaw:
164
                self.tensor[:, 6] = -self.tensor[:, 6] + np.pi
165

wuyuefeng's avatar
wuyuefeng committed
166
        if points is not None:
167
168
169
170
171
172
173
174
            assert isinstance(points, (torch.Tensor, np.ndarray, BasePoints))
            if isinstance(points, (torch.Tensor, np.ndarray)):
                if bev_direction == 'horizontal':
                    points[:, 1] = -points[:, 1]
                elif bev_direction == 'vertical':
                    points[:, 0] = -points[:, 0]
            elif isinstance(points, BasePoints):
                points.flip(bev_direction)
wuyuefeng's avatar
wuyuefeng committed
175
176
            return points

177
    def convert_to(self, dst, rt_mat=None):
Wenwei Zhang's avatar
Wenwei Zhang committed
178
        """Convert self to ``dst`` mode.
179
180

        Args:
181
            dst (:obj:`Box3DMode`): the target Box mode
182
183
184
            rt_mat (np.ndarray | torch.Tensor, optional): The rotation and
                translation matrix between different coordinates.
                Defaults to None.
Wenwei Zhang's avatar
Wenwei Zhang committed
185
                The conversion from ``src`` coordinates to ``dst`` coordinates
186
187
188
189
                usually comes along the change of sensors, e.g., from camera
                to LiDAR. This requires a transformation matrix.

        Returns:
190
            :obj:`BaseInstance3DBoxes`:
Wenwei Zhang's avatar
Wenwei Zhang committed
191
                The converted box of the same type in the ``dst`` mode.
192
193
194
195
        """
        from .box_3d_mode import Box3DMode
        return Box3DMode.convert(
            box=self, src=Box3DMode.LIDAR, dst=dst, rt_mat=rt_mat)
zhangwenwei's avatar
zhangwenwei committed
196
197

    def enlarged_box(self, extra_width):
zhangwenwei's avatar
zhangwenwei committed
198
        """Enlarge the length, width and height boxes.
zhangwenwei's avatar
zhangwenwei committed
199
200

        Args:
201
            extra_width (float | torch.Tensor): Extra width to enlarge the box.
zhangwenwei's avatar
zhangwenwei committed
202
203

        Returns:
zhangwenwei's avatar
zhangwenwei committed
204
            :obj:`LiDARInstance3DBoxes`: Enlarged boxes.
zhangwenwei's avatar
zhangwenwei committed
205
206
207
208
209
210
        """
        enlarged_boxes = self.tensor.clone()
        enlarged_boxes[:, 3:6] += extra_width * 2
        # bottom center z minus extra_width
        enlarged_boxes[:, 2] -= extra_width
        return self.new_box(enlarged_boxes)