Commit 8c7d0586 authored by zhangwenwei's avatar zhangwenwei
Browse files

More docstrings with refactor

parent ce74cb0a
...@@ -6,15 +6,13 @@ import torch ...@@ -6,15 +6,13 @@ import torch
class BaseInstance3DBoxes(object): class BaseInstance3DBoxes(object):
"""Base class for 3D Boxes """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): 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): if isinstance(tensor, torch.Tensor):
device = tensor.device device = tensor.device
else: else:
...@@ -29,7 +27,7 @@ class BaseInstance3DBoxes(object): ...@@ -29,7 +27,7 @@ class BaseInstance3DBoxes(object):
self.box_dim = box_dim self.box_dim = box_dim
self.tensor = tensor self.tensor = tensor
@abstractmethod @property
def volume(self): def volume(self):
"""Computes the volume of all the boxes. """Computes the volume of all the boxes.
...@@ -38,7 +36,7 @@ class BaseInstance3DBoxes(object): ...@@ -38,7 +36,7 @@ class BaseInstance3DBoxes(object):
""" """
return self.tensor[:, 3] * self.tensor[:, 4] * self.tensor[:, 5] return self.tensor[:, 3] * self.tensor[:, 4] * self.tensor[:, 5]
@abstractmethod @property
def bottom_center(self): def bottom_center(self):
"""Calculate the bottom center of all the boxes. """Calculate the bottom center of all the boxes.
...@@ -47,7 +45,7 @@ class BaseInstance3DBoxes(object): ...@@ -47,7 +45,7 @@ class BaseInstance3DBoxes(object):
""" """
return self.tensor[..., :3] return self.tensor[..., :3]
@abstractmethod @property
def gravity_center(self): def gravity_center(self):
"""Calculate the gravity center of all the boxes. """Calculate the gravity center of all the boxes.
...@@ -56,7 +54,7 @@ class BaseInstance3DBoxes(object): ...@@ -56,7 +54,7 @@ class BaseInstance3DBoxes(object):
""" """
pass pass
@abstractmethod @property
def corners(self): def corners(self):
"""Calculate the coordinates of corners of all the boxes. """Calculate the coordinates of corners of all the boxes.
...@@ -104,6 +102,16 @@ class BaseInstance3DBoxes(object): ...@@ -104,6 +102,16 @@ class BaseInstance3DBoxes(object):
""" """
pass 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): def nonempty(self, threshold: float = 0.0):
"""Find boxes that are non-empty. """Find boxes that are non-empty.
...@@ -123,15 +131,6 @@ class BaseInstance3DBoxes(object): ...@@ -123,15 +131,6 @@ class BaseInstance3DBoxes(object):
& (size_y > threshold) & (size_z > threshold)) & (size_y > threshold) & (size_z > threshold))
return keep 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): def __getitem__(self, item):
""" """
Note: Note:
......
...@@ -6,55 +6,56 @@ import torch ...@@ -6,55 +6,56 @@ import torch
@unique @unique
class Box3DMode(IntEnum): class Box3DMode(IntEnum):
""" r"""Enum of different ways to represent a box.
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 LIDAR = 0
"""
Coordinates in velodyne/LiDAR sensors.
up z x front
^ ^
| /
| /
left y <------ 0
"""
CAM = 1 CAM = 1
"""
Coordinates in camera.
x right
/
/
front z <------ 0
|
|
v
down y
"""
DEPTH = 2 DEPTH = 2
"""
Coordinates in Depth mode.
up z x right
^ ^
| /
| /
front y <------ 0
"""
@staticmethod @staticmethod
def convert(box, from_mode, to_mode): def convert(box, src, dst):
""" """Convert boxes from `src` mode to `dst` mode.
Args: Args:
box: can be a k-tuple, k-list or an Nxk array/tensor, where k = 7 box (tuple | list | np.ndarray | torch.Tensor):
from_mode, to_mode (BoxMode) 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: Returns:
The converted box of the same type. The converted box of the same type.
""" """
if from_mode == to_mode: if src == dst:
return box return box
original_type = type(box)
is_numpy = isinstance(box, np.ndarray) is_numpy = isinstance(box, np.ndarray)
single_box = isinstance(box, (list, tuple)) single_box = isinstance(box, (list, tuple))
if single_box: if single_box:
...@@ -70,7 +71,8 @@ class Box3DMode(IntEnum): ...@@ -70,7 +71,8 @@ class Box3DMode(IntEnum):
arr = box.clone() arr = box.clone()
# converting logic # converting logic
# TODO: add converting logic to support box conversion
original_type = type(box)
if single_box: if single_box:
return original_type(arr.flatten().tolist()) return original_type(arr.flatten().tolist())
if is_numpy: if is_numpy:
......
...@@ -6,8 +6,10 @@ from .utils import limit_period, rotation_3d_in_axis ...@@ -6,8 +6,10 @@ from .utils import limit_period, rotation_3d_in_axis
class LiDARInstance3DBoxes(BaseInstance3DBoxes): class LiDARInstance3DBoxes(BaseInstance3DBoxes):
""" """3D boxes of instances in LIDAR coordinates
This structure stores a list of boxes as a Nx7 torch.Tensor.
This structure stores a list of boxes as a NxM torch.Tensor,
where M >= 7.
It supports some common methods about boxes It supports some common methods about boxes
(`area`, `clip`, `nonempty`, etc), (`area`, `clip`, `nonempty`, etc),
and also behaves like a Tensor and also behaves like a Tensor
...@@ -20,6 +22,7 @@ class LiDARInstance3DBoxes(BaseInstance3DBoxes): ...@@ -20,6 +22,7 @@ class LiDARInstance3DBoxes(BaseInstance3DBoxes):
Each row is (x, y, z, x_size, y_size, z_size, yaw, ...). Each row is (x, y, z, x_size, y_size, z_size, yaw, ...).
""" """
@property
def gravity_center(self): def gravity_center(self):
"""Calculate the gravity center of all the boxes. """Calculate the gravity center of all the boxes.
...@@ -32,17 +35,13 @@ class LiDARInstance3DBoxes(BaseInstance3DBoxes): ...@@ -32,17 +35,13 @@ class LiDARInstance3DBoxes(BaseInstance3DBoxes):
gravity_center[:, 2] = bottom_center[:, 2] + bottom_center[:, 5] * 0.5 gravity_center[:, 2] = bottom_center[:, 2] + bottom_center[:, 5] * 0.5
return gravity_center 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. """Calculate the coordinates of corners of all the boxes.
Convert the boxes to the form of Convert the boxes to the form of
(x0y0z0, x0y0z1, x0y1z0, x0y1z1, x1y0z0, x1y0z1, x1y1z0, x1y1z1) (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: Returns:
torch.Tensor: corners of each box with size (N, 8, 3) torch.Tensor: corners of each box with size (N, 8, 3)
""" """
...@@ -52,13 +51,14 @@ class LiDARInstance3DBoxes(BaseInstance3DBoxes): ...@@ -52,13 +51,14 @@ class LiDARInstance3DBoxes(BaseInstance3DBoxes):
device=dims.device, dtype=dims.dtype) device=dims.device, dtype=dims.dtype)
corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]] 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 = 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) corners += self.tensor[:, :3].view(-1, 1, 3)
return corners return corners
@property
def nearset_bev(self): def nearset_bev(self):
"""Calculate the 2D bounding boxes in BEV without rotation """Calculate the 2D bounding boxes in BEV without rotation
......
...@@ -3,12 +3,36 @@ import torch ...@@ -3,12 +3,36 @@ import torch
def limit_period(val, offset=0.5, period=np.pi): 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 return val - torch.floor(val / period + offset) * period
def rotation_3d_in_axis(points, angles, axis=0): def rotation_3d_in_axis(points, angles, axis=0):
# points: [N, point_size, 3] """Rotate points by angles according to axis
# angles: [N]
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_sin = torch.sin(angles)
rot_cos = torch.cos(angles) rot_cos = torch.cos(angles)
ones = torch.ones_like(rot_cos) ones = torch.ones_like(rot_cos)
...@@ -32,6 +56,6 @@ def rotation_3d_in_axis(points, angles, axis=0): ...@@ -32,6 +56,6 @@ def rotation_3d_in_axis(points, angles, axis=0):
torch.stack([ones, zeros, zeros]) torch.stack([ones, zeros, zeros])
]) ])
else: 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)) 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