321 lines
12 KiB
Python
321 lines
12 KiB
Python
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from mmdet3d.core.points import BasePoints
|
||
|
|
from .base_box3d import BaseInstance3DBoxes
|
||
|
|
from .utils import limit_period, rotation_3d_in_axis
|
||
|
|
|
||
|
|
|
||
|
|
class CameraInstance3DBoxes(BaseInstance3DBoxes):
|
||
|
|
"""3D boxes of instances in CAM coordinates.
|
||
|
|
|
||
|
|
Coordinates in camera:
|
||
|
|
|
||
|
|
.. code-block:: none
|
||
|
|
|
||
|
|
z front (yaw=-0.5*pi)
|
||
|
|
/
|
||
|
|
/
|
||
|
|
0 ------> x right (yaw=0)
|
||
|
|
|
|
||
|
|
|
|
||
|
|
v
|
||
|
|
down y
|
||
|
|
|
||
|
|
The relative coordinate of bottom center in a CAM box is (0.5, 1.0, 0.5),
|
||
|
|
and the yaw is around the y axis, thus the rotation axis=1.
|
||
|
|
The yaw is 0 at the positive direction of x axis, and decreases from
|
||
|
|
the positive direction of x to the positive direction of z.
|
||
|
|
|
||
|
|
A refactor is ongoing to make the three coordinate systems
|
||
|
|
easier to understand and convert between each other.
|
||
|
|
|
||
|
|
Attributes:
|
||
|
|
tensor (torch.Tensor): Float matrix of N x box_dim.
|
||
|
|
box_dim (int): Integer indicates the dimension of a box
|
||
|
|
Each row is (x, y, z, x_size, y_size, z_size, yaw, ...).
|
||
|
|
with_yaw (bool): If True, the value of yaw will be set to 0 as minmax
|
||
|
|
boxes.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, tensor, box_dim=7, with_yaw=True, origin=(0.5, 1.0, 0.5)):
|
||
|
|
if isinstance(tensor, torch.Tensor):
|
||
|
|
device = tensor.device
|
||
|
|
else:
|
||
|
|
device = torch.device("cpu")
|
||
|
|
tensor = torch.as_tensor(tensor, dtype=torch.float32, device=device)
|
||
|
|
if tensor.numel() == 0:
|
||
|
|
# 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((0, box_dim)).to(dtype=torch.float32, device=device)
|
||
|
|
assert tensor.dim() == 2 and tensor.size(-1) == box_dim, tensor.size()
|
||
|
|
|
||
|
|
if tensor.shape[-1] == 6:
|
||
|
|
# If the dimension of boxes is 6, we expand box_dim by padding
|
||
|
|
# 0 as a fake yaw and set with_yaw to False.
|
||
|
|
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
|
||
|
|
self.with_yaw = False
|
||
|
|
else:
|
||
|
|
self.box_dim = box_dim
|
||
|
|
self.with_yaw = with_yaw
|
||
|
|
self.tensor = tensor.clone()
|
||
|
|
|
||
|
|
if origin != (0.5, 1.0, 0.5):
|
||
|
|
dst = self.tensor.new_tensor((0.5, 1.0, 0.5))
|
||
|
|
src = self.tensor.new_tensor(origin)
|
||
|
|
self.tensor[:, :3] += self.tensor[:, 3:6] * (dst - src)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def height(self):
|
||
|
|
"""torch.Tensor: A vector with height of each box."""
|
||
|
|
return self.tensor[:, 4]
|
||
|
|
|
||
|
|
@property
|
||
|
|
def top_height(self):
|
||
|
|
"""torch.Tensor: A vector with the top height of each box."""
|
||
|
|
# the positive direction is down rather than up
|
||
|
|
return self.bottom_height - self.height
|
||
|
|
|
||
|
|
@property
|
||
|
|
def bottom_height(self):
|
||
|
|
"""torch.Tensor: A vector with bottom's height of each box."""
|
||
|
|
return self.tensor[:, 1]
|
||
|
|
|
||
|
|
@property
|
||
|
|
def gravity_center(self):
|
||
|
|
"""torch.Tensor: A tensor with center of each box."""
|
||
|
|
bottom_center = self.bottom_center
|
||
|
|
gravity_center = torch.zeros_like(bottom_center)
|
||
|
|
gravity_center[:, [0, 2]] = bottom_center[:, [0, 2]]
|
||
|
|
gravity_center[:, 1] = bottom_center[:, 1] - self.tensor[:, 4] * 0.5
|
||
|
|
return gravity_center
|
||
|
|
|
||
|
|
@property
|
||
|
|
def corners(self):
|
||
|
|
"""torch.Tensor: Coordinates of corners of all the boxes in
|
||
|
|
shape (N, 8, 3).
|
||
|
|
|
||
|
|
Convert the boxes to in clockwise order, in the form of
|
||
|
|
(x0y0z0, x0y0z1, x0y1z1, x0y1z0, x1y0z0, x1y0z1, x1y1z1, x1y1z0)
|
||
|
|
|
||
|
|
.. code-block:: none
|
||
|
|
|
||
|
|
front z
|
||
|
|
/
|
||
|
|
/
|
||
|
|
(x0, y0, z1) + ----------- + (x1, y0, z1)
|
||
|
|
/| / |
|
||
|
|
/ | / |
|
||
|
|
(x0, y0, z0) + ----------- + + (x1, y1, z1)
|
||
|
|
| / . | /
|
||
|
|
| / origin | /
|
||
|
|
(x0, y1, z0) + ----------- + -------> x right
|
||
|
|
| (x1, y1, z0)
|
||
|
|
|
|
||
|
|
v
|
||
|
|
down y
|
||
|
|
"""
|
||
|
|
# TODO: rotation_3d_in_axis function do not support
|
||
|
|
# empty tensor currently.
|
||
|
|
assert len(self.tensor) != 0
|
||
|
|
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]]
|
||
|
|
# use relative origin [0.5, 1, 0.5]
|
||
|
|
corners_norm = corners_norm - dims.new_tensor([0.5, 1, 0.5])
|
||
|
|
corners = dims.view([-1, 1, 3]) * corners_norm.reshape([1, 8, 3])
|
||
|
|
|
||
|
|
# rotate around y axis
|
||
|
|
corners = rotation_3d_in_axis(corners, self.tensor[:, 6], axis=1)
|
||
|
|
corners += self.tensor[:, :3].view(-1, 1, 3)
|
||
|
|
return corners
|
||
|
|
|
||
|
|
@property
|
||
|
|
def bev(self):
|
||
|
|
"""torch.Tensor: A n x 5 tensor of 2D BEV box of each box
|
||
|
|
with rotation in XYWHR format."""
|
||
|
|
return self.tensor[:, [0, 2, 3, 5, 6]]
|
||
|
|
|
||
|
|
@property
|
||
|
|
def nearest_bev(self):
|
||
|
|
"""torch.Tensor: A tensor of 2D BEV box of each box
|
||
|
|
without rotation."""
|
||
|
|
# Obtain BEV boxes with rotation in XZWHR 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
|
||
|
|
|
||
|
|
def rotate(self, angle, points=None):
|
||
|
|
"""Rotate boxes with points (optional) with the given angle or \
|
||
|
|
rotation matrix.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
angle (float | torch.Tensor | np.ndarray):
|
||
|
|
Rotation angle or rotation matrix.
|
||
|
|
points (torch.Tensor, numpy.ndarray, :obj:`BasePoints`, optional):
|
||
|
|
Points to rotate. Defaults to None.
|
||
|
|
|
||
|
|
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``.
|
||
|
|
"""
|
||
|
|
if not isinstance(angle, torch.Tensor):
|
||
|
|
angle = self.tensor.new_tensor(angle)
|
||
|
|
assert (
|
||
|
|
angle.shape == torch.Size([3, 3]) or angle.numel() == 1
|
||
|
|
), f"invalid rotation angle shape {angle.shape}"
|
||
|
|
|
||
|
|
if angle.numel() == 1:
|
||
|
|
rot_sin = torch.sin(angle)
|
||
|
|
rot_cos = torch.cos(angle)
|
||
|
|
rot_mat_T = self.tensor.new_tensor(
|
||
|
|
[[rot_cos, 0, -rot_sin], [0, 1, 0], [rot_sin, 0, rot_cos]]
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
rot_mat_T = angle
|
||
|
|
rot_sin = rot_mat_T[2, 0]
|
||
|
|
rot_cos = rot_mat_T[0, 0]
|
||
|
|
angle = np.arctan2(rot_sin, rot_cos)
|
||
|
|
|
||
|
|
self.tensor[:, :3] = self.tensor[:, :3] @ rot_mat_T
|
||
|
|
self.tensor[:, 6] += angle
|
||
|
|
|
||
|
|
if points is not None:
|
||
|
|
if isinstance(points, torch.Tensor):
|
||
|
|
points[:, :3] = points[:, :3] @ rot_mat_T
|
||
|
|
elif isinstance(points, np.ndarray):
|
||
|
|
rot_mat_T = rot_mat_T.numpy()
|
||
|
|
points[:, :3] = np.dot(points[:, :3], rot_mat_T)
|
||
|
|
elif isinstance(points, BasePoints):
|
||
|
|
# clockwise
|
||
|
|
points.rotate(-angle)
|
||
|
|
else:
|
||
|
|
raise ValueError
|
||
|
|
return points, rot_mat_T
|
||
|
|
|
||
|
|
def flip(self, bev_direction="horizontal", points=None):
|
||
|
|
"""Flip the boxes in BEV along given BEV direction.
|
||
|
|
|
||
|
|
In CAM coordinates, it flips the x (horizontal) or z (vertical) axis.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
bev_direction (str): Flip direction (horizontal or vertical).
|
||
|
|
points (torch.Tensor, numpy.ndarray, :obj:`BasePoints`, None):
|
||
|
|
Points to flip. Defaults to None.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
torch.Tensor, numpy.ndarray or None: Flipped points.
|
||
|
|
"""
|
||
|
|
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[:, 2::7] = -self.tensor[:, 2::7]
|
||
|
|
if self.with_yaw:
|
||
|
|
self.tensor[:, 6] = -self.tensor[:, 6]
|
||
|
|
|
||
|
|
if points is not None:
|
||
|
|
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[:, 2] = -points[:, 2]
|
||
|
|
elif isinstance(points, BasePoints):
|
||
|
|
points.flip(bev_direction)
|
||
|
|
return points
|
||
|
|
|
||
|
|
def in_range_bev(self, box_range):
|
||
|
|
"""Check whether the boxes are in the given range.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
box_range (list | torch.Tensor): The range of box
|
||
|
|
(x_min, z_min, x_max, z_max).
|
||
|
|
|
||
|
|
Note:
|
||
|
|
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.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
torch.Tensor: Indicating whether each box is inside \
|
||
|
|
the reference range.
|
||
|
|
"""
|
||
|
|
in_range_flags = (
|
||
|
|
(self.tensor[:, 0] > box_range[0])
|
||
|
|
& (self.tensor[:, 2] > box_range[1])
|
||
|
|
& (self.tensor[:, 0] < box_range[2])
|
||
|
|
& (self.tensor[:, 2] < box_range[3])
|
||
|
|
)
|
||
|
|
return in_range_flags
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def height_overlaps(cls, boxes1, boxes2, mode="iou"):
|
||
|
|
"""Calculate height overlaps of two boxes.
|
||
|
|
|
||
|
|
This function calculates the height overlaps between ``boxes1`` and
|
||
|
|
``boxes2``, where ``boxes1`` and ``boxes2`` should be in the same type.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
boxes1 (:obj:`CameraInstance3DBoxes`): Boxes 1 contain N boxes.
|
||
|
|
boxes2 (:obj:`CameraInstance3DBoxes`): Boxes 2 contain M boxes.
|
||
|
|
mode (str, optional): Mode of iou calculation. Defaults to 'iou'.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
torch.Tensor: Calculated iou of boxes' heights.
|
||
|
|
"""
|
||
|
|
assert isinstance(boxes1, CameraInstance3DBoxes)
|
||
|
|
assert isinstance(boxes2, CameraInstance3DBoxes)
|
||
|
|
|
||
|
|
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)
|
||
|
|
|
||
|
|
# In camera coordinate system
|
||
|
|
# from up to down is the positive direction
|
||
|
|
heighest_of_bottom = torch.min(boxes1_bottom_height, boxes2_bottom_height)
|
||
|
|
lowest_of_top = torch.max(boxes1_top_height, boxes2_top_height)
|
||
|
|
overlaps_h = torch.clamp(heighest_of_bottom - lowest_of_top, min=0)
|
||
|
|
return overlaps_h
|
||
|
|
|
||
|
|
def convert_to(self, dst, rt_mat=None):
|
||
|
|
"""Convert self to ``dst`` mode.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
dst (:obj:`Box3DMode`): The target Box mode.
|
||
|
|
rt_mat (np.ndarray | torch.Tensor): The rotation and translation
|
||
|
|
matrix between different coordinates. 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.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
:obj:`BaseInstance3DBoxes`: \
|
||
|
|
The converted box of the same type in the ``dst`` mode.
|
||
|
|
"""
|
||
|
|
from .box_3d_mode import Box3DMode
|
||
|
|
|
||
|
|
return Box3DMode.convert(box=self, src=Box3DMode.CAM, dst=dst, rt_mat=rt_mat)
|