depth_box3d.py 10.8 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
2
3
from typing import Optional, Tuple, Union

wuyuefeng's avatar
wuyuefeng committed
4
5
import numpy as np
import torch
6
from torch import Tensor
wuyuefeng's avatar
wuyuefeng committed
7

zhangshilong's avatar
zhangshilong committed
8
from mmdet3d.structures.points import BasePoints
wuyuefeng's avatar
wuyuefeng committed
9
from .base_box3d import BaseInstance3DBoxes
10
from .utils import rotation_3d_in_axis
wuyuefeng's avatar
wuyuefeng committed
11
12
13


class DepthInstance3DBoxes(BaseInstance3DBoxes):
14
    """3D boxes of instances in DEPTH coordinates.
wuyuefeng's avatar
wuyuefeng committed
15
16

    Coordinates in Depth:
17

wuyuefeng's avatar
wuyuefeng committed
18
19
    .. code-block:: none

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

wuyuefeng's avatar
wuyuefeng committed
26
    The relative coordinate of bottom center in a Depth box is (0.5, 0.5, 0),
27
28
29
    and the yaw is around the z axis, thus the rotation axis=2. 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.
wuyuefeng's avatar
wuyuefeng committed
30

Wenwei Zhang's avatar
Wenwei Zhang committed
31
    Attributes:
32
33
34
        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, ...).
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.
wuyuefeng's avatar
wuyuefeng committed
37
    """
38
    YAW_AXIS = 2
wuyuefeng's avatar
wuyuefeng committed
39
40

    @property
41
42
43
    def corners(self) -> Tensor:
        """Convert boxes to corners in clockwise order, in the form of (x0y0z0,
        x0y0z1, x0y1z1, x0y1z0, x1y0z0, x1y0z1, x1y1z1, x1y1z0).
wuyuefeng's avatar
wuyuefeng committed
44
45
46

        .. code-block:: none

47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
                                        up z
                         front y           ^
                              /            |
                             /             |
               (x0, y1, z1) + -----------  + (x1, y1, z1)
                           /|            / |
                          / |           /  |
            (x0, y0, z1) + ----------- +   + (x1, y1, z0)
                         |  /      .   |  /
                         | / origin    | /
            (x0, y0, z0) + ----------- + --------> right x
                                       (x1, y0, z0)

        Returns:
            Tensor: A tensor with 8 corners of each box in shape (N, 8, 3).
wuyuefeng's avatar
wuyuefeng committed
62
        """
63
64
65
        if self.tensor.numel() == 0:
            return torch.empty([0, 8, 3], device=self.tensor.device)

wuyuefeng's avatar
wuyuefeng committed
66
67
68
69
70
71
        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
72
        # use relative origin (0.5, 0.5, 0)
wuyuefeng's avatar
wuyuefeng committed
73
74
75
76
        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
77
78
        corners = rotation_3d_in_axis(
            corners, self.tensor[:, 6], axis=self.YAW_AXIS)
wuyuefeng's avatar
wuyuefeng committed
79
80
81
        corners += self.tensor[:, :3].view(-1, 1, 3)
        return corners

82
83
84
85
86
87
    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]:
88
89
        """Rotate boxes with points (optional) with the given angle or rotation
        matrix.
wuyuefeng's avatar
wuyuefeng committed
90
91

        Args:
92
93
94
            angle (Tensor or np.ndarray or float): Rotation angle or rotation
                matrix.
            points (Tensor or np.ndarray or :obj:`BasePoints`, optional):
95
                Points to rotate. Defaults to None.
wuyuefeng's avatar
wuyuefeng committed
96
97

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

105
106
107
108
        assert angle.shape == torch.Size([3, 3]) or angle.numel() == 1, \
            f'invalid rotation angle shape {angle.shape}'

        if angle.numel() == 1:
109
110
111
112
113
            self.tensor[:, 0:3], rot_mat_T = rotation_3d_in_axis(
                self.tensor[:, 0:3],
                angle,
                axis=self.YAW_AXIS,
                return_mat=True)
114
        else:
115
            rot_mat_T = angle
116
117
118
            rot_sin = rot_mat_T[0, 1]
            rot_cos = rot_mat_T[0, 0]
            angle = np.arctan2(rot_sin, rot_cos)
119
            self.tensor[:, 0:3] = self.tensor[:, 0:3] @ rot_mat_T
120

wuyuefeng's avatar
wuyuefeng committed
121
        if self.with_yaw:
122
            self.tensor[:, 6] += angle
wuyuefeng's avatar
wuyuefeng committed
123
        else:
124
125
            # for axis-aligned boxes, we take the new
            # enclosing axis-aligned boxes after rotation
wuyuefeng's avatar
wuyuefeng committed
126
            corners_rot = self.corners @ rot_mat_T
wuyuefeng's avatar
wuyuefeng committed
127
128
129
130
131
132
133
134
            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
135
        if points is not None:
136
            if isinstance(points, Tensor):
wuyuefeng's avatar
wuyuefeng committed
137
138
                points[:, :3] = points[:, :3] @ rot_mat_T
            elif isinstance(points, np.ndarray):
139
                rot_mat_T = rot_mat_T.cpu().numpy()
wuyuefeng's avatar
wuyuefeng committed
140
                points[:, :3] = np.dot(points[:, :3], rot_mat_T)
141
            elif isinstance(points, BasePoints):
142
                points.rotate(rot_mat_T)
wuyuefeng's avatar
wuyuefeng committed
143
144
145
146
            else:
                raise ValueError
            return points, rot_mat_T

147
148
149
150
151
    def flip(
        self,
        bev_direction: str = 'horizontal',
        points: Optional[Union[Tensor, np.ndarray, BasePoints]] = None
    ) -> Union[Tensor, np.ndarray, BasePoints, None]:
zhangwenwei's avatar
zhangwenwei committed
152
        """Flip the boxes in BEV along given BEV direction.
wuyuefeng's avatar
wuyuefeng committed
153

154
        In Depth coordinates, it flips the x (horizontal) or y (vertical) axis.
wuyuefeng's avatar
wuyuefeng committed
155
156

        Args:
157
158
159
            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):
160
                Points to flip. Defaults to None.
wuyuefeng's avatar
wuyuefeng committed
161
162

        Returns:
163
164
165
            Tensor or np.ndarray or :obj:`BasePoints` or None: When ``points``
            is None, the function returns None, otherwise it returns the
            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
            assert isinstance(points, (Tensor, np.ndarray, BasePoints))
            if isinstance(points, (Tensor, np.ndarray)):
180
181
182
183
184
185
                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
189
190
191
    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
192
        """Convert self to ``dst`` mode.
193
194

        Args:
195
196
            dst (int): The target Box mode.
            rt_mat (Tensor or np.ndarray, optional): The rotation and
197
                translation matrix between different coordinates.
198
199
200
201
202
203
                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.
204
205

        Returns:
206
207
            :obj:`BaseInstance3DBoxes`: The converted box of the same type in
            the ``dst`` mode.
208
209
210
        """
        from .box_3d_mode import Box3DMode
        return Box3DMode.convert(
211
212
213
214
215
            box=self,
            src=Box3DMode.DEPTH,
            dst=dst,
            rt_mat=rt_mat,
            correct_yaw=correct_yaw)
wuyuefeng's avatar
wuyuefeng committed
216

217
218
219
    def enlarged_box(
            self, extra_width: Union[float, Tensor]) -> 'DepthInstance3DBoxes':
        """Enlarge the length, width and height of boxes.
220
221

        Args:
222
            extra_width (float or Tensor): Extra width to enlarge the box.
223
224

        Returns:
225
            :obj:`DepthInstance3DBoxes`: Enlarged boxes.
226
227
228
229
230
231
232
        """
        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)

233
    def get_surface_line_center(self) -> Tuple[Tensor, Tensor]:
encore-zhou's avatar
encore-zhou committed
234
235
236
        """Compute surface and line center of bounding boxes.

        Returns:
237
            Tuple[Tensor, Tensor]: Surface and line center of bounding boxes.
encore-zhou's avatar
encore-zhou committed
238
239
        """
        obj_size = self.dims
240
        center = self.gravity_center.view(-1, 1, 3)
encore-zhou's avatar
encore-zhou committed
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
        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
256
257
258
        surface_3d = (offset *
                      obj_size.view(batch_size, 1, 3).repeat(1, 6, 1)).reshape(
                          -1, 3)
encore-zhou's avatar
encore-zhou committed
259
260
261
262
263
264
265
266
267

        # 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 *
268
269
                   obj_size.view(batch_size, 1, 3).repeat(1, 12, 1)).reshape(
                       -1, 3)
encore-zhou's avatar
encore-zhou committed
270
271

        surface_rot = rot_mat_T.repeat(6, 1, 1)
272
273
        surface_3d = torch.matmul(surface_3d.unsqueeze(-2),
                                  surface_rot).squeeze(-2)
274
        surface_center = center.repeat(1, 6, 1).reshape(-1, 3) + surface_3d
encore-zhou's avatar
encore-zhou committed
275
276

        line_rot = rot_mat_T.repeat(12, 1, 1)
277
        line_3d = torch.matmul(line_3d.unsqueeze(-2), line_rot).squeeze(-2)
278
        line_center = center.repeat(1, 12, 1).reshape(-1, 3) + line_3d
encore-zhou's avatar
encore-zhou committed
279
280

        return surface_center, line_center