Skip to content

GitLab

Commit 53435c62 authored by Yezhen Cong's avatar Yezhen Cong Committed by Tai-Wang
Browse files

[Refactor] Refactor code structure and docstrings (#803) 

* refactor points_in_boxes

* Merge same functions of three boxes

* More docstring fixes and unify x/y/z size

* Add "optional" and fix "Default"

* Add "optional" and fix "Default"

* Add "optional" and fix "Default"

* Add "optional" and fix "Default"

* Add "optional" and fix "Default"

* Remove None in function param type

* Fix unittest

* Add comments for NMS functions

* Merge methods of Points

* Add unittest

* Add optional and default value

* Fix box conversion and add unittest

* Fix comments

* Add unit test

* Indent

* Fix CI

* Remove useless \\

* Remove useless \\

* Remove useless \\

* Remove useless \\

* Remove useless \\

* Add unit test for box bev

* More unit tests and refine docstrings in box_np_ops

* Fix comment

* Add deprecation warning
parent 4f36084f
Hide whitespace changes
Inline Side-by-side
Showing with 373 additions and 194 deletions
mmdet3d/ops/roiaware_pool3d/__init__.py
View file @ 53435c62
from .points_in_boxes import (points_in_boxes_batch, points_in_boxes_cpu, from .points_in_boxes import (points_in_boxes_all, points_in_boxes_cpu,
points_in_boxes_gpu) points_in_boxes_part)
from .roiaware_pool3d import RoIAwarePool3d from .roiaware_pool3d import RoIAwarePool3d
__all__ = [ __all__ = [
'RoIAwarePool3d', 'points_in_boxes_gpu', 'points_in_boxes_cpu', 'RoIAwarePool3d', 'points_in_boxes_part', 'points_in_boxes_cpu',
'points_in_boxes_batch' 'points_in_boxes_all'
] ]
mmdet3d/ops/roiaware_pool3d/points_in_boxes.py
View file @ 53435c62
...@@ -3,13 +3,13 @@ import torch ...@@ -3,13 +3,13 @@ import torch
from . import roiaware_pool3d_ext from . import roiaware_pool3d_ext
def points_in_boxes_gpu(points, boxes): def points_in_boxes_part(points, boxes):
"""Find points that are in boxes (CUDA) """Find the box in which each point is (CUDA).
Args: Args:
points (torch.Tensor): [B, M, 3], [x, y, z] in LiDAR/DEPTH coordinate points (torch.Tensor): [B, M, 3], [x, y, z] in LiDAR/DEPTH coordinate
boxes (torch.Tensor): [B, T, 7], boxes (torch.Tensor): [B, T, 7],
num_valid_boxes <= T, [x, y, z, dx, dy, dz, rz] in num_valid_boxes <= T, [x, y, z, x_size, y_size, z_size, rz] in
LiDAR/DEPTH coordinate, (x, y, z) is the bottom center LiDAR/DEPTH coordinate, (x, y, z) is the bottom center
Returns: Returns:
...@@ -43,25 +43,26 @@ def points_in_boxes_gpu(points, boxes): ...@@ -43,25 +43,26 @@ def points_in_boxes_gpu(points, boxes):
if torch.cuda.current_device() != points_device: if torch.cuda.current_device() != points_device:
torch.cuda.set_device(points_device) torch.cuda.set_device(points_device)
roiaware_pool3d_ext.points_in_boxes_gpu(boxes.contiguous(), roiaware_pool3d_ext.points_in_boxes_part(boxes.contiguous(),
points.contiguous(), points.contiguous(),
box_idxs_of_pts) box_idxs_of_pts)
return box_idxs_of_pts return box_idxs_of_pts
def points_in_boxes_cpu(points, boxes): def points_in_boxes_cpu(points, boxes):
"""Find points that are in boxes (CPU) """Find all boxes in which each point is (CPU). The CPU version of
:meth:`points_in_boxes_all`.
Args: Args:
points (torch.Tensor): [B, M, 3], [x, y, z] in points (torch.Tensor): [B, M, 3], [x, y, z] in
LiDAR/DEPTH coordinate LiDAR/DEPTH coordinate
boxes (torch.Tensor): [B, T, 7], boxes (torch.Tensor): [B, T, 7],
num_valid_boxes <= T, [x, y, z, dx, dy, dz, rz], num_valid_boxes <= T, [x, y, z, x_size, y_size, z_size, rz],
(x, y, z) is the bottom center. (x, y, z) is the bottom center.
Returns: Returns:
box_idxs_of_pts (torch.Tensor): (B, M, T), default background = 0 box_idxs_of_pts (torch.Tensor): (B, M, T), default background = 0.
""" """
assert points.shape[0] == boxes.shape[0], \ assert points.shape[0] == boxes.shape[0], \
f'Points and boxes should have the same batch size, ' \ f'Points and boxes should have the same batch size, ' \
...@@ -86,17 +87,17 @@ def points_in_boxes_cpu(points, boxes): ...@@ -86,17 +87,17 @@ def points_in_boxes_cpu(points, boxes):
return point_indices return point_indices
def points_in_boxes_batch(points, boxes): def points_in_boxes_all(points, boxes):
"""Find points that are in boxes (CUDA) """Find all boxes in which each point is (CUDA).
Args: Args:
points (torch.Tensor): [B, M, 3], [x, y, z] in LiDAR/DEPTH coordinate points (torch.Tensor): [B, M, 3], [x, y, z] in LiDAR/DEPTH coordinate
boxes (torch.Tensor): [B, T, 7], boxes (torch.Tensor): [B, T, 7],
num_valid_boxes <= T, [x, y, z, dx, dy, dz, rz], num_valid_boxes <= T, [x, y, z, x_size, y_size, z_size, rz],
(x, y, z) is the bottom center. (x, y, z) is the bottom center.
Returns: Returns:
box_idxs_of_pts (torch.Tensor): (B, M, T), default background = 0 box_idxs_of_pts (torch.Tensor): (B, M, T), default background = 0.
""" """
assert boxes.shape[0] == points.shape[0], \ assert boxes.shape[0] == points.shape[0], \
f'Points and boxes should have the same batch size, ' \ f'Points and boxes should have the same batch size, ' \
...@@ -120,8 +121,8 @@ def points_in_boxes_batch(points, boxes): ...@@ -120,8 +121,8 @@ def points_in_boxes_batch(points, boxes):
if torch.cuda.current_device() != points_device: if torch.cuda.current_device() != points_device:
torch.cuda.set_device(points_device) torch.cuda.set_device(points_device)
roiaware_pool3d_ext.points_in_boxes_batch(boxes.contiguous(), roiaware_pool3d_ext.points_in_boxes_all(boxes.contiguous(),
points.contiguous(), points.contiguous(),
box_idxs_of_pts) box_idxs_of_pts)
return box_idxs_of_pts return box_idxs_of_pts
mmdet3d/ops/roiaware_pool3d/src/points_in_boxes_cpu.cpp
View file @ 53435c62
...@@ -23,23 +23,23 @@ inline void lidar_to_local_coords_cpu(float shift_x, float shift_y, float rz, ...@@ -23,23 +23,23 @@ inline void lidar_to_local_coords_cpu(float shift_x, float shift_y, float rz,
inline int check_pt_in_box3d_cpu(const float *pt, const float *box3d, inline int check_pt_in_box3d_cpu(const float *pt, const float *box3d,
float &local_x, float &local_y) { float &local_x, float &local_y) {
// param pt: (x, y, z) // param pt: (x, y, z)
// param box3d: (cx, cy, cz, w, l, h, rz) in LiDAR coordinate, cz in the // param box3d: (cx, cy, cz, x_size, y_size, z_size, rz) in LiDAR coordinate, cz in the
// bottom center // bottom center
float x = pt[0], y = pt[1], z = pt[2]; float x = pt[0], y = pt[1], z = pt[2];
float cx = box3d[0], cy = box3d[1], cz = box3d[2]; float cx = box3d[0], cy = box3d[1], cz = box3d[2];
float dx = box3d[3], dy = box3d[4], dz = box3d[5], rz = box3d[6]; float x_size = box3d[3], y_size = box3d[4], z_size = box3d[5], rz = box3d[6];
cz += dz / 2.0; // shift to the center since cz in box3d is the bottom center cz += z_size / 2.0; // shift to the center since cz in box3d is the bottom center
if (fabsf(z - cz) > dz / 2.0) return 0; if (fabsf(z - cz) > z_size / 2.0) return 0;
lidar_to_local_coords_cpu(x - cx, y - cy, rz, local_x, local_y); lidar_to_local_coords_cpu(x - cx, y - cy, rz, local_x, local_y);
float in_flag = (local_x > -dx / 2.0) & (local_x < dx / 2.0) & float in_flag = (local_x > -x_size / 2.0) & (local_x < x_size / 2.0) &
(local_y > -dy / 2.0) & (local_y < dy / 2.0); (local_y > -y_size / 2.0) & (local_y < y_size / 2.0);
return in_flag; return in_flag;
} }
int points_in_boxes_cpu(at::Tensor boxes_tensor, at::Tensor pts_tensor, int points_in_boxes_cpu(at::Tensor boxes_tensor, at::Tensor pts_tensor,
at::Tensor pts_indices_tensor) { at::Tensor pts_indices_tensor) {
// params boxes: (N, 7) [x, y, z, w, l, h, rz] in LiDAR coordinate, z is the // params boxes: (N, 7) [x, y, z, x_size, y_size, z_size, rz] in LiDAR coordinate, z is the
// bottom center, each box DO NOT overlaps params pts: (npoints, 3) [x, y, z] // bottom center, each box DO NOT overlaps params pts: (npoints, 3) [x, y, z]
// in LiDAR coordinate params pts_indices: (N, npoints) // in LiDAR coordinate params pts_indices: (N, npoints)
......
mmdet3d/ops/roiaware_pool3d/src/points_in_boxes_cuda.cu
View file @ 53435c62
...@@ -32,25 +32,25 @@ __device__ inline void lidar_to_local_coords(float shift_x, float shift_y, ...@@ -32,25 +32,25 @@ __device__ inline void lidar_to_local_coords(float shift_x, float shift_y,
__device__ inline int check_pt_in_box3d(const float *pt, const float *box3d, __device__ inline int check_pt_in_box3d(const float *pt, const float *box3d,
float &local_x, float &local_y) { float &local_x, float &local_y) {
// param pt: (x, y, z) // param pt: (x, y, z)
// param box3d: (cx, cy, cz, w, l, h, rz) in LiDAR coordinate, cz in the // param box3d: (cx, cy, cz, x_size, y_size, z_size, rz) in LiDAR coordinate, cz in the
// bottom center // bottom center
float x = pt[0], y = pt[1], z = pt[2]; float x = pt[0], y = pt[1], z = pt[2];
float cx = box3d[0], cy = box3d[1], cz = box3d[2]; float cx = box3d[0], cy = box3d[1], cz = box3d[2];
float dx = box3d[3], dy = box3d[4], dz = box3d[5], rz = box3d[6]; float x_size = box3d[3], y_size = box3d[4], z_size = box3d[5], rz = box3d[6];
cz += dz / 2.0; // shift to the center since cz in box3d is the bottom center cz += z_size / 2.0; // shift to the center since cz in box3d is the bottom center
if (fabsf(z - cz) > dz / 2.0) return 0; if (fabsf(z - cz) > z_size / 2.0) return 0;
lidar_to_local_coords(x - cx, y - cy, rz, local_x, local_y); lidar_to_local_coords(x - cx, y - cy, rz, local_x, local_y);
float in_flag = (local_x > -dx / 2.0) & (local_x < dx / 2.0) & float in_flag = (local_x > -x_size / 2.0) & (local_x < x_size / 2.0) &
(local_y > -dy / 2.0) & (local_y < dy / 2.0); (local_y > -y_size / 2.0) & (local_y < y_size / 2.0);
return in_flag; return in_flag;
} }
__global__ void points_in_boxes_kernel(int batch_size, int boxes_num, __global__ void points_in_boxes_part_kernel(int batch_size, int boxes_num,
int pts_num, const float *boxes, int pts_num, const float *boxes,
const float *pts, const float *pts,
int *box_idx_of_points) { int *box_idx_of_points) {
// params boxes: (B, N, 7) [x, y, z, dx, dy, dz, rz] in LiDAR coordinate, z is // params boxes: (B, N, 7) [x, y, z, x_size, y_size, z_size, rz] in LiDAR coordinate, z is
// the bottom center, each box DO NOT overlaps params pts: (B, npoints, 3) [x, // the bottom center, each box DO NOT overlaps params pts: (B, npoints, 3) [x,
// y, z] in LiDAR coordinate params boxes_idx_of_points: (B, npoints), default // y, z] in LiDAR coordinate params boxes_idx_of_points: (B, npoints), default
// -1 // -1
...@@ -74,11 +74,11 @@ __global__ void points_in_boxes_kernel(int batch_size, int boxes_num, ...@@ -74,11 +74,11 @@ __global__ void points_in_boxes_kernel(int batch_size, int boxes_num,
} }
} }
__global__ void points_in_boxes_batch_kernel(int batch_size, int boxes_num, __global__ void points_in_boxes_all_kernel(int batch_size, int boxes_num,
int pts_num, const float *boxes, int pts_num, const float *boxes,
const float *pts, const float *pts,
int *box_idx_of_points) { int *box_idx_of_points) {
// params boxes: (B, N, 7) [x, y, z, dx, dy, dz, rz] in LiDAR coordinate, z is // params boxes: (B, N, 7) [x, y, z, x_size, y_size, z_size, rz] in LiDAR coordinate, z is
// the bottom center, each box DO NOT overlaps params pts: (B, npoints, 3) [x, // the bottom center, each box DO NOT overlaps params pts: (B, npoints, 3) [x,
// y, z] in LiDAR coordinate params boxes_idx_of_points: (B, npoints), default // y, z] in LiDAR coordinate params boxes_idx_of_points: (B, npoints), default
// -1 // -1
...@@ -102,10 +102,10 @@ __global__ void points_in_boxes_batch_kernel(int batch_size, int boxes_num, ...@@ -102,10 +102,10 @@ __global__ void points_in_boxes_batch_kernel(int batch_size, int boxes_num,
} }
} }
void points_in_boxes_launcher(int batch_size, int boxes_num, int pts_num, void points_in_boxes_part_launcher(int batch_size, int boxes_num, int pts_num,
const float *boxes, const float *pts, const float *boxes, const float *pts,
int *box_idx_of_points) { int *box_idx_of_points) {
// params boxes: (B, N, 7) [x, y, z, dx, dy, dz, rz] in LiDAR coordinate, z is // params boxes: (B, N, 7) [x, y, z, x_size, y_size, z_size, rz] in LiDAR coordinate, z is
// the bottom center, each box DO NOT overlaps params pts: (B, npoints, 3) [x, // the bottom center, each box DO NOT overlaps params pts: (B, npoints, 3) [x,
// y, z] in LiDAR coordinate params boxes_idx_of_points: (B, npoints), default // y, z] in LiDAR coordinate params boxes_idx_of_points: (B, npoints), default
// -1 // -1
...@@ -113,8 +113,8 @@ void points_in_boxes_launcher(int batch_size, int boxes_num, int pts_num, ...@@ -113,8 +113,8 @@ void points_in_boxes_launcher(int batch_size, int boxes_num, int pts_num,
dim3 blocks(DIVUP(pts_num, THREADS_PER_BLOCK), batch_size); dim3 blocks(DIVUP(pts_num, THREADS_PER_BLOCK), batch_size);
dim3 threads(THREADS_PER_BLOCK); dim3 threads(THREADS_PER_BLOCK);
points_in_boxes_kernel<<<blocks, threads>>>(batch_size, boxes_num, pts_num, points_in_boxes_part_kernel<<<blocks, threads>>>(batch_size, boxes_num, pts_num,
boxes, pts, box_idx_of_points); boxes, pts, box_idx_of_points);
err = cudaGetLastError(); err = cudaGetLastError();
if (cudaSuccess != err) { if (cudaSuccess != err) {
...@@ -127,17 +127,17 @@ void points_in_boxes_launcher(int batch_size, int boxes_num, int pts_num, ...@@ -127,17 +127,17 @@ void points_in_boxes_launcher(int batch_size, int boxes_num, int pts_num,
#endif #endif
} }
void points_in_boxes_batch_launcher(int batch_size, int boxes_num, int pts_num, void points_in_boxes_all_launcher(int batch_size, int boxes_num, int pts_num,
const float *boxes, const float *pts, const float *boxes, const float *pts,
int *box_idx_of_points) { int *box_idx_of_points) {
// params boxes: (B, N, 7) [x, y, z, dx, dy, dz, rz] in LiDAR coordinate, z is // params boxes: (B, N, 7) [x, y, z, x_size, y_size, z_size, rz] in LiDAR coordinate, z is
// the bottom center, each box params pts: (B, npoints, 3) [x, y, z] in // the bottom center, each box params pts: (B, npoints, 3) [x, y, z] in
// LiDAR coordinate params boxes_idx_of_points: (B, npoints), default -1 // LiDAR coordinate params boxes_idx_of_points: (B, npoints), default -1
cudaError_t err; cudaError_t err;
dim3 blocks(DIVUP(pts_num, THREADS_PER_BLOCK), batch_size); dim3 blocks(DIVUP(pts_num, THREADS_PER_BLOCK), batch_size);
dim3 threads(THREADS_PER_BLOCK); dim3 threads(THREADS_PER_BLOCK);
points_in_boxes_batch_kernel<<<blocks, threads>>>( points_in_boxes_all_kernel<<<blocks, threads>>>(
batch_size, boxes_num, pts_num, boxes, pts, box_idx_of_points); batch_size, boxes_num, pts_num, boxes, pts, box_idx_of_points);
err = cudaGetLastError(); err = cudaGetLastError();
...@@ -151,9 +151,9 @@ void points_in_boxes_batch_launcher(int batch_size, int boxes_num, int pts_num, ...@@ -151,9 +151,9 @@ void points_in_boxes_batch_launcher(int batch_size, int boxes_num, int pts_num,
#endif #endif
} }
int points_in_boxes_gpu(at::Tensor boxes_tensor, at::Tensor pts_tensor, int points_in_boxes_part(at::Tensor boxes_tensor, at::Tensor pts_tensor,
at::Tensor box_idx_of_points_tensor) { at::Tensor box_idx_of_points_tensor) {
// params boxes: (B, N, 7) [x, y, z, dx, dy, dz, rz] in LiDAR coordinate, z is // params boxes: (B, N, 7) [x, y, z, x_size, y_size, z_size, rz] in LiDAR coordinate, z is
// the bottom center, each box DO NOT overlaps params pts: (B, npoints, 3) [x, // the bottom center, each box DO NOT overlaps params pts: (B, npoints, 3) [x,
// y, z] in LiDAR coordinate params boxes_idx_of_points: (B, npoints), default // y, z] in LiDAR coordinate params boxes_idx_of_points: (B, npoints), default
// -1 // -1
...@@ -170,15 +170,15 @@ int points_in_boxes_gpu(at::Tensor boxes_tensor, at::Tensor pts_tensor, ...@@ -170,15 +170,15 @@ int points_in_boxes_gpu(at::Tensor boxes_tensor, at::Tensor pts_tensor,
const float *pts = pts_tensor.data_ptr<float>(); const float *pts = pts_tensor.data_ptr<float>();
int *box_idx_of_points = box_idx_of_points_tensor.data_ptr<int>(); int *box_idx_of_points = box_idx_of_points_tensor.data_ptr<int>();
points_in_boxes_launcher(batch_size, boxes_num, pts_num, boxes, pts, points_in_boxes_part_launcher(batch_size, boxes_num, pts_num, boxes, pts,
box_idx_of_points); box_idx_of_points);
return 1; return 1;
} }
int points_in_boxes_batch(at::Tensor boxes_tensor, at::Tensor pts_tensor, int points_in_boxes_all(at::Tensor boxes_tensor, at::Tensor pts_tensor,
at::Tensor box_idx_of_points_tensor) { at::Tensor box_idx_of_points_tensor) {
// params boxes: (B, N, 7) [x, y, z, dx, dy, dz, rz] in LiDAR coordinate, z is // params boxes: (B, N, 7) [x, y, z, x_size, y_size, z_size, rz] in LiDAR coordinate, z is
// the bottom center. params pts: (B, npoints, 3) [x, y, z] in LiDAR // the bottom center. params pts: (B, npoints, 3) [x, y, z] in LiDAR
// coordinate params boxes_idx_of_points: (B, npoints), default -1 // coordinate params boxes_idx_of_points: (B, npoints), default -1
...@@ -194,8 +194,8 @@ int points_in_boxes_batch(at::Tensor boxes_tensor, at::Tensor pts_tensor, ...@@ -194,8 +194,8 @@ int points_in_boxes_batch(at::Tensor boxes_tensor, at::Tensor pts_tensor,
const float *pts = pts_tensor.data_ptr<float>(); const float *pts = pts_tensor.data_ptr<float>();
int *box_idx_of_points = box_idx_of_points_tensor.data_ptr<int>(); int *box_idx_of_points = box_idx_of_points_tensor.data_ptr<int>();
points_in_boxes_batch_launcher(batch_size, boxes_num, pts_num, boxes, pts, points_in_boxes_all_launcher(batch_size, boxes_num, pts_num, boxes, pts,
box_idx_of_points); box_idx_of_points);
return 1; return 1;
} }
mmdet3d/ops/roiaware_pool3d/src/roiaware_pool3d.cpp
View file @ 53435c62
...@@ -40,16 +40,16 @@ int roiaware_pool3d_gpu_backward(at::Tensor pts_idx_of_voxels, ...@@ -40,16 +40,16 @@ int roiaware_pool3d_gpu_backward(at::Tensor pts_idx_of_voxels,
int points_in_boxes_cpu(at::Tensor boxes_tensor, at::Tensor pts_tensor, int points_in_boxes_cpu(at::Tensor boxes_tensor, at::Tensor pts_tensor,
at::Tensor pts_indices_tensor); at::Tensor pts_indices_tensor);
int points_in_boxes_gpu(at::Tensor boxes_tensor, at::Tensor pts_tensor, int points_in_boxes_part(at::Tensor boxes_tensor, at::Tensor pts_tensor,
at::Tensor box_idx_of_points_tensor); at::Tensor box_idx_of_points_tensor);
int points_in_boxes_batch(at::Tensor boxes_tensor, at::Tensor pts_tensor, int points_in_boxes_all(at::Tensor boxes_tensor, at::Tensor pts_tensor,
at::Tensor box_idx_of_points_tensor); at::Tensor box_idx_of_points_tensor);
int roiaware_pool3d_gpu(at::Tensor rois, at::Tensor pts, at::Tensor pts_feature, int roiaware_pool3d_gpu(at::Tensor rois, at::Tensor pts, at::Tensor pts_feature,
at::Tensor argmax, at::Tensor pts_idx_of_voxels, at::Tensor argmax, at::Tensor pts_idx_of_voxels,
at::Tensor pooled_features, int pool_method) { at::Tensor pooled_features, int pool_method) {
// params rois: (N, 7) [x, y, z, w, l, h, ry] in LiDAR coordinate // params rois: (N, 7) [x, y, z, x_size, y_size, z_size, ry] in LiDAR coordinate
// params pts: (npoints, 3) [x, y, z] in LiDAR coordinate // params pts: (npoints, 3) [x, y, z] in LiDAR coordinate
// params pts_feature: (npoints, C) // params pts_feature: (npoints, C)
// params argmax: (N, out_x, out_y, out_z, C) // params argmax: (N, out_x, out_y, out_z, C)
...@@ -127,10 +127,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { ...@@ -127,10 +127,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &roiaware_pool3d_gpu, "roiaware pool3d forward (CUDA)"); m.def("forward", &roiaware_pool3d_gpu, "roiaware pool3d forward (CUDA)");
m.def("backward", &roiaware_pool3d_gpu_backward, m.def("backward", &roiaware_pool3d_gpu_backward,
"roiaware pool3d backward (CUDA)"); "roiaware pool3d backward (CUDA)");
m.def("points_in_boxes_gpu", &points_in_boxes_gpu, m.def("points_in_boxes_part", &points_in_boxes_part,
"points_in_boxes_gpu forward (CUDA)"); "points_in_boxes_part forward (CUDA)");
m.def("points_in_boxes_batch", &points_in_boxes_batch, m.def("points_in_boxes_all", &points_in_boxes_all,
"points_in_boxes_batch forward (CUDA)"); "points_in_boxes_all forward (CUDA)");
m.def("points_in_boxes_cpu", &points_in_boxes_cpu, m.def("points_in_boxes_cpu", &points_in_boxes_cpu,
"points_in_boxes_cpu forward (CPU)"); "points_in_boxes_cpu forward (CPU)");
} }
mmdet3d/ops/roiaware_pool3d/src/roiaware_pool3d_kernel.cu
View file @ 53435c62
...@@ -25,17 +25,17 @@ __device__ inline void lidar_to_local_coords(float shift_x, float shift_y, ...@@ -25,17 +25,17 @@ __device__ inline void lidar_to_local_coords(float shift_x, float shift_y,
__device__ inline int check_pt_in_box3d(const float *pt, const float *box3d, __device__ inline int check_pt_in_box3d(const float *pt, const float *box3d,
float &local_x, float &local_y) { float &local_x, float &local_y) {
// param pt: (x, y, z) // param pt: (x, y, z)
// param box3d: (cx, cy, cz, dx, dy, dz, rz) in LiDAR coordinate, cz in the // param box3d: (cx, cy, cz, x_size, y_size, z_size, rz) in LiDAR coordinate, cz in the
// bottom center // bottom center
float x = pt[0], y = pt[1], z = pt[2]; float x = pt[0], y = pt[1], z = pt[2];
float cx = box3d[0], cy = box3d[1], cz = box3d[2]; float cx = box3d[0], cy = box3d[1], cz = box3d[2];
float dx = box3d[3], dy = box3d[4], dz = box3d[5], rz = box3d[6]; float x_size = box3d[3], y_size = box3d[4], z_size = box3d[5], rz = box3d[6];
cz += dz / 2.0; // shift to the center since cz in box3d is the bottom center cz += z_size / 2.0; // shift to the center since cz in box3d is the bottom center
if (fabsf(z - cz) > dz / 2.0) return 0; if (fabsf(z - cz) > z_size / 2.0) return 0;
lidar_to_local_coords(x - cx, y - cy, rz, local_x, local_y); lidar_to_local_coords(x - cx, y - cy, rz, local_x, local_y);
float in_flag = (local_x > -dx / 2.0) & (local_x < dx / 2.0) & float in_flag = (local_x > -x_size / 2.0) & (local_x < x_size / 2.0) &
(local_y > -dy / 2.0) & (local_y < dy / 2.0); (local_y > -y_size / 2.0) & (local_y < y_size / 2.0);
return in_flag; return in_flag;
} }
...@@ -43,7 +43,7 @@ __global__ void generate_pts_mask_for_box3d(int boxes_num, int pts_num, ...@@ -43,7 +43,7 @@ __global__ void generate_pts_mask_for_box3d(int boxes_num, int pts_num,
int out_x, int out_y, int out_z, int out_x, int out_y, int out_z,
const float *rois, const float *pts, const float *rois, const float *pts,
int *pts_mask) { int *pts_mask) {
// params rois: (N, 7) [x, y, z, dx, dy, dz, rz] in LiDAR coordinate // params rois: (N, 7) [x, y, z, x_size, y_size, z_size, rz] in LiDAR coordinate
// params pts: (npoints, 3) [x, y, z] // params pts: (npoints, 3) [x, y, z]
// params pts_mask: (N, npoints): -1 means point does not in this box, // params pts_mask: (N, npoints): -1 means point does not in this box,
// otherwise: encode (x_idxs, y_idxs, z_idxs) by binary bit // otherwise: encode (x_idxs, y_idxs, z_idxs) by binary bit
...@@ -61,14 +61,14 @@ __global__ void generate_pts_mask_for_box3d(int boxes_num, int pts_num, ...@@ -61,14 +61,14 @@ __global__ void generate_pts_mask_for_box3d(int boxes_num, int pts_num,
pts_mask[0] = -1; pts_mask[0] = -1;
if (cur_in_flag > 0) { if (cur_in_flag > 0) {
float local_z = pts[2] - rois[2]; float local_z = pts[2] - rois[2];
float dx = rois[3], dy = rois[4], dz = rois[5]; float x_size = rois[3], y_size = rois[4], z_size = rois[5];
float x_res = dx / out_x; float x_res = x_size / out_x;
float y_res = dy / out_y; float y_res = y_size / out_y;
float z_res = dz / out_z; float z_res = z_size / out_z;
unsigned int x_idx = int((local_x + dx / 2) / x_res); unsigned int x_idx = int((local_x + x_size / 2) / x_res);
unsigned int y_idx = int((local_y + dy / 2) / y_res); unsigned int y_idx = int((local_y + y_size / 2) / y_res);
unsigned int z_idx = int(local_z / z_res); unsigned int z_idx = int(local_z / z_res);
x_idx = min(max(x_idx, 0), out_x - 1); x_idx = min(max(x_idx, 0), out_x - 1);
...@@ -229,7 +229,7 @@ void roiaware_pool3d_launcher(int boxes_num, int pts_num, int channels, ...@@ -229,7 +229,7 @@ void roiaware_pool3d_launcher(int boxes_num, int pts_num, int channels,
const float *pts_feature, int *argmax, const float *pts_feature, int *argmax,
int *pts_idx_of_voxels, float *pooled_features, int *pts_idx_of_voxels, float *pooled_features,
int pool_method) { int pool_method) {
// params rois: (N, 7) [x, y, z, dx, dy, dz, rz] in LiDAR coordinate // params rois: (N, 7) [x, y, z, x_size, y_size, z_size, rz] in LiDAR coordinate
// params pts: (npoints, 3) [x, y, z] in LiDAR coordinate // params pts: (npoints, 3) [x, y, z] in LiDAR coordinate
// params pts_feature: (npoints, C) // params pts_feature: (npoints, C)
// params argmax: (N, out_x, out_y, out_z, C) // params argmax: (N, out_x, out_y, out_z, C)
......
mmdet3d/ops/sparse_block.py
View file @ 53435c62
...@@ -14,12 +14,12 @@ class SparseBottleneck(Bottleneck, spconv.SparseModule): ...@@ -14,12 +14,12 @@ class SparseBottleneck(Bottleneck, spconv.SparseModule):
Args: Args:
inplanes (int): inplanes of block. inplanes (int): inplanes of block.
planes (int): planes of block. planes (int): planes of block.
stride (int): stride of the first block. Default: 1 stride (int, optional): stride of the first block. Default: 1.
downsample (None | Module): down sample module for block. downsample (Module, optional): down sample module for block.
conv_cfg (dict): dictionary to construct and config conv layer. conv_cfg (dict, optional): dictionary to construct and config conv
Default: None layer. Default: None.
norm_cfg (dict): dictionary to construct and config norm layer. norm_cfg (dict, optional): dictionary to construct and config norm
Default: dict(type='BN') layer. Default: dict(type='BN').
""" """
expansion = 4 expansion = 4
...@@ -73,12 +73,12 @@ class SparseBasicBlock(BasicBlock, spconv.SparseModule): ...@@ -73,12 +73,12 @@ class SparseBasicBlock(BasicBlock, spconv.SparseModule):
Args: Args:
inplanes (int): inplanes of block. inplanes (int): inplanes of block.
planes (int): planes of block. planes (int): planes of block.
stride (int): stride of the first block. Default: 1 stride (int, optional): stride of the first block. Default: 1.
downsample (None | Module): down sample module for block. downsample (Module, optional): down sample module for block.
conv_cfg (dict): dictionary to construct and config conv layer. conv_cfg (dict, optional): dictionary to construct and config conv
Default: None layer. Default: None.
norm_cfg (dict): dictionary to construct and config norm layer. norm_cfg (dict, optional): dictionary to construct and config norm
Default: dict(type='BN') layer. Default: dict(type='BN').
""" """
expansion = 1 expansion = 1
......
tests/test_models/test_common_modules/test_roiaware_pool3d.py
View file @ 53435c62
...@@ -3,9 +3,9 @@ import numpy as np ...@@ -3,9 +3,9 @@ import numpy as np
import pytest import pytest
import torch import torch
from mmdet3d.ops.roiaware_pool3d import (RoIAwarePool3d, points_in_boxes_batch, from mmdet3d.ops.roiaware_pool3d import (RoIAwarePool3d, points_in_boxes_all,
points_in_boxes_cpu, points_in_boxes_cpu,
points_in_boxes_gpu) points_in_boxes_part)
def test_RoIAwarePool3d(): def test_RoIAwarePool3d():
...@@ -42,7 +42,7 @@ def test_RoIAwarePool3d(): ...@@ -42,7 +42,7 @@ def test_RoIAwarePool3d():
torch.tensor(49.750).cuda(), 1e-3) torch.tensor(49.750).cuda(), 1e-3)
def test_points_in_boxes_gpu(): def test_points_in_boxes_part():
if not torch.cuda.is_available(): if not torch.cuda.is_available():
pytest.skip('test requires GPU and torch+cuda') pytest.skip('test requires GPU and torch+cuda')
boxes = torch.tensor( boxes = torch.tensor(
...@@ -58,7 +58,7 @@ def test_points_in_boxes_gpu(): ...@@ -58,7 +58,7 @@ def test_points_in_boxes_gpu():
[0, 0, 0], [6, 7, 8], [-2, -3, -4], [6, 4, 9]]], [0, 0, 0], [6, 7, 8], [-2, -3, -4], [6, 4, 9]]],
dtype=torch.float32).cuda() # points (b, m, 3) in lidar coordinate dtype=torch.float32).cuda() # points (b, m, 3) in lidar coordinate
point_indices = points_in_boxes_gpu(points=pts, boxes=boxes) point_indices = points_in_boxes_part(points=pts, boxes=boxes)
expected_point_indices = torch.tensor( expected_point_indices = torch.tensor(
[[0, 0, 0, 0, 0, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1, -1]], [[0, 0, 0, 0, 0, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1, -1]],
dtype=torch.int32).cuda() dtype=torch.int32).cuda()
...@@ -71,7 +71,7 @@ def test_points_in_boxes_gpu(): ...@@ -71,7 +71,7 @@ def test_points_in_boxes_gpu():
[[[4, 6.928, 0], [6.928, 4, 0], [4, -6.928, 0], [6.928, -4, 0], [[[4, 6.928, 0], [6.928, 4, 0], [4, -6.928, 0], [6.928, -4, 0],
[-4, 6.928, 0], [-6.928, 4, 0], [-4, -6.928, 0], [-6.928, -4, 0]]], [-4, 6.928, 0], [-6.928, 4, 0], [-4, -6.928, 0], [-6.928, -4, 0]]],
dtype=torch.float32).cuda() dtype=torch.float32).cuda()
point_indices = points_in_boxes_gpu(points=pts, boxes=boxes) point_indices = points_in_boxes_part(points=pts, boxes=boxes)
expected_point_indices = torch.tensor([[-1, -1, 0, -1, 0, -1, -1, -1]], expected_point_indices = torch.tensor([[-1, -1, 0, -1, 0, -1, -1, -1]],
dtype=torch.int32).cuda() dtype=torch.int32).cuda()
assert (point_indices == expected_point_indices).all() assert (point_indices == expected_point_indices).all()
...@@ -80,7 +80,7 @@ def test_points_in_boxes_gpu(): ...@@ -80,7 +80,7 @@ def test_points_in_boxes_gpu():
pts = pts.to('cuda:1') pts = pts.to('cuda:1')
boxes = boxes.to('cuda:1') boxes = boxes.to('cuda:1')
expected_point_indices = expected_point_indices.to('cuda:1') expected_point_indices = expected_point_indices.to('cuda:1')
point_indices = points_in_boxes_gpu(points=pts, boxes=boxes) point_indices = points_in_boxes_part(points=pts, boxes=boxes)
assert point_indices.shape == torch.Size([2, 8]) assert point_indices.shape == torch.Size([2, 8])
assert (point_indices == expected_point_indices).all() assert (point_indices == expected_point_indices).all()
...@@ -119,7 +119,7 @@ def test_points_in_boxes_cpu(): ...@@ -119,7 +119,7 @@ def test_points_in_boxes_cpu():
assert (point_indices == expected_point_indices).all() assert (point_indices == expected_point_indices).all()
def test_points_in_boxes_batch(): def test_points_in_boxes_all():
if not torch.cuda.is_available(): if not torch.cuda.is_available():
pytest.skip('test requires GPU and torch+cuda') pytest.skip('test requires GPU and torch+cuda')
...@@ -136,7 +136,7 @@ def test_points_in_boxes_batch(): ...@@ -136,7 +136,7 @@ def test_points_in_boxes_batch():
], [-21.3, -52, -5], [0, 0, 0], [6, 7, 8], [-2, -3, -4]]], ], [-21.3, -52, -5], [0, 0, 0], [6, 7, 8], [-2, -3, -4]]],
dtype=torch.float32).cuda() # points (n, 3) in lidar coordinate dtype=torch.float32).cuda() # points (n, 3) in lidar coordinate
point_indices = points_in_boxes_batch(points=pts, boxes=boxes) point_indices = points_in_boxes_all(points=pts, boxes=boxes)
expected_point_indices = torch.tensor( expected_point_indices = torch.tensor(
[[[1, 0], [1, 0], [1, 0], [1, 0], [1, 0], [0, 1], [0, 0], [0, 0], [[[1, 0], [1, 0], [1, 0], [1, 0], [1, 0], [0, 1], [0, 0], [0, 0],
[0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0]]], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0]]],
...@@ -148,6 +148,6 @@ def test_points_in_boxes_batch(): ...@@ -148,6 +148,6 @@ def test_points_in_boxes_batch():
pts = pts.to('cuda:1') pts = pts.to('cuda:1')
boxes = boxes.to('cuda:1') boxes = boxes.to('cuda:1')
expected_point_indices = expected_point_indices.to('cuda:1') expected_point_indices = expected_point_indices.to('cuda:1')
point_indices = points_in_boxes_batch(points=pts, boxes=boxes) point_indices = points_in_boxes_all(points=pts, boxes=boxes)
assert point_indices.shape == torch.Size([1, 15, 2]) assert point_indices.shape == torch.Size([1, 15, 2])
assert (point_indices == expected_point_indices).all() assert (point_indices == expected_point_indices).all()
tests/test_models/test_forward.py
View file @ 53435c62
...@@ -148,7 +148,7 @@ def _demo_mm_inputs(input_shape=(1, 3, 300, 300), ...@@ -148,7 +148,7 @@ def _demo_mm_inputs(input_shape=(1, 3, 300, 300),
input_shape (tuple): input_shape (tuple):
input batch dimensions input batch dimensions
num_items (None | List[int]): num_items (List[int]):
specifies the number of boxes in each batch item specifies the number of boxes in each batch item
num_classes (int): num_classes (int):
......
tests/test_models/test_heads/test_heads.py
View file @ 53435c62
...@@ -1144,7 +1144,7 @@ def test_groupfree3d_head(): ...@@ -1144,7 +1144,7 @@ def test_groupfree3d_head():
assert ret_dict['s5.sem_scores'].shape == torch.Size([2, 256, 18]) assert ret_dict['s5.sem_scores'].shape == torch.Size([2, 256, 18])
# test losses # test losses
points = [torch.rand([50000, 4], device='cuda') for i in range(2)] points = [torch.rand([5000, 4], device='cuda') for i in range(2)]
gt_bbox1 = torch.rand([10, 7], dtype=torch.float32).cuda() gt_bbox1 = torch.rand([10, 7], dtype=torch.float32).cuda()
gt_bbox2 = torch.rand([10, 7], dtype=torch.float32).cuda() gt_bbox2 = torch.rand([10, 7], dtype=torch.float32).cuda()
...@@ -1152,12 +1152,12 @@ def test_groupfree3d_head(): ...@@ -1152,12 +1152,12 @@ def test_groupfree3d_head():
gt_bbox2 = DepthInstance3DBoxes(gt_bbox2) gt_bbox2 = DepthInstance3DBoxes(gt_bbox2)
gt_bboxes = [gt_bbox1, gt_bbox2] gt_bboxes = [gt_bbox1, gt_bbox2]
pts_instance_mask_1 = torch.randint(0, 10, [50000], device='cuda') pts_instance_mask_1 = torch.randint(0, 10, [5000], device='cuda')
pts_instance_mask_2 = torch.randint(0, 10, [50000], device='cuda') pts_instance_mask_2 = torch.randint(0, 10, [5000], device='cuda')
pts_instance_mask = [pts_instance_mask_1, pts_instance_mask_2] pts_instance_mask = [pts_instance_mask_1, pts_instance_mask_2]
pts_semantic_mask_1 = torch.randint(0, 19, [50000], device='cuda') pts_semantic_mask_1 = torch.randint(0, 19, [5000], device='cuda')
pts_semantic_mask_2 = torch.randint(0, 19, [50000], device='cuda') pts_semantic_mask_2 = torch.randint(0, 19, [5000], device='cuda')
pts_semantic_mask = [pts_semantic_mask_1, pts_semantic_mask_2] pts_semantic_mask = [pts_semantic_mask_1, pts_semantic_mask_2]
labels_1 = torch.randint(0, 18, [10], device='cuda') labels_1 = torch.randint(0, 18, [10], device='cuda')
...@@ -1178,7 +1178,7 @@ def test_groupfree3d_head(): ...@@ -1178,7 +1178,7 @@ def test_groupfree3d_head():
# test multiclass_nms_single # test multiclass_nms_single
obj_scores = torch.rand([256], device='cuda') obj_scores = torch.rand([256], device='cuda')
sem_scores = torch.rand([256, 18], device='cuda') sem_scores = torch.rand([256, 18], device='cuda')
points = torch.rand([50000, 3], device='cuda') points = torch.rand([5000, 3], device='cuda')
bbox = torch.rand([256, 7], device='cuda') bbox = torch.rand([256, 7], device='cuda')
input_meta = dict(box_type_3d=DepthInstance3DBoxes) input_meta = dict(box_type_3d=DepthInstance3DBoxes)
bbox_selected, score_selected, labels = \ bbox_selected, score_selected, labels = \
...@@ -1193,9 +1193,9 @@ def test_groupfree3d_head(): ...@@ -1193,9 +1193,9 @@ def test_groupfree3d_head():
assert labels.shape[0] >= 0 assert labels.shape[0] >= 0
# test get_boxes # test get_boxes
points = torch.rand([1, 50000, 3], device='cuda') points = torch.rand([1, 5000, 3], device='cuda')
seed_points = torch.rand([1, 1024, 3], device='cuda') seed_points = torch.rand([1, 1024, 3], device='cuda')
seed_indices = torch.randint(0, 50000, [1, 1024], device='cuda') seed_indices = torch.randint(0, 5000, [1, 1024], device='cuda')
obj_scores = torch.rand([1, 256, 1], device='cuda') obj_scores = torch.rand([1, 256, 1], device='cuda')
center = torch.rand([1, 256, 3], device='cuda') center = torch.rand([1, 256, 3], device='cuda')
dir_class = torch.rand([1, 256, 1], device='cuda') dir_class = torch.rand([1, 256, 1], device='cuda')
......
tests/test_utils/test_box3d.py
View file @ 53435c62
...@@ -5,9 +5,9 @@ import torch ...@@ -5,9 +5,9 @@ import torch
import unittest import unittest
from mmdet3d.core.bbox import (BaseInstance3DBoxes, Box3DMode, from mmdet3d.core.bbox import (BaseInstance3DBoxes, Box3DMode,
CameraInstance3DBoxes, DepthInstance3DBoxes, CameraInstance3DBoxes, Coord3DMode,
LiDARInstance3DBoxes, bbox3d2roi, DepthInstance3DBoxes, LiDARInstance3DBoxes,
bbox3d_mapping_back) bbox3d2roi, bbox3d_mapping_back)
from mmdet3d.core.bbox.structures.utils import (get_box_type, limit_period, from mmdet3d.core.bbox.structures.utils import (get_box_type, limit_period,
points_cam2img, points_cam2img,
rotation_3d_in_axis, rotation_3d_in_axis,
...@@ -409,6 +409,13 @@ def test_lidar_boxes3d(): ...@@ -409,6 +409,13 @@ def test_lidar_boxes3d():
assert torch.allclose(boxes.tensor, expected_tensor) assert torch.allclose(boxes.tensor, expected_tensor)
# test bbox in_range_bev # test bbox in_range_bev
expected_tensor = torch.tensor(
[[1.1282, -3.0508, 1.7598, 3.4090, -1.2079],
[8.0981, -4.9332, 1.5486, 4.0325, -1.3479],
[27.6424, -7.2409, 1.4782, 2.2425, 1.8421],
[20.0183, -28.4773, 1.5687, 3.4995, 1.9621],
[28.2147, -16.5020, 1.7497, 3.7911, -2.5179]])
assert torch.allclose(boxes.bev, expected_tensor, atol=1e-3)
expected_tensor = torch.tensor([1, 1, 1, 1, 1], dtype=torch.bool) expected_tensor = torch.tensor([1, 1, 1, 1, 1], dtype=torch.bool)
mask = boxes.in_range_bev([0., -40., 70.4, 40.]) mask = boxes.in_range_bev([0., -40., 70.4, 40.])
assert (mask == expected_tensor).all() assert (mask == expected_tensor).all()
...@@ -1000,6 +1007,14 @@ def test_camera_boxes3d(): ...@@ -1000,6 +1007,14 @@ def test_camera_boxes3d():
mask = boxes.in_range_3d([-2, -5, 0, 20, 2, 22]) mask = boxes.in_range_3d([-2, -5, 0, 20, 2, 22])
assert (mask == expected_tensor).all() assert (mask == expected_tensor).all()
expected_tensor = torch.tensor(
[[3.0508, 1.1282, 1.7598, 3.4090, -5.9203],
[4.9332, 8.0981, 1.5486, 4.0325, -6.0603],
[7.2409, 27.6424, 1.4782, 2.2425, -2.8703],
[28.4773, 20.0183, 1.5687, 3.4995, -2.7503],
[16.5020, 28.2147, 1.7497, 3.7911, -0.9471]])
assert torch.allclose(boxes.bev, expected_tensor, atol=1e-3)
# test properties # test properties
assert torch.allclose(boxes.bottom_center, boxes.tensor[:, :3]) assert torch.allclose(boxes.bottom_center, boxes.tensor[:, :3])
expected_tensor = ( expected_tensor = (
...@@ -1389,6 +1404,11 @@ def test_depth_boxes3d(): ...@@ -1389,6 +1404,11 @@ def test_depth_boxes3d():
mask = boxes.nonempty() mask = boxes.nonempty()
assert (mask == expected_tensor).all() assert (mask == expected_tensor).all()
# test bbox in_range
expected_tensor = torch.tensor([0, 1], dtype=torch.bool)
mask = boxes.in_range_3d([1, 0, -2, 2, 1, 5])
assert (mask == expected_tensor).all()
expected_tensor = torch.tensor([[[-0.1030, 0.6649, 0.1056], expected_tensor = torch.tensor([[[-0.1030, 0.6649, 0.1056],
[-0.1030, 0.6649, 0.3852], [-0.1030, 0.6649, 0.3852],
[-0.1030, 0.9029, 0.3852], [-0.1030, 0.9029, 0.3852],
...@@ -1409,7 +1429,7 @@ def test_depth_boxes3d(): ...@@ -1409,7 +1429,7 @@ def test_depth_boxes3d():
# test points in boxes # test points in boxes
if torch.cuda.is_available(): if torch.cuda.is_available():
box_idxs_of_pts = boxes.points_in_boxes_batch(points.cuda()) box_idxs_of_pts = boxes.points_in_boxes_all(points.cuda())
expected_idxs_of_pts = torch.tensor( expected_idxs_of_pts = torch.tensor(
[[0, 0], [0, 0], [0, 0], [0, 0], [0, 0]], [[0, 0], [0, 0], [0, 0], [0, 0], [0, 0]],
device='cuda:0', device='cuda:0',
...@@ -1467,23 +1487,25 @@ def test_depth_boxes3d(): ...@@ -1467,23 +1487,25 @@ def test_depth_boxes3d():
def test_rotation_3d_in_axis(): def test_rotation_3d_in_axis():
# # clockwise # clockwise
# points = torch.tensor([[[-0.4599, -0.0471, 0.0000], points = torch.tensor([[[-0.4599, -0.0471, 0.0000],
# [-0.4599, -0.0471, 1.8433], [-0.4599, -0.0471, 1.8433],
# [-0.4599, 0.0471, 1.8433]], [-0.4599, 0.0471, 1.8433]],
# [[-0.2555, -0.2683, 0.0000], [[-0.2555, -0.2683, 0.0000],
# [-0.2555, -0.2683, 0.9072], [-0.2555, -0.2683, 0.9072],
# [-0.2555, 0.2683, 0.9072]]]) [-0.2555, 0.2683, 0.9072]]])
# rotated = rotation_3d_in_axis( rotated = rotation_3d_in_axis(
# points, torch.tensor([-np.pi / 10, np.pi / 10]), points,
# axis=0, clockwise=True) torch.tensor([-np.pi / 10, np.pi / 10]),
# expected_rotated = torch.tensor([[[0.0000, -0.4228, -0.1869], axis=0,
# [1.8433, -0.4228, -0.1869], clockwise=True)
# [1.8433, -0.4519, -0.0973]], expected_rotated = torch.tensor(
# [[0.0000, -0.3259, -0.1762], [[[-0.4599, -0.0448, -0.0146], [-0.4599, -0.6144, 1.7385],
# [0.9072, -0.3259, -0.1762], [-0.4599, -0.5248, 1.7676]],
# [0.9072, -0.1601, 0.3341]]]) [[-0.2555, -0.2552, 0.0829], [-0.2555, 0.0252, 0.9457],
# assert torch.allclose(rotated, expected_rotated, 1e-3) [-0.2555, 0.5355, 0.7799]]],
dtype=torch.float32)
assert torch.allclose(rotated, expected_rotated, atol=1e-3)
# anti-clockwise with return rotation mat # anti-clockwise with return rotation mat
points = torch.tensor([[[-0.4599, -0.0471, 0.0000], points = torch.tensor([[[-0.4599, -0.0471, 0.0000],
...@@ -1622,3 +1644,128 @@ def test_points_cam2img(): ...@@ -1622,3 +1644,128 @@ def test_points_cam2img():
point_2d_res = points_cam2img(points, proj_mat) point_2d_res = points_cam2img(points, proj_mat)
expected_point_2d_res = torch.from_numpy(expected_point_2d_res) expected_point_2d_res = torch.from_numpy(expected_point_2d_res)
assert torch.allclose(point_2d_res, expected_point_2d_res, 1e-3) assert torch.allclose(point_2d_res, expected_point_2d_res, 1e-3)
point_2d_res = points_cam2img(points, proj_mat, with_depth=True)
expected_point_2d_res = torch.tensor([[0.5832, 0.6496, 1.7577],
[0.6146, 0.7910, 1.5477],
[0.6994, 0.7782, 2.0091],
[0.5623, 0.6303, 1.8739],
[0.4359, 0.6532, 1.2056]])
assert torch.allclose(point_2d_res, expected_point_2d_res, 1e-3)
def test_points_in_boxes():
if not torch.cuda.is_available():
pytest.skip('test requires GPU and torch+cuda')
lidar_pts = torch.tensor([[1.0, 4.3, 0.1], [1.0, 4.4,
0.1], [1.1, 4.3, 0.1],
[0.9, 4.3, 0.1], [1.0, -0.3, 0.1],
[1.0, -0.4, 0.1], [2.9, 0.1, 6.0],
[-0.9, 3.9, 6.0]]).cuda()
lidar_boxes = torch.tensor([[1.0, 2.0, 0.0, 4.0, 4.0, 6.0, np.pi / 6],
[1.0, 2.0, 0.0, 4.0, 4.0, 6.0, np.pi / 2],
[1.0, 2.0, 0.0, 4.0, 4.0, 6.0, 7 * np.pi / 6],
[1.0, 2.0, 0.0, 4.0, 4.0, 6.0, -np.pi / 6]],
dtype=torch.float32).cuda()
lidar_boxes = LiDARInstance3DBoxes(lidar_boxes)
point_indices = lidar_boxes.points_in_boxes_all(lidar_pts)
expected_point_indices = torch.tensor(
[[1, 0, 1, 1], [0, 0, 0, 0], [1, 0, 1, 0], [0, 0, 0, 1], [1, 0, 1, 1],
[0, 0, 0, 0], [0, 1, 0, 0], [0, 1, 0, 0]],
dtype=torch.int32).cuda()
assert point_indices.shape == torch.Size([8, 4])
assert (point_indices == expected_point_indices).all()
lidar_pts = torch.tensor([[1.0, 4.3, 0.1], [1.0, 4.4,
0.1], [1.1, 4.3, 0.1],
[0.9, 4.3, 0.1], [1.0, -0.3, 0.1],
[1.0, -0.4, 0.1], [2.9, 0.1, 6.0],
[-0.9, 3.9, 6.0]]).cuda()
lidar_boxes = torch.tensor([[1.0, 2.0, 0.0, 4.0, 4.0, 6.0, np.pi / 6],
[1.0, 2.0, 0.0, 4.0, 4.0, 6.0, np.pi / 2],
[1.0, 2.0, 0.0, 4.0, 4.0, 6.0, 7 * np.pi / 6],
[1.0, 2.0, 0.0, 4.0, 4.0, 6.0, -np.pi / 6]],
dtype=torch.float32).cuda()
lidar_boxes = LiDARInstance3DBoxes(lidar_boxes)
point_indices = lidar_boxes.points_in_boxes_part(lidar_pts)
expected_point_indices = torch.tensor([0, -1, 0, 3, 0, -1, 1, 1],
dtype=torch.int32).cuda()
assert point_indices.shape == torch.Size([8])
assert (point_indices == expected_point_indices).all()
depth_boxes = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 0.3],
[-10.0, 23.0, 16.0, 10, 20, 20, 0.5]],
dtype=torch.float32).cuda()
depth_boxes = DepthInstance3DBoxes(depth_boxes)
depth_pts = torch.tensor(
[[[1, 2, 3.3], [1.2, 2.5, 3.0], [0.8, 2.1, 3.5], [1.6, 2.6, 3.6],
[0.8, 1.2, 3.9], [-9.2, 21.0, 18.2], [3.8, 7.9, 6.3],
[4.7, 3.5, -12.2], [3.8, 7.6, -2], [-10.6, -12.9, -20], [
-16, -18, 9
], [-21.3, -52, -5], [0, 0, 0], [6, 7, 8], [-2, -3, -4]]],
dtype=torch.float32).cuda()
point_indices = depth_boxes.points_in_boxes_all(depth_pts)
expected_point_indices = torch.tensor(
[[1, 0], [1, 0], [1, 0], [1, 0], [1, 0], [0, 1], [0, 0], [0, 0],
[0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0]],
dtype=torch.int32).cuda()
assert point_indices.shape == torch.Size([15, 2])
assert (point_indices == expected_point_indices).all()
point_indices = depth_boxes.points_in_boxes_part(depth_pts)
expected_point_indices = torch.tensor(
[0, 0, 0, 0, 0, 1, -1, -1, -1, -1, -1, -1, -1, -1, -1],
dtype=torch.int32).cuda()
assert point_indices.shape == torch.Size([15])
assert (point_indices == expected_point_indices).all()
depth_boxes = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 0.3],
[-10.0, 23.0, 16.0, 10, 20, 20, 0.5],
[1.0, 2.0, 0.0, 4.0, 4.0, 6.0, np.pi / 6],
[1.0, 2.0, 0.0, 4.0, 4.0, 6.0, np.pi / 2],
[1.0, 2.0, 0.0, 4.0, 4.0, 6.0, 7 * np.pi / 6],
[1.0, 2.0, 0.0, 4.0, 4.0, 6.0, -np.pi / 6]],
dtype=torch.float32).cuda()
cam_boxes = DepthInstance3DBoxes(depth_boxes).convert_to(Box3DMode.CAM)
depth_pts = torch.tensor(
[[1, 2, 3.3], [1.2, 2.5, 3.0], [0.8, 2.1, 3.5], [1.6, 2.6, 3.6],
[0.8, 1.2, 3.9], [-9.2, 21.0, 18.2], [3.8, 7.9, 6.3],
[4.7, 3.5, -12.2], [3.8, 7.6, -2], [-10.6, -12.9, -20], [-16, -18, 9],
[-21.3, -52, -5], [0, 0, 0], [6, 7, 8], [-2, -3, -4], [1.0, 4.3, 0.1],
[1.0, 4.4, 0.1], [1.1, 4.3, 0.1], [0.9, 4.3, 0.1], [1.0, -0.3, 0.1],
[1.0, -0.4, 0.1], [2.9, 0.1, 6.0], [-0.9, 3.9, 6.0]],
dtype=torch.float32).cuda()
cam_pts = DepthPoints(depth_pts).convert_to(Coord3DMode.CAM).tensor
point_indices = cam_boxes.points_in_boxes_all(cam_pts)
expected_point_indices = torch.tensor(
[[1, 0, 1, 1, 1, 1], [1, 0, 1, 1, 1, 1], [1, 0, 1, 1, 1, 1],
[1, 0, 1, 1, 1, 1], [1, 0, 1, 1, 1, 1], [0, 1, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0],
[0, 0, 0, 1, 0, 1], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0],
[0, 0, 1, 1, 1, 1], [0, 0, 0, 1, 0, 0], [0, 0, 0, 1, 0, 1],
[0, 0, 1, 1, 1, 0], [0, 0, 1, 1, 1, 1], [0, 0, 0, 1, 0, 0],
[1, 0, 0, 0, 0, 0], [1, 0, 0, 0, 0, 0]],
dtype=torch.int32).cuda()
assert point_indices.shape == torch.Size([23, 6])
assert (point_indices == expected_point_indices).all()
point_indices = cam_boxes.points_in_boxes_batch(cam_pts)
assert (point_indices == expected_point_indices).all()
point_indices = cam_boxes.points_in_boxes_part(cam_pts)
expected_point_indices = torch.tensor([
0, 0, 0, 0, 0, 1, -1, -1, -1, -1, -1, -1, 3, -1, -1, 2, 3, 3, 2, 2, 3,
0, 0
],
dtype=torch.int32).cuda()
assert point_indices.shape == torch.Size([23])
assert (point_indices == expected_point_indices).all()
point_indices = cam_boxes.points_in_boxes(cam_pts)
assert (point_indices == expected_point_indices).all()
tests/test_utils/test_box_np_ops.py
View file @ 53435c62
...@@ -63,3 +63,21 @@ def test_center_to_corner_box2d(): ...@@ -63,3 +63,21 @@ def test_center_to_corner_box2d():
expected_corner = np.array([[[-4.24264, -1.41421], [1.41421, 4.24264], expected_corner = np.array([[[-4.24264, -1.41421], [1.41421, 4.24264],
[4.24264, 1.41421], [-1.41421, -4.24264]]]) [4.24264, 1.41421], [-1.41421, -4.24264]]])
assert np.allclose(corner, expected_corner) assert np.allclose(corner, expected_corner)
def test_points_in_convex_polygon_jit():
from mmdet3d.core.bbox.box_np_ops import points_in_convex_polygon_jit
points = np.array([[0.4, 0.4], [0.5, 0.5], [0.6, 0.6]])
polygons = np.array([[[1.0, 0.0], [0.0, 1.0], [0.0, 0.5], [0.0, 0.0]],
[[1.0, 0.0], [1.0, 1.0], [0.5, 1.0], [0.0, 1.0]],
[[1.0, 0.0], [0.0, 1.0], [-1.0, 0.0], [0.0, -1.0]]])
res = points_in_convex_polygon_jit(points, polygons)
expected_res = np.array([[1, 0, 1], [0, 0, 0], [0, 1, 0]]).astype(np.bool)
assert np.allclose(res, expected_res)
polygons = np.array([[[0.0, 0.0], [0.0, 1.0], [0.5, 0.5], [1.0, 0.0]],
[[0.0, 1.0], [1.0, 1.0], [1.0, 0.5], [1.0, 0.0]],
[[1.0, 0.0], [0.0, -1.0], [-1.0, 0.0], [0.0, 1.1]]])
res = points_in_convex_polygon_jit(points, polygons, clockwise=True)
expected_res = np.array([[1, 0, 1], [0, 0, 1], [0, 1, 0]]).astype(np.bool)
assert np.allclose(res, expected_res)
tests/test_utils/test_coord_3d_mode.py
View file @ 53435c62
...@@ -263,11 +263,11 @@ def test_boxes_conversion(): ...@@ -263,11 +263,11 @@ def test_boxes_conversion():
convert_depth_boxes = Coord3DMode.convert(cam_boxes, Coord3DMode.CAM, convert_depth_boxes = Coord3DMode.convert(cam_boxes, Coord3DMode.CAM,
Coord3DMode.DEPTH) Coord3DMode.DEPTH)
expected_tensor = torch.tensor( expected_tensor = torch.tensor(
[[1.7802, 1.7501, 2.5162, 1.7500, 1.6500, 3.3900, -1.4800], [[1.7802, -1.7501, -2.5162, 1.7500, 1.6500, 3.3900, -1.4800],
[8.9594, 1.6357, 2.4567, 1.5400, 1.5700, 4.0100, -1.6200], [8.9594, -1.6357, -2.4567, 1.5400, 1.5700, 4.0100, -1.6200],
[28.2967, 1.3033, -0.5558, 1.4700, 1.4800, 2.2300, 1.5700], [28.2967, -1.3033, 0.5558, 1.4700, 1.4800, 2.2300, 1.5700],
[26.6690, 1.7361, 21.8230, 1.5600, 1.4000, 3.4800, 1.6900], [26.6690, -1.7361, -21.8230, 1.5600, 1.4000, 3.4800, 1.6900],
[31.3198, 1.6218, 8.1621, 1.7400, 1.4800, 3.7700, -2.7900]]) [31.3198, -1.6218, -8.1621, 1.7400, 1.4800, 3.7700, -2.7900]])
assert torch.allclose(expected_tensor, convert_depth_boxes.tensor, 1e-3) assert torch.allclose(expected_tensor, convert_depth_boxes.tensor, 1e-3)
# test LIDAR to CAM and DEPTH # test LIDAR to CAM and DEPTH
...@@ -327,11 +327,11 @@ def test_boxes_conversion(): ...@@ -327,11 +327,11 @@ def test_boxes_conversion():
convert_cam_boxes = Coord3DMode.convert(depth_boxes, Coord3DMode.DEPTH, convert_cam_boxes = Coord3DMode.convert(depth_boxes, Coord3DMode.DEPTH,
Coord3DMode.CAM) Coord3DMode.CAM)
expected_tensor = torch.tensor( expected_tensor = torch.tensor(
[[1.7802, -1.7501, -2.5162, 1.7500, 1.6500, 3.3900, -1.4800], [[1.7802, 1.7501, 2.5162, 1.7500, 1.6500, 3.3900, -1.4800],
[8.9594, -1.6357, -2.4567, 1.5400, 1.5700, 4.0100, -1.6200], [8.9594, 1.6357, 2.4567, 1.5400, 1.5700, 4.0100, -1.6200],
[28.2967, -1.3033, 0.5558, 1.4700, 1.4800, 2.2300, 1.5700], [28.2967, 1.3033, -0.5558, 1.4700, 1.4800, 2.2300, 1.5700],
[26.6690, -1.7361, -21.8230, 1.5600, 1.4000, 3.4800, 1.6900], [26.6690, 1.7361, 21.8230, 1.5600, 1.4000, 3.4800, 1.6900],
[31.3198, -1.6218, -8.1621, 1.7400, 1.4800, 3.7700, -2.7900]]) [31.3198, 1.6218, 8.1621, 1.7400, 1.4800, 3.7700, -2.7900]])
assert torch.allclose(expected_tensor, convert_cam_boxes.tensor, 1e-3) assert torch.allclose(expected_tensor, convert_cam_boxes.tensor, 1e-3)
convert_lidar_boxes = Coord3DMode.convert(depth_boxes, Coord3DMode.DEPTH, convert_lidar_boxes = Coord3DMode.convert(depth_boxes, Coord3DMode.DEPTH,
......
tests/test_utils/test_points.py
View file @ 53435c62
...@@ -66,6 +66,7 @@ def test_base_points(): ...@@ -66,6 +66,7 @@ def test_base_points():
]]) ]])
assert torch.allclose(expected_tensor, base_points.tensor) assert torch.allclose(expected_tensor, base_points.tensor)
assert torch.allclose(expected_tensor[:, :2], base_points.bev)
assert torch.allclose(expected_tensor[:, :3], base_points.coord) assert torch.allclose(expected_tensor[:, :3], base_points.coord)
assert torch.allclose(expected_tensor[:, 3:6], base_points.color) assert torch.allclose(expected_tensor[:, 3:6], base_points.color)
assert torch.allclose(expected_tensor[:, 6], base_points.height) assert torch.allclose(expected_tensor[:, 6], base_points.height)
...@@ -327,6 +328,7 @@ def test_cam_points(): ...@@ -327,6 +328,7 @@ def test_cam_points():
]]) ]])
assert torch.allclose(expected_tensor, cam_points.tensor) assert torch.allclose(expected_tensor, cam_points.tensor)
assert torch.allclose(expected_tensor[:, [0, 2]], cam_points.bev)
assert torch.allclose(expected_tensor[:, :3], cam_points.coord) assert torch.allclose(expected_tensor[:, :3], cam_points.coord)
assert torch.allclose(expected_tensor[:, 3:6], cam_points.color) assert torch.allclose(expected_tensor[:, 3:6], cam_points.color)
assert torch.allclose(expected_tensor[:, 6], cam_points.height) assert torch.allclose(expected_tensor[:, 6], cam_points.height)
...@@ -603,6 +605,7 @@ def test_lidar_points(): ...@@ -603,6 +605,7 @@ def test_lidar_points():
]]) ]])
assert torch.allclose(expected_tensor, lidar_points.tensor) assert torch.allclose(expected_tensor, lidar_points.tensor)
assert torch.allclose(expected_tensor[:, :2], lidar_points.bev)
assert torch.allclose(expected_tensor[:, :3], lidar_points.coord) assert torch.allclose(expected_tensor[:, :3], lidar_points.coord)
assert torch.allclose(expected_tensor[:, 3:6], lidar_points.color) assert torch.allclose(expected_tensor[:, 3:6], lidar_points.color)
assert torch.allclose(expected_tensor[:, 6], lidar_points.height) assert torch.allclose(expected_tensor[:, 6], lidar_points.height)
...@@ -879,6 +882,7 @@ def test_depth_points(): ...@@ -879,6 +882,7 @@ def test_depth_points():
]]) ]])
assert torch.allclose(expected_tensor, depth_points.tensor) assert torch.allclose(expected_tensor, depth_points.tensor)
assert torch.allclose(expected_tensor[:, :2], depth_points.bev)
assert torch.allclose(expected_tensor[:, :3], depth_points.coord) assert torch.allclose(expected_tensor[:, :3], depth_points.coord)
assert torch.allclose(expected_tensor[:, 3:6], depth_points.color) assert torch.allclose(expected_tensor[:, 3:6], depth_points.color)
assert torch.allclose(expected_tensor[:, 6], depth_points.height) assert torch.allclose(expected_tensor[:, 6], depth_points.height)
......
tools/create_data.py
View file @ 53435c62
...@@ -61,7 +61,8 @@ def nuscenes_data_prep(root_path, ...@@ -61,7 +61,8 @@ def nuscenes_data_prep(root_path,
version (str): Dataset version. version (str): Dataset version.
dataset_name (str): The dataset class name. dataset_name (str): The dataset class name.
out_dir (str): Output directory of the groundtruth database info. out_dir (str): Output directory of the groundtruth database info.
max_sweeps (int): Number of input consecutive frames. Default: 10 max_sweeps (int, optional): Number of input consecutive frames.
Default: 10
""" """
nuscenes_converter.create_nuscenes_infos( nuscenes_converter.create_nuscenes_infos(
root_path, info_prefix, version=version, max_sweeps=max_sweeps) root_path, info_prefix, version=version, max_sweeps=max_sweeps)
...@@ -152,8 +153,9 @@ def waymo_data_prep(root_path, ...@@ -152,8 +153,9 @@ def waymo_data_prep(root_path,
info_prefix (str): The prefix of info filenames. info_prefix (str): The prefix of info filenames.
out_dir (str): Output directory of the generated info file. out_dir (str): Output directory of the generated info file.
workers (int): Number of threads to be used. workers (int): Number of threads to be used.
max_sweeps (int): Number of input consecutive frames. Default: 5 \ max_sweeps (int, optional): Number of input consecutive frames.
Here we store pose information of these frames for later use. Default: 5. Here we store pose information of these frames
for later use.
""" """
from tools.data_converter import waymo_converter as waymo from tools.data_converter import waymo_converter as waymo
......
tools/data_converter/create_gt_database.py
View file @ 53435c62
...@@ -126,19 +126,19 @@ def create_groundtruth_database(dataset_class_name, ...@@ -126,19 +126,19 @@ def create_groundtruth_database(dataset_class_name,
dataset_class_name （str): Name of the input dataset. dataset_class_name （str): Name of the input dataset.
data_path (str): Path of the data. data_path (str): Path of the data.
info_prefix (str): Prefix of the info file. info_prefix (str): Prefix of the info file.
info_path (str): Path of the info file. info_path (str, optional): Path of the info file.
Default: None. Default: None.
mask_anno_path (str): Path of the mask_anno. mask_anno_path (str, optional): Path of the mask_anno.
Default: None. Default: None.
used_classes (list[str]): Classes have been used. used_classes (list[str], optional): Classes have been used.
Default: None. Default: None.
database_save_path (str): Path to save database. database_save_path (str, optional): Path to save database.
Default: None. Default: None.
db_info_save_path (str): Path to save db_info. db_info_save_path (str, optional): Path to save db_info.
Default: None. Default: None.
relative_path (bool): Whether to use relative path. relative_path (bool, optional): Whether to use relative path.
Default: True. Default: True.
with_mask (bool): Whether to use mask. with_mask (bool, optional): Whether to use mask.
Default: False. Default: False.
""" """
print(f'Create GT Database of {dataset_class_name}') print(f'Create GT Database of {dataset_class_name}')
......
tools/data_converter/indoor_converter.py
View file @ 53435c62
...@@ -19,10 +19,11 @@ def create_indoor_info_file(data_path, ...@@ -19,10 +19,11 @@ def create_indoor_info_file(data_path,
Args: Args:
data_path (str): Path of the data. data_path (str): Path of the data.
pkl_prefix (str): Prefix of the pkl to be saved. Default: 'sunrgbd'. pkl_prefix (str, optional): Prefix of the pkl to be saved.
save_path (str): Path of the pkl to be saved. Default: None. Default: 'sunrgbd'.
use_v1 (bool): Whether to use v1. Default: False. save_path (str, optional): Path of the pkl to be saved. Default: None.
workers (int): Number of threads to be used. Default: 4. use_v1 (bool, optional): Whether to use v1. Default: False.
workers (int, optional): Number of threads to be used. Default: 4.
""" """
assert os.path.exists(data_path) assert os.path.exists(data_path)
assert pkl_prefix in ['sunrgbd', 'scannet', 's3dis'], \ assert pkl_prefix in ['sunrgbd', 'scannet', 's3dis'], \
......
tools/data_converter/kitti_converter.py
View file @ 53435c62
...@@ -159,7 +159,7 @@ def create_waymo_info_file(data_path, ...@@ -159,7 +159,7 @@ def create_waymo_info_file(data_path,
Args: Args:
data_path (str): Path of the data root. data_path (str): Path of the data root.
pkl_prefix (str): Prefix of the info file to be generated. pkl_prefix (str): Prefix of the info file to be generated.
save_path (str | None): Path to save the info file. save_path (str): Path to save the info file.
relative_path (bool): Whether to use relative path. relative_path (bool): Whether to use relative path.
max_sweeps (int): Max sweeps before the detection frame to be used. max_sweeps (int): Max sweeps before the detection frame to be used.
""" """
...@@ -238,11 +238,13 @@ def _create_reduced_point_cloud(data_path, ...@@ -238,11 +238,13 @@ def _create_reduced_point_cloud(data_path,
Args: Args:
data_path (str): Path of original data. data_path (str): Path of original data.
info_path (str): Path of data info. info_path (str): Path of data info.
save_path (str | None): Path to save reduced point cloud data. save_path (str, optional): Path to save reduced point cloud
Default: None. data. Default: None.
back (bool): Whether to flip the points to back. back (bool, optional): Whether to flip the points to back.
num_features (int): Number of point features. Default: 4. Default: False.
front_camera_id (int): The referenced/front camera ID. Default: 2. num_features (int, optional): Number of point features. Default: 4.
front_camera_id (int, optional): The referenced/front camera ID.
Default: 2.
""" """
kitti_infos = mmcv.load(info_path) kitti_infos = mmcv.load(info_path)
...@@ -298,14 +300,16 @@ def create_reduced_point_cloud(data_path, ...@@ -298,14 +300,16 @@ def create_reduced_point_cloud(data_path,
Args: Args:
data_path (str): Path of original data. data_path (str): Path of original data.
pkl_prefix (str): Prefix of info files. pkl_prefix (str): Prefix of info files.
train_info_path (str | None): Path of training set info. train_info_path (str, optional): Path of training set info.
Default: None.
val_info_path (str, optional): Path of validation set info.
Default: None. Default: None.
val_info_path (str | None): Path of validation set info. test_info_path (str, optional): Path of test set info.
Default: None. Default: None.
test_info_path (str | None): Path of test set info. save_path (str, optional): Path to save reduced point cloud data.
Default: None. Default: None.
save_path (str | None): Path to save reduced point cloud data. with_back (bool, optional): Whether to flip the points to back.
with_back (bool): Whether to flip the points to back. Default: False.
""" """
if train_info_path is None: if train_info_path is None:
train_info_path = Path(data_path) / f'{pkl_prefix}_infos_train.pkl' train_info_path = Path(data_path) / f'{pkl_prefix}_infos_train.pkl'
...@@ -335,7 +339,8 @@ def export_2d_annotation(root_path, info_path, mono3d=True): ...@@ -335,7 +339,8 @@ def export_2d_annotation(root_path, info_path, mono3d=True):
Args: Args:
root_path (str): Root path of the raw data. root_path (str): Root path of the raw data.
info_path (str): Path of the info file. info_path (str): Path of the info file.
mono3d (bool): Whether to export mono3d annotation. Default: True. mono3d (bool, optional): Whether to export mono3d annotation.
Default: True.
""" """
# get bbox annotations for camera # get bbox annotations for camera
kitti_infos = mmcv.load(info_path) kitti_infos = mmcv.load(info_path)
...@@ -381,8 +386,8 @@ def get_2d_boxes(info, occluded, mono3d=True): ...@@ -381,8 +386,8 @@ def get_2d_boxes(info, occluded, mono3d=True):
Args: Args:
info: Information of the given sample data. info: Information of the given sample data.
occluded: Integer (0, 1, 2, 3) indicating occlusion state: \ occluded: Integer (0, 1, 2, 3) indicating occlusion state:
0 = fully visible, 1 = partly occluded, 2 = largely occluded, \ 0 = fully visible, 1 = partly occluded, 2 = largely occluded,
3 = unknown, -1 = DontCare 3 = unknown, -1 = DontCare
mono3d (bool): Whether to get boxes with mono3d annotation. mono3d (bool): Whether to get boxes with mono3d annotation.
...@@ -508,7 +513,7 @@ def generate_record(ann_rec, x1, y1, x2, y2, sample_data_token, filename): ...@@ -508,7 +513,7 @@ def generate_record(ann_rec, x1, y1, x2, y2, sample_data_token, filename):
- area (float): 2d box area - area (float): 2d box area
- category_name (str): category name - category_name (str): category name
- category_id (int): category id - category_id (int): category id
- bbox (list[float]): left x, top y, dx, dy of 2d box - bbox (list[float]): left x, top y, x_size, y_size of 2d box
- iscrowd (int): whether the area is crowd - iscrowd (int): whether the area is crowd
""" """
repro_rec = OrderedDict() repro_rec = OrderedDict()
......
tools/data_converter/lyft_converter.py
View file @ 53435c62
...@@ -26,10 +26,10 @@ def create_lyft_infos(root_path, ...@@ -26,10 +26,10 @@ def create_lyft_infos(root_path,
Args: Args:
root_path (str): Path of the data root. root_path (str): Path of the data root.
info_prefix (str): Prefix of the info file to be generated. info_prefix (str): Prefix of the info file to be generated.
version (str): Version of the data. version (str, optional): Version of the data.
Default: 'v1.01-train' Default: 'v1.01-train'.
max_sweeps (int): Max number of sweeps. max_sweeps (int, optional): Max number of sweeps.
Default: 10 Default: 10.
""" """
lyft = Lyft( lyft = Lyft(
data_path=osp.join(root_path, version), data_path=osp.join(root_path, version),
...@@ -101,9 +101,9 @@ def _fill_trainval_infos(lyft, ...@@ -101,9 +101,9 @@ def _fill_trainval_infos(lyft,
lyft (:obj:`LyftDataset`): Dataset class in the Lyft dataset. lyft (:obj:`LyftDataset`): Dataset class in the Lyft dataset.
train_scenes (list[str]): Basic information of training scenes. train_scenes (list[str]): Basic information of training scenes.
val_scenes (list[str]): Basic information of validation scenes. val_scenes (list[str]): Basic information of validation scenes.
test (bool): Whether use the test mode. In the test mode, no test (bool, optional): Whether use the test mode. In the test mode, no
annotations can be accessed. Default: False. annotations can be accessed. Default: False.
max_sweeps (int): Max number of sweeps. Default: 10. max_sweeps (int, optional): Max number of sweeps. Default: 10.
Returns: Returns:
tuple[list[dict]]: Information of training set and tuple[list[dict]]: Information of training set and
...@@ -194,7 +194,7 @@ def _fill_trainval_infos(lyft, ...@@ -194,7 +194,7 @@ def _fill_trainval_infos(lyft,
# we need to convert box size to # we need to convert box size to
# the format of our lidar coordinate system # the format of our lidar coordinate system
# which is dx, dy, dz (corresponding to l, w, h) # which is x_size, y_size, z_size (corresponding to l, w, h)
gt_boxes = np.concatenate([locs, dims[:, [1, 0, 2]], rots], axis=1) gt_boxes = np.concatenate([locs, dims[:, [1, 0, 2]], rots], axis=1)
assert len(gt_boxes) == len( assert len(gt_boxes) == len(
annotations), f'{len(gt_boxes)}, {len(annotations)}' annotations), f'{len(gt_boxes)}, {len(annotations)}'
......
tools/data_converter/nuscenes_converter.py
View file @ 53435c62
...@@ -34,10 +34,10 @@ def create_nuscenes_infos(root_path, ...@@ -34,10 +34,10 @@ def create_nuscenes_infos(root_path,
Args: Args:
root_path (str): Path of the data root. root_path (str): Path of the data root.
info_prefix (str): Prefix of the info file to be generated. info_prefix (str): Prefix of the info file to be generated.
version (str): Version of the data. version (str, optional): Version of the data.
Default: 'v1.0-trainval' Default: 'v1.0-trainval'.
max_sweeps (int): Max number of sweeps. max_sweeps (int, optional): Max number of sweeps.
Default: 10 Default: 10.
""" """
from nuscenes.nuscenes import NuScenes from nuscenes.nuscenes import NuScenes
nusc = NuScenes(version=version, dataroot=root_path, verbose=True) nusc = NuScenes(version=version, dataroot=root_path, verbose=True)
...@@ -152,9 +152,9 @@ def _fill_trainval_infos(nusc, ...@@ -152,9 +152,9 @@ def _fill_trainval_infos(nusc,
nusc (:obj:`NuScenes`): Dataset class in the nuScenes dataset. nusc (:obj:`NuScenes`): Dataset class in the nuScenes dataset.
train_scenes (list[str]): Basic information of training scenes. train_scenes (list[str]): Basic information of training scenes.
val_scenes (list[str]): Basic information of validation scenes. val_scenes (list[str]): Basic information of validation scenes.
test (bool): Whether use the test mode. In the test mode, no test (bool, optional): Whether use the test mode. In test mode, no
annotations can be accessed. Default: False. annotations can be accessed. Default: False.
max_sweeps (int): Max number of sweeps. Default: 10. max_sweeps (int, optional): Max number of sweeps. Default: 10.
Returns: Returns:
tuple[list[dict]]: Information of training set and validation set tuple[list[dict]]: Information of training set and validation set
...@@ -251,7 +251,7 @@ def _fill_trainval_infos(nusc, ...@@ -251,7 +251,7 @@ def _fill_trainval_infos(nusc,
names = np.array(names) names = np.array(names)
# we need to convert box size to # we need to convert box size to
# the format of our lidar coordinate system # the format of our lidar coordinate system
# which is dx, dy, dz (corresponding to l, w, h) # which is x_size, y_size, z_size (corresponding to l, w, h)
gt_boxes = np.concatenate([locs, dims[:, [1, 0, 2]], rots], axis=1) gt_boxes = np.concatenate([locs, dims[:, [1, 0, 2]], rots], axis=1)
assert len(gt_boxes) == len( assert len(gt_boxes) == len(
annotations), f'{len(gt_boxes)}, {len(annotations)}' annotations), f'{len(gt_boxes)}, {len(annotations)}'
...@@ -291,7 +291,7 @@ def obtain_sensor2top(nusc, ...@@ -291,7 +291,7 @@ def obtain_sensor2top(nusc,
e2g_t (np.ndarray): Translation from ego to global in shape (1, 3). e2g_t (np.ndarray): Translation from ego to global in shape (1, 3).
e2g_r_mat (np.ndarray): Rotation matrix from ego to global e2g_r_mat (np.ndarray): Rotation matrix from ego to global
in shape (3, 3). in shape (3, 3).
sensor_type (str): Sensor to calibrate. Default: 'lidar'. sensor_type (str, optional): Sensor to calibrate. Default: 'lidar'.
Returns: Returns:
sweep (dict): Sweep information after transformation. sweep (dict): Sweep information after transformation.
...@@ -340,7 +340,8 @@ def export_2d_annotation(root_path, info_path, version, mono3d=True): ...@@ -340,7 +340,8 @@ def export_2d_annotation(root_path, info_path, version, mono3d=True):
root_path (str): Root path of the raw data. root_path (str): Root path of the raw data.
info_path (str): Path of the info file. info_path (str): Path of the info file.
version (str): Dataset version. version (str): Dataset version.
mono3d (bool): Whether to export mono3d annotation. Default: True. mono3d (bool, optional): Whether to export mono3d annotation.
Default: True.
""" """
# get bbox annotations for camera # get bbox annotations for camera
camera_types = [ camera_types = [
...@@ -404,7 +405,7 @@ def get_2d_boxes(nusc, ...@@ -404,7 +405,7 @@ def get_2d_boxes(nusc,
"""Get the 2D annotation records for a given `sample_data_token`. """Get the 2D annotation records for a given `sample_data_token`.
Args: Args:
sample_data_token (str): Sample data token belonging to a camera \ sample_data_token (str): Sample data token belonging to a camera
keyframe. keyframe.
visibilities (list[str]): Visibility filter. visibilities (list[str]): Visibility filter.
mono3d (bool): Whether to get boxes with mono3d annotation. mono3d (bool): Whether to get boxes with mono3d annotation.
......
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