"vscode:/vscode.git/clone" did not exist on "90cc9baade7b4e298d220b0f8735baeeeec0f4bb"
Commit 8c7d0586 authored by zhangwenwei's avatar zhangwenwei
Browse files

More docstrings with refactor

parent ce74cb0a
......@@ -6,15 +6,13 @@ import torch
class BaseInstance3DBoxes(object):
"""Base class for 3D Boxes
Args:
tensor (torch.Tensor | np.ndarray): a Nxbox_dim matrix.
box_dim (int): number of the dimension of a box
Each row is (x, y, z, x_size, y_size, z_size, yaw).
"""
def __init__(self, tensor, box_dim=7):
"""
Args:
tensor (torch.Tensor | np.ndarray): a Nxbox_dim matrix.
box_dim (int): number of the dimension of a box
Each row is (x, y, z, x_size, y_size, z_size, yaw).
"""
if isinstance(tensor, torch.Tensor):
device = tensor.device
else:
......@@ -29,7 +27,7 @@ class BaseInstance3DBoxes(object):
self.box_dim = box_dim
self.tensor = tensor
@abstractmethod
@property
def volume(self):
"""Computes the volume of all the boxes.
......@@ -38,7 +36,7 @@ class BaseInstance3DBoxes(object):
"""
return self.tensor[:, 3] * self.tensor[:, 4] * self.tensor[:, 5]
@abstractmethod
@property
def bottom_center(self):
"""Calculate the bottom center of all the boxes.
......@@ -47,7 +45,7 @@ class BaseInstance3DBoxes(object):
"""
return self.tensor[..., :3]
@abstractmethod
@property
def gravity_center(self):
"""Calculate the gravity center of all the boxes.
......@@ -56,7 +54,7 @@ class BaseInstance3DBoxes(object):
"""
pass
@abstractmethod
@property
def corners(self):
"""Calculate the coordinates of corners of all the boxes.
......@@ -104,6 +102,16 @@ class BaseInstance3DBoxes(object):
"""
pass
@abstractmethod
def scale(self, scale_factors):
"""Scale the box with horizontal and vertical scaling factors
Args:
scale_factors (float | torch.Tensor | list[float]):
scale factors to scale the boxes.
"""
pass
def nonempty(self, threshold: float = 0.0):
"""Find boxes that are non-empty.
......@@ -123,15 +131,6 @@ class BaseInstance3DBoxes(object):
& (size_y > threshold) & (size_z > threshold))
return keep
def scale(self, scale_factors):
"""Scale the box with horizontal and vertical scaling factors
Args:
scale_factors (float | torch.Tensor | list[float]):
scale factors to scale the boxes.
"""
pass
def __getitem__(self, item):
"""
Note:
......
......@@ -6,55 +6,56 @@ import torch
@unique
class Box3DMode(IntEnum):
"""
Enum of different ways to represent a box.
r"""Enum of different ways to represent a box.
Coordinates in velodyne/LiDAR sensors:
.. code-block:: none
up z x front
^ ^
| /
| /
left y <------ 0
Coordinates in camera:
.. code-block:: none
x right
/
/
front z <------ 0
|
|
v
down y
Coordinates in Depth mode:
.. code-block:: none
up z x right
^ ^
| /
| /
front y <------ 0
"""
LIDAR = 0
"""
Coordinates in velodyne/LiDAR sensors.
up z x front
^ ^
| /
| /
left y <------ 0
"""
CAM = 1
"""
Coordinates in camera.
x right
/
/
front z <------ 0
|
|
v
down y
"""
DEPTH = 2
"""
Coordinates in Depth mode.
up z x right
^ ^
| /
| /
front y <------ 0
"""
@staticmethod
def convert(box, from_mode, to_mode):
"""
def convert(box, src, dst):
"""Convert boxes from `src` mode to `dst` mode.
Args:
box: can be a k-tuple, k-list or an Nxk array/tensor, where k = 7
from_mode, to_mode (BoxMode)
box (tuple | list | np.ndarray | torch.Tensor):
can be a k-tuple, k-list or an Nxk array/tensor, where k = 7
src (BoxMode): the src Box mode
dst (BoxMode): the target Box mode
Returns:
The converted box of the same type.
"""
if from_mode == to_mode:
if src == dst:
return box
original_type = type(box)
is_numpy = isinstance(box, np.ndarray)
single_box = isinstance(box, (list, tuple))
if single_box:
......@@ -70,7 +71,8 @@ class Box3DMode(IntEnum):
arr = box.clone()
# converting logic
# TODO: add converting logic to support box conversion
original_type = type(box)
if single_box:
return original_type(arr.flatten().tolist())
if is_numpy:
......
......@@ -6,8 +6,10 @@ from .utils import limit_period, rotation_3d_in_axis
class LiDARInstance3DBoxes(BaseInstance3DBoxes):
"""
This structure stores a list of boxes as a Nx7 torch.Tensor.
"""3D boxes of instances in LIDAR coordinates
This structure stores a list of boxes as a NxM torch.Tensor,
where M >= 7.
It supports some common methods about boxes
(`area`, `clip`, `nonempty`, etc),
and also behaves like a Tensor
......@@ -20,6 +22,7 @@ class LiDARInstance3DBoxes(BaseInstance3DBoxes):
Each row is (x, y, z, x_size, y_size, z_size, yaw, ...).
"""
@property
def gravity_center(self):
"""Calculate the gravity center of all the boxes.
......@@ -32,17 +35,13 @@ class LiDARInstance3DBoxes(BaseInstance3DBoxes):
gravity_center[:, 2] = bottom_center[:, 2] + bottom_center[:, 5] * 0.5
return gravity_center
def corners(self, origin=[0.5, 1.0, 0.5], axis=1):
@property
def corners(self):
"""Calculate the coordinates of corners of all the boxes.
Convert the boxes to the form of
(x0y0z0, x0y0z1, x0y1z0, x0y1z1, x1y0z0, x1y0z1, x1y1z0, x1y1z1)
Args:
origin (list[float]): origin point relate to smallest point.
use [0.5, 1.0, 0.5] in camera and [0.5, 0.5, 0] in lidar.
axis (int): rotation axis. 1 for camera and 2 for lidar.
Returns:
torch.Tensor: corners of each box with size (N, 8, 3)
"""
......@@ -52,13 +51,14 @@ class LiDARInstance3DBoxes(BaseInstance3DBoxes):
device=dims.device, dtype=dims.dtype)
corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]]
corners_norm = corners_norm - dims.new_tensor(origin)
corners_norm = corners_norm - dims.new_tensor([0.5, 1.0, 0.5])
corners = dims.view([-1, 1, 3]) * corners_norm.reshape([1, 2**3, 3])
corners = rotation_3d_in_axis(corners, self.tensor[:, 6], axis=axis)
corners = rotation_3d_in_axis(corners, self.tensor[:, 6], axis=1)
corners += self.tensor[:, :3].view(-1, 1, 3)
return corners
@property
def nearset_bev(self):
"""Calculate the 2D bounding boxes in BEV without rotation
......
......@@ -3,12 +3,36 @@ import torch
def limit_period(val, offset=0.5, period=np.pi):
"""Limit the value into a period for periodic function.
Args:
val (torch.Tensor): The value to be converted
offset (float, optional): Offset to set the value range.
Defaults to 0.5.
period ([type], optional): Period of the value. Defaults to np.pi.
Returns:
torch.Tensor: value in the range of
[-offset * period, (1-offset) * period]
"""
return val - torch.floor(val / period + offset) * period
def rotation_3d_in_axis(points, angles, axis=0):
# points: [N, point_size, 3]
# angles: [N]
"""Rotate points by angles according to axis
Args:
points (torch.Tensor): Points of shape (N, M, 3).
angles (torch.Tensor): Vector of angles in shape (N,)
axis (int, optional): The axis to be rotated. Defaults to 0.
Raises:
ValueError: when the axis is not in range [0, 1, 2], it will
raise value error.
Returns:
torch.Tensor: rotated points in shape (N, M, 3)
"""
rot_sin = torch.sin(angles)
rot_cos = torch.cos(angles)
ones = torch.ones_like(rot_cos)
......@@ -32,6 +56,6 @@ def rotation_3d_in_axis(points, angles, axis=0):
torch.stack([ones, zeros, zeros])
])
else:
raise ValueError('axis should in range')
raise ValueError(f'axis should in range [0, 1, 2], got {axis}')
return torch.einsum('aij,jka->aik', (points, rot_mat_T))
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment