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

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


class DepthInstance3DBoxes(BaseInstance3DBoxes):
liyinhao's avatar
liyinhao committed
11
    """3D boxes of instances in Depth coordinates.
wuyuefeng's avatar
wuyuefeng committed
12
13

    Coordinates in Depth:
14

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

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

wuyuefeng's avatar
wuyuefeng committed
23
    The relative coordinate of bottom center in a Depth box is (0.5, 0.5, 0),
wuyuefeng's avatar
wuyuefeng committed
24
    and the yaw is around the z axis, thus the rotation axis=2.
25
    The yaw is 0 at the positive direction of x axis, and decreases from
wuyuefeng's avatar
wuyuefeng committed
26
    the positive direction of x to the positive direction of y.
27
28
29
30
31
    Also note that rotation of DepthInstance3DBoxes is counterclockwise,
    which is reverse to the definition of the yaw angle (clockwise).

    A refactor is ongoing to make the three coordinate systems
    easier to understand and convert between each other.
wuyuefeng's avatar
wuyuefeng committed
32

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

    @property
    def gravity_center(self):
44
        """torch.Tensor: A tensor with center of each box in shape (N, 3)."""
wuyuefeng's avatar
wuyuefeng committed
45
46
47
48
49
50
51
52
        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

    @property
    def corners(self):
Wenwei Zhang's avatar
Wenwei Zhang committed
53
54
        """torch.Tensor: Coordinates of corners of all the boxes
        in shape (N, 8, 3).
wuyuefeng's avatar
wuyuefeng committed
55
56

        Convert the boxes to corners in clockwise order, in form of
Wenwei Zhang's avatar
Wenwei Zhang committed
57
        ``(x0y0z0, x0y0z1, x0y1z1, x0y1z0, x1y0z0, x1y0z1, x1y1z1, x1y1z0)``
wuyuefeng's avatar
wuyuefeng committed
58
59
60
61
62
63
64
65
66
67
68
69

        .. code-block:: none

                                           up z
                            front y           ^
                                 /            |
                                /             |
                  (x0, y1, z1) + -----------  + (x1, y1, z1)
                              /|            / |
                             / |           /  |
               (x0, y0, z1) + ----------- +   + (x1, y1, z0)
                            |  /      .   |  /
70
                            | / origin    | /
wuyuefeng's avatar
wuyuefeng committed
71
72
73
               (x0, y0, z0) + ----------- + --------> right x
                                          (x1, y0, z0)
        """
74
75
76
        if self.tensor.numel() == 0:
            return torch.empty([0, 8, 3], device=self.tensor.device)

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

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

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

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

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

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

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

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

wuyuefeng's avatar
wuyuefeng committed
127
        if self.with_yaw:
128
            self.tensor[:, 6] += angle
wuyuefeng's avatar
wuyuefeng committed
129
        else:
130
131
            # for axis-aligned boxes, we take the new
            # enclosing axis-aligned boxes after rotation
wuyuefeng's avatar
wuyuefeng committed
132
            corners_rot = self.corners @ rot_mat_T
wuyuefeng's avatar
wuyuefeng committed
133
134
135
136
137
138
139
140
            new_x_size = corners_rot[..., 0].max(
                dim=1, keepdim=True)[0] - corners_rot[..., 0].min(
                    dim=1, keepdim=True)[0]
            new_y_size = corners_rot[..., 1].max(
                dim=1, keepdim=True)[0] - corners_rot[..., 1].min(
                    dim=1, keepdim=True)[0]
            self.tensor[:, 3:5] = torch.cat((new_x_size, new_y_size), dim=-1)

wuyuefeng's avatar
wuyuefeng committed
141
142
143
144
        if points is not None:
            if isinstance(points, torch.Tensor):
                points[:, :3] = points[:, :3] @ rot_mat_T
            elif isinstance(points, np.ndarray):
145
                rot_mat_T = rot_mat_T.cpu().numpy()
wuyuefeng's avatar
wuyuefeng committed
146
                points[:, :3] = np.dot(points[:, :3], rot_mat_T)
147
            elif isinstance(points, BasePoints):
148
                points.rotate(rot_mat_T)
wuyuefeng's avatar
wuyuefeng committed
149
150
151
152
153
            else:
                raise ValueError
            return points, rot_mat_T

    def flip(self, bev_direction='horizontal', points=None):
zhangwenwei's avatar
zhangwenwei committed
154
        """Flip the boxes in BEV along given BEV direction.
wuyuefeng's avatar
wuyuefeng committed
155
156
157
158

        In Depth coordinates, it flips x (horizontal) or y (vertical) axis.

        Args:
159
160
161
            bev_direction (str, optional): Flip direction
                (horizontal or vertical). Defaults to 'horizontal'.
            points (torch.Tensor | np.ndarray | :obj:`BasePoints`, optional):
162
                Points to flip. Defaults to None.
wuyuefeng's avatar
wuyuefeng committed
163
164
165

        Returns:
            torch.Tensor, numpy.ndarray or None: Flipped points.
wuyuefeng's avatar
wuyuefeng committed
166
167
168
169
170
171
172
173
174
175
176
        """
        assert bev_direction in ('horizontal', 'vertical')
        if bev_direction == 'horizontal':
            self.tensor[:, 0::7] = -self.tensor[:, 0::7]
            if self.with_yaw:
                self.tensor[:, 6] = -self.tensor[:, 6] + np.pi
        elif bev_direction == 'vertical':
            self.tensor[:, 1::7] = -self.tensor[:, 1::7]
            if self.with_yaw:
                self.tensor[:, 6] = -self.tensor[:, 6]

wuyuefeng's avatar
wuyuefeng committed
177
        if points is not None:
178
179
180
181
182
183
184
185
            assert isinstance(points, (torch.Tensor, np.ndarray, BasePoints))
            if isinstance(points, (torch.Tensor, np.ndarray)):
                if bev_direction == 'horizontal':
                    points[:, 0] = -points[:, 0]
                elif bev_direction == 'vertical':
                    points[:, 1] = -points[:, 1]
            elif isinstance(points, BasePoints):
                points.flip(bev_direction)
wuyuefeng's avatar
wuyuefeng committed
186
187
            return points

188
    def convert_to(self, dst, rt_mat=None):
Wenwei Zhang's avatar
Wenwei Zhang committed
189
        """Convert self to ``dst`` mode.
190
191

        Args:
192
            dst (:obj:`Box3DMode`): The target Box mode.
193
194
195
            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
196
                The conversion from ``src`` coordinates to ``dst`` coordinates
197
198
199
200
                usually comes along the change of sensors, e.g., from camera
                to LiDAR. This requires a transformation matrix.

        Returns:
201
            :obj:`DepthInstance3DBoxes`:
Wenwei Zhang's avatar
Wenwei Zhang committed
202
                The converted box of the same type in the ``dst`` mode.
203
204
205
206
        """
        from .box_3d_mode import Box3DMode
        return Box3DMode.convert(
            box=self, src=Box3DMode.DEPTH, dst=dst, rt_mat=rt_mat)
wuyuefeng's avatar
wuyuefeng committed
207

208
209
210
211
212
213
214
    def enlarged_box(self, extra_width):
        """Enlarge the length, width and height boxes.

        Args:
            extra_width (float | torch.Tensor): Extra width to enlarge the box.

        Returns:
215
            :obj:`DepthInstance3DBoxes`: Enlarged boxes.
216
217
218
219
220
221
222
        """
        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)

encore-zhou's avatar
encore-zhou committed
223
224
225
226
227
228
229
    def get_surface_line_center(self):
        """Compute surface and line center of bounding boxes.

        Returns:
            torch.Tensor: Surface and line center of bounding boxes.
        """
        obj_size = self.dims
230
        center = self.gravity_center.view(-1, 1, 3)
encore-zhou's avatar
encore-zhou committed
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
        batch_size = center.shape[0]

        rot_sin = torch.sin(-self.yaw)
        rot_cos = torch.cos(-self.yaw)
        rot_mat_T = self.yaw.new_zeros(tuple(list(self.yaw.shape) + [3, 3]))
        rot_mat_T[..., 0, 0] = rot_cos
        rot_mat_T[..., 0, 1] = -rot_sin
        rot_mat_T[..., 1, 0] = rot_sin
        rot_mat_T[..., 1, 1] = rot_cos
        rot_mat_T[..., 2, 2] = 1

        # Get the object surface center
        offset = obj_size.new_tensor([[0, 0, 1], [0, 0, -1], [0, 1, 0],
                                      [0, -1, 0], [1, 0, 0], [-1, 0, 0]])
        offset = offset.view(1, 6, 3) / 2
246
247
248
        surface_3d = (offset *
                      obj_size.view(batch_size, 1, 3).repeat(1, 6, 1)).reshape(
                          -1, 3)
encore-zhou's avatar
encore-zhou committed
249
250
251
252
253
254
255
256
257

        # Get the object line center
        offset = obj_size.new_tensor([[1, 0, 1], [-1, 0, 1], [0, 1, 1],
                                      [0, -1, 1], [1, 0, -1], [-1, 0, -1],
                                      [0, 1, -1], [0, -1, -1], [1, 1, 0],
                                      [1, -1, 0], [-1, 1, 0], [-1, -1, 0]])
        offset = offset.view(1, 12, 3) / 2

        line_3d = (offset *
258
259
                   obj_size.view(batch_size, 1, 3).repeat(1, 12, 1)).reshape(
                       -1, 3)
encore-zhou's avatar
encore-zhou committed
260
261

        surface_rot = rot_mat_T.repeat(6, 1, 1)
262
263
        surface_3d = torch.matmul(surface_3d.unsqueeze(-2),
                                  surface_rot).squeeze(-2)
264
        surface_center = center.repeat(1, 6, 1).reshape(-1, 3) + surface_3d
encore-zhou's avatar
encore-zhou committed
265
266

        line_rot = rot_mat_T.repeat(12, 1, 1)
267
        line_3d = torch.matmul(line_3d.unsqueeze(-2), line_rot).squeeze(-2)
268
        line_center = center.repeat(1, 12, 1).reshape(-1, 3) + line_3d
encore-zhou's avatar
encore-zhou committed
269
270

        return surface_center, line_center