Unverified Commit 333536f6 authored by Wenwei Zhang's avatar Wenwei Zhang Committed by GitHub
Browse files

Release v1.0.0rc1

parents 9c7270d0 f747daab
#pragma once
#include <torch/extension.h>
typedef enum { SUM = 0, MEAN = 1, MAX = 2 } reduce_t;
namespace voxelization {
int hard_voxelize_cpu(const at::Tensor &points, at::Tensor &voxels,
at::Tensor &coors, at::Tensor &num_points_per_voxel,
const std::vector<float> voxel_size,
const std::vector<float> coors_range,
const int max_points, const int max_voxels,
const int NDim = 3);
void dynamic_voxelize_cpu(const at::Tensor &points, at::Tensor &coors,
const std::vector<float> voxel_size,
const std::vector<float> coors_range,
const int NDim = 3);
std::vector<at::Tensor> dynamic_point_to_voxel_cpu(
const at::Tensor &points, const at::Tensor &voxel_mapping,
const std::vector<float> voxel_size, const std::vector<float> coors_range);
#ifdef WITH_CUDA
int hard_voxelize_gpu(const at::Tensor &points, at::Tensor &voxels,
at::Tensor &coors, at::Tensor &num_points_per_voxel,
const std::vector<float> voxel_size,
const std::vector<float> coors_range,
const int max_points, const int max_voxels,
const int NDim = 3);
int nondisterministic_hard_voxelize_gpu(const at::Tensor &points, at::Tensor &voxels,
at::Tensor &coors, at::Tensor &num_points_per_voxel,
const std::vector<float> voxel_size,
const std::vector<float> coors_range,
const int max_points, const int max_voxels,
const int NDim = 3);
void dynamic_voxelize_gpu(const at::Tensor &points, at::Tensor &coors,
const std::vector<float> voxel_size,
const std::vector<float> coors_range,
const int NDim = 3);
std::vector<torch::Tensor> dynamic_point_to_voxel_forward_gpu(const torch::Tensor &feats,
const torch::Tensor &coors,
const reduce_t reduce_type);
void dynamic_point_to_voxel_backward_gpu(torch::Tensor &grad_feats,
const torch::Tensor &grad_reduced_feats,
const torch::Tensor &feats,
const torch::Tensor &reduced_feats,
const torch::Tensor &coors_idx,
const torch::Tensor &reduce_count,
const reduce_t reduce_type);
#endif
// Interface for Python
inline int hard_voxelize(const at::Tensor &points, at::Tensor &voxels,
at::Tensor &coors, at::Tensor &num_points_per_voxel,
const std::vector<float> voxel_size,
const std::vector<float> coors_range,
const int max_points, const int max_voxels,
const int NDim = 3, const bool deterministic = true) {
if (points.device().is_cuda()) {
#ifdef WITH_CUDA
if (deterministic) {
return hard_voxelize_gpu(points, voxels, coors, num_points_per_voxel,
voxel_size, coors_range, max_points, max_voxels,
NDim);
}
return nondisterministic_hard_voxelize_gpu(points, voxels, coors, num_points_per_voxel,
voxel_size, coors_range, max_points, max_voxels,
NDim);
#else
AT_ERROR("Not compiled with GPU support");
#endif
}
return hard_voxelize_cpu(points, voxels, coors, num_points_per_voxel,
voxel_size, coors_range, max_points, max_voxels,
NDim);
}
inline void dynamic_voxelize(const at::Tensor &points, at::Tensor &coors,
const std::vector<float> voxel_size,
const std::vector<float> coors_range,
const int NDim = 3) {
if (points.device().is_cuda()) {
#ifdef WITH_CUDA
return dynamic_voxelize_gpu(points, coors, voxel_size, coors_range, NDim);
#else
AT_ERROR("Not compiled with GPU support");
#endif
}
return dynamic_voxelize_cpu(points, coors, voxel_size, coors_range, NDim);
}
inline reduce_t convert_reduce_type(const std::string &reduce_type) {
if (reduce_type == "max")
return reduce_t::MAX;
else if (reduce_type == "sum")
return reduce_t::SUM;
else if (reduce_type == "mean")
return reduce_t::MEAN;
else TORCH_CHECK(false, "do not support reduce type " + reduce_type)
return reduce_t::SUM;
}
inline std::vector<torch::Tensor> dynamic_point_to_voxel_forward(const torch::Tensor &feats,
const torch::Tensor &coors,
const std::string &reduce_type) {
if (feats.device().is_cuda()) {
#ifdef WITH_CUDA
return dynamic_point_to_voxel_forward_gpu(feats, coors, convert_reduce_type(reduce_type));
#else
TORCH_CHECK(false, "Not compiled with GPU support");
#endif
}
TORCH_CHECK(false, "do not support cpu yet");
return std::vector<torch::Tensor>();
}
inline void dynamic_point_to_voxel_backward(torch::Tensor &grad_feats,
const torch::Tensor &grad_reduced_feats,
const torch::Tensor &feats,
const torch::Tensor &reduced_feats,
const torch::Tensor &coors_idx,
const torch::Tensor &reduce_count,
const std::string &reduce_type) {
if (grad_feats.device().is_cuda()) {
#ifdef WITH_CUDA
dynamic_point_to_voxel_backward_gpu(
grad_feats, grad_reduced_feats, feats, reduced_feats, coors_idx, reduce_count,
convert_reduce_type(reduce_type));
return;
#else
TORCH_CHECK(false, "Not compiled with GPU support");
#endif
}
TORCH_CHECK(false, "do not support cpu yet");
}
} // namespace voxelization
#include <ATen/TensorUtils.h>
#include <torch/extension.h>
// #include "voxelization.h"
namespace {
template <typename T, typename T_int>
void dynamic_voxelize_kernel(const torch::TensorAccessor<T, 2> points,
torch::TensorAccessor<T_int, 2> coors,
const std::vector<float> voxel_size,
const std::vector<float> coors_range,
const std::vector<int> grid_size,
const int num_points, const int num_features,
const int NDim) {
const int ndim_minus_1 = NDim - 1;
bool failed = false;
// int coor[NDim];
int* coor = new int[NDim]();
int c;
for (int i = 0; i < num_points; ++i) {
failed = false;
for (int j = 0; j < NDim; ++j) {
c = floor((points[i][j] - coors_range[j]) / voxel_size[j]);
// necessary to rm points out of range
if ((c < 0 || c >= grid_size[j])) {
failed = true;
break;
}
coor[ndim_minus_1 - j] = c;
}
for (int k = 0; k < NDim; ++k) {
if (failed)
coors[i][k] = -1;
else
coors[i][k] = coor[k];
}
}
delete[] coor;
return;
}
template <typename T, typename T_int>
void hard_voxelize_kernel(const torch::TensorAccessor<T, 2> points,
torch::TensorAccessor<T, 3> voxels,
torch::TensorAccessor<T_int, 2> coors,
torch::TensorAccessor<T_int, 1> num_points_per_voxel,
torch::TensorAccessor<T_int, 3> coor_to_voxelidx,
int& voxel_num, const std::vector<float> voxel_size,
const std::vector<float> coors_range,
const std::vector<int> grid_size,
const int max_points, const int max_voxels,
const int num_points, const int num_features,
const int NDim) {
// declare a temp coors
at::Tensor temp_coors = at::zeros(
{num_points, NDim}, at::TensorOptions().dtype(at::kInt).device(at::kCPU));
// First use dynamic voxelization to get coors,
// then check max points/voxels constraints
dynamic_voxelize_kernel<T, int>(points, temp_coors.accessor<int, 2>(),
voxel_size, coors_range, grid_size,
num_points, num_features, NDim);
int voxelidx, num;
auto coor = temp_coors.accessor<int, 2>();
for (int i = 0; i < num_points; ++i) {
// T_int* coor = temp_coors.data_ptr<int>() + i * NDim;
if (coor[i][0] == -1) continue;
voxelidx = coor_to_voxelidx[coor[i][0]][coor[i][1]][coor[i][2]];
// record voxel
if (voxelidx == -1) {
voxelidx = voxel_num;
if (max_voxels != -1 && voxel_num >= max_voxels) continue;
voxel_num += 1;
coor_to_voxelidx[coor[i][0]][coor[i][1]][coor[i][2]] = voxelidx;
for (int k = 0; k < NDim; ++k) {
coors[voxelidx][k] = coor[i][k];
}
}
// put points into voxel
num = num_points_per_voxel[voxelidx];
if (max_points == -1 || num < max_points) {
for (int k = 0; k < num_features; ++k) {
voxels[voxelidx][num][k] = points[i][k];
}
num_points_per_voxel[voxelidx] += 1;
}
}
return;
}
} // namespace
namespace voxelization {
int hard_voxelize_cpu(const at::Tensor& points, at::Tensor& voxels,
at::Tensor& coors, at::Tensor& num_points_per_voxel,
const std::vector<float> voxel_size,
const std::vector<float> coors_range,
const int max_points, const int max_voxels,
const int NDim = 3) {
// current version tooks about 0.02s_0.03s for one frame on cpu
// check device
AT_ASSERTM(points.device().is_cpu(), "points must be a CPU tensor");
std::vector<int> grid_size(NDim);
const int num_points = points.size(0);
const int num_features = points.size(1);
for (int i = 0; i < NDim; ++i) {
grid_size[i] =
round((coors_range[NDim + i] - coors_range[i]) / voxel_size[i]);
}
// coors, num_points_per_voxel, coor_to_voxelidx are int Tensor
// printf("cpu coor_to_voxelidx size: [%d, %d, %d]\n", grid_size[2],
// grid_size[1], grid_size[0]);
at::Tensor coor_to_voxelidx =
-at::ones({grid_size[2], grid_size[1], grid_size[0]}, coors.options());
int voxel_num = 0;
AT_DISPATCH_FLOATING_TYPES_AND_HALF(
points.scalar_type(), "hard_voxelize_forward", [&] {
hard_voxelize_kernel<scalar_t, int>(
points.accessor<scalar_t, 2>(), voxels.accessor<scalar_t, 3>(),
coors.accessor<int, 2>(), num_points_per_voxel.accessor<int, 1>(),
coor_to_voxelidx.accessor<int, 3>(), voxel_num, voxel_size,
coors_range, grid_size, max_points, max_voxels, num_points,
num_features, NDim);
});
return voxel_num;
}
void dynamic_voxelize_cpu(const at::Tensor& points, at::Tensor& coors,
const std::vector<float> voxel_size,
const std::vector<float> coors_range,
const int NDim = 3) {
// check device
AT_ASSERTM(points.device().is_cpu(), "points must be a CPU tensor");
std::vector<int> grid_size(NDim);
const int num_points = points.size(0);
const int num_features = points.size(1);
for (int i = 0; i < NDim; ++i) {
grid_size[i] =
round((coors_range[NDim + i] - coors_range[i]) / voxel_size[i]);
}
// coors, num_points_per_voxel, coor_to_voxelidx are int Tensor
AT_DISPATCH_FLOATING_TYPES_AND_HALF(
points.scalar_type(), "hard_voxelize_forward", [&] {
dynamic_voxelize_kernel<scalar_t, int>(
points.accessor<scalar_t, 2>(), coors.accessor<int, 2>(),
voxel_size, coors_range, grid_size, num_points, num_features, NDim);
});
return;
}
} // namespace voxelization
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/types.h>
#include <ATen/cuda/CUDAApplyUtils.cuh>
#define CHECK_CUDA(x) \
TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) \
TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) \
CHECK_CUDA(x); \
CHECK_CONTIGUOUS(x)
namespace {
int const threadsPerBlock = sizeof(unsigned long long) * 8;
}
#define CUDA_1D_KERNEL_LOOP(i, n) \
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; \
i += blockDim.x * gridDim.x)
template <typename T, typename T_int>
__global__ void dynamic_voxelize_kernel(
const T* points, T_int* coors, const float voxel_x, const float voxel_y,
const float voxel_z, const float coors_x_min, const float coors_y_min,
const float coors_z_min, const float coors_x_max, const float coors_y_max,
const float coors_z_max, const int grid_x, const int grid_y,
const int grid_z, const int num_points, const int num_features,
const int NDim) {
// const int index = blockIdx.x * threadsPerBlock + threadIdx.x;
CUDA_1D_KERNEL_LOOP(index, num_points) {
// To save some computation
auto points_offset = points + index * num_features;
auto coors_offset = coors + index * NDim;
int c_x = floor((points_offset[0] - coors_x_min) / voxel_x);
if (c_x < 0 || c_x >= grid_x) {
coors_offset[0] = -1;
return;
}
int c_y = floor((points_offset[1] - coors_y_min) / voxel_y);
if (c_y < 0 || c_y >= grid_y) {
coors_offset[0] = -1;
coors_offset[1] = -1;
return;
}
int c_z = floor((points_offset[2] - coors_z_min) / voxel_z);
if (c_z < 0 || c_z >= grid_z) {
coors_offset[0] = -1;
coors_offset[1] = -1;
coors_offset[2] = -1;
} else {
coors_offset[0] = c_z;
coors_offset[1] = c_y;
coors_offset[2] = c_x;
}
}
}
template <typename T, typename T_int>
__global__ void assign_point_to_voxel(const int nthreads, const T* points,
T_int* point_to_voxelidx,
T_int* coor_to_voxelidx, T* voxels,
const int max_points,
const int num_features,
const int num_points, const int NDim) {
CUDA_1D_KERNEL_LOOP(thread_idx, nthreads) {
// const int index = blockIdx.x * threadsPerBlock + threadIdx.x;
int index = thread_idx / num_features;
int num = point_to_voxelidx[index];
int voxelidx = coor_to_voxelidx[index];
if (num > -1 && voxelidx > -1) {
auto voxels_offset =
voxels + voxelidx * max_points * num_features + num * num_features;
int k = thread_idx % num_features;
voxels_offset[k] = points[thread_idx];
}
}
}
template <typename T, typename T_int>
__global__ void assign_voxel_coors(const int nthreads, T_int* coor,
T_int* point_to_voxelidx,
T_int* coor_to_voxelidx, T_int* voxel_coors,
const int num_points, const int NDim) {
CUDA_1D_KERNEL_LOOP(thread_idx, nthreads) {
// const int index = blockIdx.x * threadsPerBlock + threadIdx.x;
// if (index >= num_points) return;
int index = thread_idx / NDim;
int num = point_to_voxelidx[index];
int voxelidx = coor_to_voxelidx[index];
if (num == 0 && voxelidx > -1) {
auto coors_offset = voxel_coors + voxelidx * NDim;
int k = thread_idx % NDim;
coors_offset[k] = coor[thread_idx];
}
}
}
template <typename T_int>
__global__ void point_to_voxelidx_kernel(const T_int* coor,
T_int* point_to_voxelidx,
T_int* point_to_pointidx,
const int max_points,
const int max_voxels,
const int num_points, const int NDim) {
CUDA_1D_KERNEL_LOOP(index, num_points) {
auto coor_offset = coor + index * NDim;
// skip invalid points
if ((index >= num_points) || (coor_offset[0] == -1)) return;
int num = 0;
int coor_x = coor_offset[0];
int coor_y = coor_offset[1];
int coor_z = coor_offset[2];
// only calculate the coors before this coor[index]
for (int i = 0; i < index; ++i) {
auto prev_coor = coor + i * NDim;
if (prev_coor[0] == -1) continue;
// Find all previous points that have the same coors
// if find the same coor, record it
if ((prev_coor[0] == coor_x) && (prev_coor[1] == coor_y) &&
(prev_coor[2] == coor_z)) {
num++;
if (num == 1) {
// point to the same coor that first show up
point_to_pointidx[index] = i;
} else if (num >= max_points) {
// out of boundary
return;
}
}
}
if (num == 0) {
point_to_pointidx[index] = index;
}
if (num < max_points) {
point_to_voxelidx[index] = num;
}
}
}
template <typename T_int>
__global__ void determin_voxel_num(
// const T_int* coor,
T_int* num_points_per_voxel, T_int* point_to_voxelidx,
T_int* point_to_pointidx, T_int* coor_to_voxelidx, T_int* voxel_num,
const int max_points, const int max_voxels, const int num_points) {
// only calculate the coors before this coor[index]
for (int i = 0; i < num_points; ++i) {
// if (coor[i][0] == -1)
// continue;
int point_pos_in_voxel = point_to_voxelidx[i];
// record voxel
if (point_pos_in_voxel == -1) {
// out of max_points or invalid point
continue;
} else if (point_pos_in_voxel == 0) {
// record new voxel
int voxelidx = voxel_num[0];
if (voxel_num[0] >= max_voxels) continue;
voxel_num[0] += 1;
coor_to_voxelidx[i] = voxelidx;
num_points_per_voxel[voxelidx] = 1;
} else {
int point_idx = point_to_pointidx[i];
int voxelidx = coor_to_voxelidx[point_idx];
if (voxelidx != -1) {
coor_to_voxelidx[i] = voxelidx;
num_points_per_voxel[voxelidx] += 1;
}
}
}
}
__global__ void nondisterministic_get_assign_pos(
const int nthreads, const int32_t *coors_map, int32_t *pts_id,
int32_t *coors_count, int32_t *reduce_count, int32_t *coors_order) {
CUDA_1D_KERNEL_LOOP(thread_idx, nthreads) {
int coors_idx = coors_map[thread_idx];
if (coors_idx > -1) {
int32_t coors_pts_pos = atomicAdd(&reduce_count[coors_idx], 1);
pts_id[thread_idx] = coors_pts_pos;
if (coors_pts_pos == 0) {
coors_order[coors_idx] = atomicAdd(coors_count, 1);
}
}
}
}
template<typename T>
__global__ void nondisterministic_assign_point_voxel(
const int nthreads, const T *points, const int32_t *coors_map,
const int32_t *pts_id, const int32_t *coors_in,
const int32_t *reduce_count, const int32_t *coors_order,
T *voxels, int32_t *coors, int32_t *pts_count, const int max_voxels,
const int max_points, const int num_features, const int NDim) {
CUDA_1D_KERNEL_LOOP(thread_idx, nthreads) {
int coors_idx = coors_map[thread_idx];
int coors_pts_pos = pts_id[thread_idx];
if (coors_idx > -1) {
int coors_pos = coors_order[coors_idx];
if (coors_pos < max_voxels && coors_pts_pos < max_points) {
auto voxels_offset =
voxels + (coors_pos * max_points + coors_pts_pos) * num_features;
auto points_offset = points + thread_idx * num_features;
for (int k = 0; k < num_features; k++) {
voxels_offset[k] = points_offset[k];
}
if (coors_pts_pos == 0) {
pts_count[coors_pos] = min(reduce_count[coors_idx], max_points);
auto coors_offset = coors + coors_pos * NDim;
auto coors_in_offset = coors_in + coors_idx * NDim;
for (int k = 0; k < NDim; k++) {
coors_offset[k] = coors_in_offset[k];
}
}
}
}
}
}
namespace voxelization {
int hard_voxelize_gpu(const at::Tensor& points, at::Tensor& voxels,
at::Tensor& coors, at::Tensor& num_points_per_voxel,
const std::vector<float> voxel_size,
const std::vector<float> coors_range,
const int max_points, const int max_voxels,
const int NDim = 3) {
// current version tooks about 0.04s for one frame on cpu
// check device
CHECK_INPUT(points);
at::cuda::CUDAGuard device_guard(points.device());
const int num_points = points.size(0);
const int num_features = points.size(1);
const float voxel_x = voxel_size[0];
const float voxel_y = voxel_size[1];
const float voxel_z = voxel_size[2];
const float coors_x_min = coors_range[0];
const float coors_y_min = coors_range[1];
const float coors_z_min = coors_range[2];
const float coors_x_max = coors_range[3];
const float coors_y_max = coors_range[4];
const float coors_z_max = coors_range[5];
const int grid_x = round((coors_x_max - coors_x_min) / voxel_x);
const int grid_y = round((coors_y_max - coors_y_min) / voxel_y);
const int grid_z = round((coors_z_max - coors_z_min) / voxel_z);
// map points to voxel coors
at::Tensor temp_coors =
at::zeros({num_points, NDim}, points.options().dtype(at::kInt));
dim3 grid(std::min(at::cuda::ATenCeilDiv(num_points, 512), 4096));
dim3 block(512);
// 1. link point to corresponding voxel coors
AT_DISPATCH_ALL_TYPES(
points.scalar_type(), "hard_voxelize_kernel", ([&] {
dynamic_voxelize_kernel<scalar_t, int>
<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
points.contiguous().data_ptr<scalar_t>(),
temp_coors.contiguous().data_ptr<int>(), voxel_x, voxel_y,
voxel_z, coors_x_min, coors_y_min, coors_z_min, coors_x_max,
coors_y_max, coors_z_max, grid_x, grid_y, grid_z, num_points,
num_features, NDim);
}));
cudaDeviceSynchronize();
AT_CUDA_CHECK(cudaGetLastError());
// 2. map point to the idx of the corresponding voxel, find duplicate coor
// create some temporary variables
auto point_to_pointidx = -at::ones(
{
num_points,
},
points.options().dtype(at::kInt));
auto point_to_voxelidx = -at::ones(
{
num_points,
},
points.options().dtype(at::kInt));
dim3 map_grid(std::min(at::cuda::ATenCeilDiv(num_points, 512), 4096));
dim3 map_block(512);
AT_DISPATCH_ALL_TYPES(
temp_coors.scalar_type(), "determin_duplicate", ([&] {
point_to_voxelidx_kernel<int>
<<<map_grid, map_block, 0, at::cuda::getCurrentCUDAStream()>>>(
temp_coors.contiguous().data_ptr<int>(),
point_to_voxelidx.contiguous().data_ptr<int>(),
point_to_pointidx.contiguous().data_ptr<int>(), max_points,
max_voxels, num_points, NDim);
}));
cudaDeviceSynchronize();
AT_CUDA_CHECK(cudaGetLastError());
// 3. determine voxel num and voxel's coor index
// make the logic in the CUDA device could accelerate about 10 times
auto coor_to_voxelidx = -at::ones(
{
num_points,
},
points.options().dtype(at::kInt));
auto voxel_num = at::zeros(
{
1,
},
points.options().dtype(at::kInt)); // must be zero from the beginning
AT_DISPATCH_ALL_TYPES(
temp_coors.scalar_type(), "determin_duplicate", ([&] {
determin_voxel_num<int><<<1, 1, 0, at::cuda::getCurrentCUDAStream()>>>(
num_points_per_voxel.contiguous().data_ptr<int>(),
point_to_voxelidx.contiguous().data_ptr<int>(),
point_to_pointidx.contiguous().data_ptr<int>(),
coor_to_voxelidx.contiguous().data_ptr<int>(),
voxel_num.contiguous().data_ptr<int>(), max_points, max_voxels,
num_points);
}));
cudaDeviceSynchronize();
AT_CUDA_CHECK(cudaGetLastError());
// 4. copy point features to voxels
// Step 4 & 5 could be parallel
auto pts_output_size = num_points * num_features;
dim3 cp_grid(std::min(at::cuda::ATenCeilDiv(pts_output_size, 512), 4096));
dim3 cp_block(512);
AT_DISPATCH_ALL_TYPES(
points.scalar_type(), "assign_point_to_voxel", ([&] {
assign_point_to_voxel<float, int>
<<<cp_grid, cp_block, 0, at::cuda::getCurrentCUDAStream()>>>(
pts_output_size, points.contiguous().data_ptr<float>(),
point_to_voxelidx.contiguous().data_ptr<int>(),
coor_to_voxelidx.contiguous().data_ptr<int>(),
voxels.contiguous().data_ptr<float>(), max_points, num_features,
num_points, NDim);
}));
// cudaDeviceSynchronize();
// AT_CUDA_CHECK(cudaGetLastError());
// 5. copy coors of each voxels
auto coors_output_size = num_points * NDim;
dim3 coors_cp_grid(
std::min(at::cuda::ATenCeilDiv(coors_output_size, 512), 4096));
dim3 coors_cp_block(512);
AT_DISPATCH_ALL_TYPES(
points.scalar_type(), "assign_point_to_voxel", ([&] {
assign_voxel_coors<float, int><<<coors_cp_grid, coors_cp_block, 0,
at::cuda::getCurrentCUDAStream()>>>(
coors_output_size, temp_coors.contiguous().data_ptr<int>(),
point_to_voxelidx.contiguous().data_ptr<int>(),
coor_to_voxelidx.contiguous().data_ptr<int>(),
coors.contiguous().data_ptr<int>(), num_points, NDim);
}));
cudaDeviceSynchronize();
AT_CUDA_CHECK(cudaGetLastError());
auto voxel_num_cpu = voxel_num.to(at::kCPU);
int voxel_num_int = voxel_num_cpu.data_ptr<int>()[0];
return voxel_num_int;
}
int nondisterministic_hard_voxelize_gpu(
const at::Tensor &points, at::Tensor &voxels,
at::Tensor &coors, at::Tensor &num_points_per_voxel,
const std::vector<float> voxel_size,
const std::vector<float> coors_range,
const int max_points, const int max_voxels,
const int NDim = 3) {
CHECK_INPUT(points);
at::cuda::CUDAGuard device_guard(points.device());
const int num_points = points.size(0);
const int num_features = points.size(1);
if (num_points == 0)
return 0;
const float voxel_x = voxel_size[0];
const float voxel_y = voxel_size[1];
const float voxel_z = voxel_size[2];
const float coors_x_min = coors_range[0];
const float coors_y_min = coors_range[1];
const float coors_z_min = coors_range[2];
const float coors_x_max = coors_range[3];
const float coors_y_max = coors_range[4];
const float coors_z_max = coors_range[5];
const int grid_x = round((coors_x_max - coors_x_min) / voxel_x);
const int grid_y = round((coors_y_max - coors_y_min) / voxel_y);
const int grid_z = round((coors_z_max - coors_z_min) / voxel_z);
// map points to voxel coors
at::Tensor temp_coors =
at::zeros({num_points, NDim}, points.options().dtype(torch::kInt32));
dim3 grid(std::min(at::cuda::ATenCeilDiv(num_points, 512), 4096));
dim3 block(512);
// 1. link point to corresponding voxel coors
AT_DISPATCH_ALL_TYPES(
points.scalar_type(), "hard_voxelize_kernel", ([&] {
dynamic_voxelize_kernel<scalar_t, int>
<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
points.contiguous().data_ptr<scalar_t>(),
temp_coors.contiguous().data_ptr<int>(), voxel_x, voxel_y,
voxel_z, coors_x_min, coors_y_min, coors_z_min, coors_x_max,
coors_y_max, coors_z_max, grid_x, grid_y, grid_z, num_points,
num_features, NDim);
}));
at::Tensor coors_map;
at::Tensor coors_count;
at::Tensor coors_order;
at::Tensor reduce_count;
at::Tensor pts_id;
auto coors_clean = temp_coors.masked_fill(temp_coors.lt(0).any(-1, true), -1);
std::tie(temp_coors, coors_map, reduce_count) =
at::unique_dim(coors_clean, 0, true, true, false);
if (temp_coors.index({0, 0}).lt(0).item<bool>()) {
// the first element of temp_coors is (-1,-1,-1) and should be removed
temp_coors = temp_coors.slice(0, 1);
coors_map = coors_map - 1;
}
int num_coors = temp_coors.size(0);
temp_coors = temp_coors.to(torch::kInt32);
coors_map = coors_map.to(torch::kInt32);
coors_count = coors_map.new_zeros(1);
coors_order = coors_map.new_empty(num_coors);
reduce_count = coors_map.new_zeros(num_coors);
pts_id = coors_map.new_zeros(num_points);
dim3 cp_grid(std::min(at::cuda::ATenCeilDiv(num_points, 512), 4096));
dim3 cp_block(512);
AT_DISPATCH_ALL_TYPES(points.scalar_type(), "get_assign_pos", ([&] {
nondisterministic_get_assign_pos<<<cp_grid, cp_block, 0,
at::cuda::getCurrentCUDAStream()>>>(
num_points,
coors_map.contiguous().data_ptr<int32_t>(),
pts_id.contiguous().data_ptr<int32_t>(),
coors_count.contiguous().data_ptr<int32_t>(),
reduce_count.contiguous().data_ptr<int32_t>(),
coors_order.contiguous().data_ptr<int32_t>());
}));
AT_DISPATCH_ALL_TYPES(
points.scalar_type(), "assign_point_to_voxel", ([&] {
nondisterministic_assign_point_voxel<scalar_t>
<<<cp_grid, cp_block, 0, at::cuda::getCurrentCUDAStream()>>>(
num_points, points.contiguous().data_ptr<scalar_t>(),
coors_map.contiguous().data_ptr<int32_t>(),
pts_id.contiguous().data_ptr<int32_t>(),
temp_coors.contiguous().data_ptr<int32_t>(),
reduce_count.contiguous().data_ptr<int32_t>(),
coors_order.contiguous().data_ptr<int32_t>(),
voxels.contiguous().data_ptr<scalar_t>(),
coors.contiguous().data_ptr<int32_t>(),
num_points_per_voxel.contiguous().data_ptr<int32_t>(),
max_voxels, max_points,
num_features, NDim);
}));
AT_CUDA_CHECK(cudaGetLastError());
return max_voxels < num_coors ? max_voxels : num_coors;
}
void dynamic_voxelize_gpu(const at::Tensor& points, at::Tensor& coors,
const std::vector<float> voxel_size,
const std::vector<float> coors_range,
const int NDim = 3) {
// current version tooks about 0.04s for one frame on cpu
// check device
CHECK_INPUT(points);
at::cuda::CUDAGuard device_guard(points.device());
const int num_points = points.size(0);
const int num_features = points.size(1);
const float voxel_x = voxel_size[0];
const float voxel_y = voxel_size[1];
const float voxel_z = voxel_size[2];
const float coors_x_min = coors_range[0];
const float coors_y_min = coors_range[1];
const float coors_z_min = coors_range[2];
const float coors_x_max = coors_range[3];
const float coors_y_max = coors_range[4];
const float coors_z_max = coors_range[5];
const int grid_x = round((coors_x_max - coors_x_min) / voxel_x);
const int grid_y = round((coors_y_max - coors_y_min) / voxel_y);
const int grid_z = round((coors_z_max - coors_z_min) / voxel_z);
const int col_blocks = at::cuda::ATenCeilDiv(num_points, threadsPerBlock);
dim3 blocks(col_blocks);
dim3 threads(threadsPerBlock);
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
AT_DISPATCH_ALL_TYPES(points.scalar_type(), "dynamic_voxelize_kernel", [&] {
dynamic_voxelize_kernel<scalar_t, int><<<blocks, threads, 0, stream>>>(
points.contiguous().data_ptr<scalar_t>(),
coors.contiguous().data_ptr<int>(), voxel_x, voxel_y, voxel_z,
coors_x_min, coors_y_min, coors_z_min, coors_x_max, coors_y_max,
coors_z_max, grid_x, grid_y, grid_z, num_points, num_features, NDim);
});
cudaDeviceSynchronize();
AT_CUDA_CHECK(cudaGetLastError());
return;
}
} // namespace voxelization
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import torch
from torch import nn
from torch.autograd import Function
from torch.nn.modules.utils import _pair
from .voxel_layer import dynamic_voxelize, hard_voxelize
class _Voxelization(Function):
@staticmethod
def forward(ctx,
points,
voxel_size,
coors_range,
max_points=35,
max_voxels=20000,
deterministic=True):
"""convert kitti points(N, >=3) to voxels.
Args:
points: [N, ndim] float tensor. points[:, :3] contain xyz points
and points[:, 3:] contain other information like reflectivity
voxel_size: [3] list/tuple or array, float. xyz, indicate voxel
size
coors_range: [6] list/tuple or array, float. indicate voxel
range. format: xyzxyz, minmax
max_points: int. indicate maximum points contained in a voxel. if
max_points=-1, it means using dynamic_voxelize
max_voxels: int. indicate maximum voxels this function create.
for second, 20000 is a good choice. Users should shuffle points
before call this function because max_voxels may drop points.
deterministic: bool. whether to invoke the non-deterministic
version of hard-voxelization implementations. non-deterministic
version is considerablly fast but is not deterministic. only
affects hard voxelization. default True. for more information
of this argument and the implementation insights, please refer
to the following links:
https://github.com/open-mmlab/mmdetection3d/issues/894
https://github.com/open-mmlab/mmdetection3d/pull/904
it is an experimental feature and we will appreciate it if
you could share with us the failing cases.
Returns:
voxels: [M, max_points, ndim] float tensor. only contain points
and returned when max_points != -1.
coordinates: [M, 3] int32 tensor, always returned.
num_points_per_voxel: [M] int32 tensor. Only returned when
max_points != -1.
"""
if max_points == -1 or max_voxels == -1:
coors = points.new_zeros(size=(points.size(0), 3), dtype=torch.int)
dynamic_voxelize(points, coors, voxel_size, coors_range, 3)
return coors
else:
voxels = points.new_zeros(
size=(max_voxels, max_points, points.size(1)))
coors = points.new_zeros(size=(max_voxels, 3), dtype=torch.int)
num_points_per_voxel = points.new_zeros(
size=(max_voxels, ), dtype=torch.int)
voxel_num = hard_voxelize(points, voxels, coors,
num_points_per_voxel, voxel_size,
coors_range, max_points, max_voxels, 3,
deterministic)
# select the valid voxels
voxels_out = voxels[:voxel_num]
coors_out = coors[:voxel_num]
num_points_per_voxel_out = num_points_per_voxel[:voxel_num]
return voxels_out, coors_out, num_points_per_voxel_out
voxelization = _Voxelization.apply
class Voxelization(nn.Module):
def __init__(self,
voxel_size,
point_cloud_range,
max_num_points,
max_voxels=20000,
deterministic=True):
super(Voxelization, self).__init__()
"""
Args:
voxel_size (list): list [x, y, z] size of three dimension
point_cloud_range (list):
[x_min, y_min, z_min, x_max, y_max, z_max]
max_num_points (int): max number of points per voxel
max_voxels (tuple or int): max number of voxels in
(training, testing) time
deterministic: bool. whether to invoke the non-deterministic
version of hard-voxelization implementations. non-deterministic
version is considerablly fast but is not deterministic. only
affects hard voxelization. default True. for more information
of this argument and the implementation insights, please refer
to the following links:
https://github.com/open-mmlab/mmdetection3d/issues/894
https://github.com/open-mmlab/mmdetection3d/pull/904
it is an experimental feature and we will appreciate it if
you could share with us the failing cases.
"""
self.voxel_size = voxel_size
self.point_cloud_range = point_cloud_range
self.max_num_points = max_num_points
if isinstance(max_voxels, tuple):
self.max_voxels = max_voxels
else:
self.max_voxels = _pair(max_voxels)
self.deterministic = deterministic
point_cloud_range = torch.tensor(
point_cloud_range, dtype=torch.float32)
# [0, -40, -3, 70.4, 40, 1]
voxel_size = torch.tensor(voxel_size, dtype=torch.float32)
grid_size = (point_cloud_range[3:] -
point_cloud_range[:3]) / voxel_size
grid_size = torch.round(grid_size).long()
input_feat_shape = grid_size[:2]
self.grid_size = grid_size
# the origin shape is as [x-len, y-len, z-len]
# [w, h, d] -> [d, h, w]
self.pcd_shape = [*input_feat_shape, 1][::-1]
def forward(self, input):
"""
Args:
input: NC points
"""
if self.training:
max_voxels = self.max_voxels[0]
else:
max_voxels = self.max_voxels[1]
return voxelization(input, self.voxel_size, self.point_cloud_range,
self.max_num_points, max_voxels,
self.deterministic)
def __repr__(self):
tmpstr = self.__class__.__name__ + '('
tmpstr += 'voxel_size=' + str(self.voxel_size)
tmpstr += ', point_cloud_range=' + str(self.point_cloud_range)
tmpstr += ', max_num_points=' + str(self.max_num_points)
tmpstr += ', max_voxels=' + str(self.max_voxels)
tmpstr += ', deterministic=' + str(self.deterministic)
tmpstr += ')'
return tmpstr
# Copyright (c) Open-MMLab. All rights reserved.
__version__ = '1.0.0rc0'
__version__ = '1.0.0rc1'
short_version = __version__
......
mmcv-full>=1.3.8,<=1.5.0
mmcv-full>=1.4.8,<=1.5.0
mmdet>=2.19.0,<=3.0.0
mmsegmentation>=0.20.0,<=1.0.0
mmcv
mmdet
mmcv>=1.4.8
mmdet>=2.19.0
mmsegmentation>=0.20.1
torch
torchvision
......@@ -224,97 +224,5 @@ if __name__ == '__main__':
'build': parse_requirements('requirements/build.txt'),
'optional': parse_requirements('requirements/optional.txt'),
},
ext_modules=[
make_cuda_ext(
name='sparse_conv_ext',
module='mmdet3d.ops.spconv',
extra_include_path=[
# PyTorch 1.5 uses ninjia, which requires absolute path
# of included files, relative path will cause failure.
os.path.abspath(
os.path.join(*'mmdet3d.ops.spconv'.split('.'),
'include/'))
],
sources=[
'src/all.cc',
'src/reordering.cc',
'src/reordering_cuda.cu',
'src/indice.cc',
'src/indice_cuda.cu',
'src/maxpool.cc',
'src/maxpool_cuda.cu',
],
extra_args=['-w', '-std=c++14']),
make_cuda_ext(
name='iou3d_cuda',
module='mmdet3d.ops.iou3d',
sources=[
'src/iou3d.cpp',
'src/iou3d_kernel.cu',
]),
make_cuda_ext(
name='voxel_layer',
module='mmdet3d.ops.voxel',
sources=[
'src/voxelization.cpp',
'src/scatter_points_cpu.cpp',
'src/scatter_points_cuda.cu',
'src/voxelization_cpu.cpp',
'src/voxelization_cuda.cu',
]),
make_cuda_ext(
name='roiaware_pool3d_ext',
module='mmdet3d.ops.roiaware_pool3d',
sources=[
'src/roiaware_pool3d.cpp',
'src/points_in_boxes_cpu.cpp',
],
sources_cuda=[
'src/roiaware_pool3d_kernel.cu',
'src/points_in_boxes_cuda.cu',
]),
make_cuda_ext(
name='roipoint_pool3d_ext',
module='mmdet3d.ops.roipoint_pool3d',
sources=['src/roipoint_pool3d.cpp'],
sources_cuda=['src/roipoint_pool3d_kernel.cu']),
make_cuda_ext(
name='ball_query_ext',
module='mmdet3d.ops.ball_query',
sources=['src/ball_query.cpp'],
sources_cuda=['src/ball_query_cuda.cu']),
make_cuda_ext(
name='knn_ext',
module='mmdet3d.ops.knn',
sources=['src/knn.cpp'],
sources_cuda=['src/knn_cuda.cu']),
make_cuda_ext(
name='assign_score_withk_ext',
module='mmdet3d.ops.paconv',
sources=['src/assign_score_withk.cpp'],
sources_cuda=['src/assign_score_withk_cuda.cu']),
make_cuda_ext(
name='group_points_ext',
module='mmdet3d.ops.group_points',
sources=['src/group_points.cpp'],
sources_cuda=['src/group_points_cuda.cu']),
make_cuda_ext(
name='interpolate_ext',
module='mmdet3d.ops.interpolate',
sources=['src/interpolate.cpp'],
sources_cuda=[
'src/three_interpolate_cuda.cu', 'src/three_nn_cuda.cu'
]),
make_cuda_ext(
name='furthest_point_sample_ext',
module='mmdet3d.ops.furthest_point_sample',
sources=['src/furthest_point_sample.cpp'],
sources_cuda=['src/furthest_point_sample_cuda.cu']),
make_cuda_ext(
name='gather_points_ext',
module='mmdet3d.ops.gather_points',
sources=['src/gather_points.cpp'],
sources_cuda=['src/gather_points_cuda.cu'])
],
cmdclass={'build_ext': BuildExtension},
zip_safe=False)
......@@ -250,9 +250,12 @@ def test_evaluate():
metric = ['mAP']
result = dict(boxes_3d=boxes_3d, labels_3d=labels_3d, scores_3d=scores_3d)
ap_dict = kitti_dataset.evaluate([result], metric)
assert np.isclose(ap_dict['KITTI/Overall_3D_easy'], 3.0303030303030307)
assert np.isclose(ap_dict['KITTI/Overall_3D_moderate'], 3.0303030303030307)
assert np.isclose(ap_dict['KITTI/Overall_3D_hard'], 3.0303030303030307)
assert np.isclose(ap_dict['KITTI/Overall_3D_AP11_easy'],
3.0303030303030307)
assert np.isclose(ap_dict['KITTI/Overall_3D_AP11_moderate'],
3.0303030303030307)
assert np.isclose(ap_dict['KITTI/Overall_3D_AP11_hard'],
3.0303030303030307)
def test_show():
......
......@@ -207,9 +207,11 @@ def test_evaluate():
metric = ['mAP']
ap_dict = kitti_dataset.evaluate(results, metric)
assert np.isclose(ap_dict['img_bbox/KITTI/Overall_3D_easy'], 3.0303)
assert np.isclose(ap_dict['img_bbox/KITTI/Overall_3D_moderate'], 6.0606)
assert np.isclose(ap_dict['img_bbox/KITTI/Overall_3D_hard'], 6.0606)
assert np.isclose(ap_dict['img_bbox2d/KITTI/Overall_2D_easy'], 3.0303)
assert np.isclose(ap_dict['img_bbox2d/KITTI/Overall_2D_moderate'], 6.0606)
assert np.isclose(ap_dict['img_bbox2d/KITTI/Overall_2D_hard'], 6.0606)
assert np.isclose(ap_dict['img_bbox/KITTI/Overall_3D_AP11_easy'], 3.0303)
assert np.isclose(ap_dict['img_bbox/KITTI/Overall_3D_AP11_moderate'],
6.0606)
assert np.isclose(ap_dict['img_bbox/KITTI/Overall_3D_AP11_hard'], 6.0606)
assert np.isclose(ap_dict['img_bbox2d/KITTI/Overall_2D_AP11_easy'], 3.0303)
assert np.isclose(ap_dict['img_bbox2d/KITTI/Overall_2D_AP11_moderate'],
6.0606)
assert np.isclose(ap_dict['img_bbox2d/KITTI/Overall_2D_AP11_hard'], 6.0606)
......@@ -118,12 +118,12 @@ def test_getitem():
assert lyft_dataset.CLASSES == ('car', 'pedestrian')
import tempfile
tmp_file = tempfile.NamedTemporaryFile()
with open(tmp_file.name, 'w') as f:
f.write('car\npedestrian\n')
with tempfile.TemporaryDirectory() as tmpdir:
path = tmpdir + 'classes.txt'
with open(path, 'w') as f:
f.write('car\npedestrian\n')
lyft_dataset = LyftDataset(
ann_file, None, root_path, classes=tmp_file.name)
lyft_dataset = LyftDataset(ann_file, None, root_path, classes=path)
assert lyft_dataset.CLASSES != original_classes
assert lyft_dataset.CLASSES == ['car', 'pedestrian']
......
......@@ -179,15 +179,16 @@ def test_seg_getitem():
# test load classes from file
import tempfile
tmp_file = tempfile.NamedTemporaryFile()
with open(tmp_file.name, 'w') as f:
f.write('beam\nwindow\n')
with tempfile.TemporaryDirectory() as tmpdir:
path = tmpdir + 'classes.txt'
with open(path, 'w') as f:
f.write('beam\nwindow\n')
s3dis_dataset = S3DISSegDataset(
data_root=root_path,
ann_files=ann_file,
pipeline=None,
classes=tmp_file.name,
classes=path,
scene_idxs=scene_idxs)
assert s3dis_dataset.CLASSES != original_classes
assert s3dis_dataset.CLASSES == ['beam', 'window']
......
......@@ -5,7 +5,8 @@ import numpy as np
import pytest
import torch
from mmdet3d.datasets import ScanNetDataset, ScanNetSegDataset
from mmdet3d.datasets import (ScanNetDataset, ScanNetInstanceSegDataset,
ScanNetSegDataset)
def test_getitem():
......@@ -111,12 +112,13 @@ def test_getitem():
# Test load classes from file
import tempfile
tmp_file = tempfile.NamedTemporaryFile()
with open(tmp_file.name, 'w') as f:
f.write('cabinet\nbed\n')
with tempfile.TemporaryDirectory() as tmpdir:
path = tmpdir + 'classes.txt'
with open(path, 'w') as f:
f.write('cabinet\nbed\n')
scannet_dataset = ScanNetDataset(
root_path, ann_file, pipeline=None, classes=tmp_file.name)
root_path, ann_file, pipeline=None, classes=path)
assert scannet_dataset.CLASSES != original_classes
assert scannet_dataset.CLASSES == ['cabinet', 'bed']
......@@ -496,15 +498,16 @@ def test_seg_getitem():
# test load classes from file
import tempfile
tmp_file = tempfile.NamedTemporaryFile()
with open(tmp_file.name, 'w') as f:
f.write('cabinet\nchair\n')
with tempfile.TemporaryDirectory() as tmpdir:
path = tmpdir + 'classes.txt'
with open(path, 'w') as f:
f.write('cabinet\nchair\n')
scannet_dataset = ScanNetSegDataset(
data_root=root_path,
ann_file=ann_file,
pipeline=None,
classes=tmp_file.name,
classes=path,
scene_idxs=scene_idxs)
assert scannet_dataset.CLASSES != original_classes
assert scannet_dataset.CLASSES == ['cabinet', 'chair']
......@@ -685,3 +688,155 @@ def test_seg_format_results():
expected_txt_path = osp.join(tmp_dir.name, 'results', 'scene0000_00.txt')
assert np.all(result_files[0]['seg_mask'] == expected_label)
mmcv.check_file_exist(expected_txt_path)
def test_instance_seg_getitem():
np.random.seed(0)
root_path = './tests/data/scannet/'
ann_file = './tests/data/scannet/scannet_infos.pkl'
class_names = ('cabinet', 'bed', 'chair', 'sofa', 'table', 'door',
'window', 'bookshelf', 'picture', 'counter', 'desk',
'curtain', 'refrigerator', 'showercurtrain', 'toilet',
'sink', 'bathtub', 'garbagebin')
train_pipeline = [
dict(
type='LoadPointsFromFile',
coord_type='DEPTH',
shift_height=False,
use_color=True,
load_dim=6,
use_dim=[0, 1, 2, 3, 4, 5]),
dict(
type='LoadAnnotations3D',
with_bbox_3d=False,
with_label_3d=False,
with_mask_3d=True,
with_seg_3d=True),
dict(
type='PointSegClassMapping',
valid_cat_ids=(3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33,
34, 36, 39),
max_cat_id=40),
dict(type='NormalizePointsColor', color_mean=None),
dict(type='DefaultFormatBundle3D', class_names=class_names),
dict(
type='Collect3D',
keys=['points', 'pts_semantic_mask', 'pts_instance_mask'])
]
scannet_dataset = ScanNetInstanceSegDataset(
data_root=root_path,
ann_file=ann_file,
pipeline=train_pipeline,
classes=class_names,
test_mode=False)
expected_points = torch.tensor([[
-3.4742e+00, 7.8792e-01, 1.7397e+00, 3.3725e-01, 3.5294e-01, 3.0588e-01
], [
2.7216e+00, 3.4164e+00, 2.4572e+00, 6.6275e-01, 6.2745e-01, 5.1373e-01
],
[
1.3404e+00, -1.4675e+00, -4.4059e-02,
3.8431e-01, 3.6078e-01, 3.5686e-01
],
[
-3.0335e+00, 2.7273e+00, 1.5181e+00,
2.3137e-01, 1.6078e-01, 8.2353e-02
],
[
-4.3207e-01, 1.8154e+00, 1.7455e-01,
4.0392e-01, 3.8039e-01, 4.1961e-01
]])
data = scannet_dataset[0]
points = data['points']._data[:5]
pts_semantic_mask = data['pts_semantic_mask']._data[:5]
pts_instance_mask = data['pts_instance_mask']._data[:5]
expected_semantic_mask = np.array([11, 18, 18, 0, 4])
expected_instance_mask = np.array([6, 56, 10, 9, 35])
assert torch.allclose(points, expected_points, 1e-2)
assert np.all(pts_semantic_mask.numpy() == expected_semantic_mask)
assert np.all(pts_instance_mask.numpy() == expected_instance_mask)
def test_instance_seg_evaluate():
root_path = './tests/data/scannet'
ann_file = './tests/data/scannet/scannet_infos.pkl'
class_names = ('cabinet', 'bed', 'chair', 'sofa', 'table', 'door',
'window', 'bookshelf', 'picture', 'counter', 'desk',
'curtain', 'refrigerator', 'showercurtrain', 'toilet',
'sink', 'bathtub', 'garbagebin')
test_pipeline = [
dict(
type='LoadPointsFromFile',
coord_type='DEPTH',
shift_height=False,
use_color=True,
load_dim=6,
use_dim=[0, 1, 2, 3, 4, 5]),
dict(type='NormalizePointsColor', color_mean=None),
dict(type='DefaultFormatBundle3D', class_names=class_names),
dict(type='Collect3D', keys=['points'])
]
scannet_dataset = ScanNetInstanceSegDataset(
data_root=root_path,
ann_file=ann_file,
pipeline=test_pipeline,
test_mode=True)
pred_mask = torch.tensor([
1, -1, -1, -1, 7, 11, 2, -1, 1, 10, -1, -1, 5, -1, -1, -1, -1, 1, -1,
-1, -1, -1, 0, -1, 1, -1, 12, -1, -1, -1, 8, 5, 1, 5, 2, -1, -1, -1,
-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 1, 8, -1, -1, -1,
0, 4, 3, -1, 9, -1, -1, 6, -1, -1, -1, -1, 13, -1, -1, 5, -1, 5, -1,
-1, 9, 0, 5, -1, -1, 2, 3, 4, -1, -1, -1, 2, -1, -1, -1, 5, 9, -1, 1,
-1, 4, 10, 4, -1
]).long()
pred_labels = torch.tensor(
[4, 11, 11, 10, 0, 3, 12, 4, 14, 1, 0, 0, 0, 5, 5]).long()
pred_scores = torch.tensor([.99 for _ in range(len(pred_labels))])
results = [
dict(
instance_mask=pred_mask,
instance_label=pred_labels,
instance_score=torch.tensor(pred_scores))
]
eval_pipeline = [
dict(
type='LoadPointsFromFile',
coord_type='DEPTH',
shift_height=False,
use_color=True,
load_dim=6,
use_dim=[0, 1, 2, 3, 4, 5]),
dict(
type='LoadAnnotations3D',
with_bbox_3d=False,
with_label_3d=False,
with_mask_3d=True,
with_seg_3d=True),
dict(
type='PointSegClassMapping',
valid_cat_ids=(3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33,
34, 36, 39),
max_cat_id=40),
dict(type='NormalizePointsColor', color_mean=None),
dict(type='DefaultFormatBundle3D', class_names=class_names),
dict(
type='Collect3D',
keys=['points', 'pts_semantic_mask', 'pts_instance_mask'])
]
# We add options here as default min_region_size
# is much bigger than test instances.
ret_dict = scannet_dataset.evaluate(
results,
pipeline=eval_pipeline,
options=dict(min_region_sizes=np.array([1])))
assert abs(ret_dict['all_ap'] - 0.90625) < 0.001
assert abs(ret_dict['all_ap_50%'] - 0.90625) < 0.001
assert abs(ret_dict['all_ap_25%'] - 0.94444) < 0.001
assert abs(ret_dict['classes']['cabinet']['ap25%'] - 1.0) < 0.001
assert abs(ret_dict['classes']['cabinet']['ap50%'] - 0.65625) < 0.001
assert abs(ret_dict['classes']['door']['ap25%'] - 0.5) < 0.001
assert abs(ret_dict['classes']['door']['ap50%'] - 0.5) < 0.001
......@@ -89,6 +89,9 @@ def _generate_sunrgbd_multi_modality_dataset_config():
def test_getitem():
from os import path as osp
np.random.seed(0)
root_path, ann_file, class_names, pipelines, modality = \
_generate_sunrgbd_dataset_config()
......@@ -107,7 +110,8 @@ def test_getitem():
pcd_rotation_expected = np.array([[0.99889565, 0.04698427, 0.],
[-0.04698427, 0.99889565, 0.],
[0., 0., 1.]])
assert file_name == './tests/data/sunrgbd/points/000001.bin'
expected_file_name = osp.join('./tests/data/sunrgbd', 'points/000001.bin')
assert file_name == expected_file_name
assert pcd_horizontal_flip is False
assert abs(pcd_scale_factor - 0.9770964398016714) < 1e-5
assert np.allclose(pcd_rotation, pcd_rotation_expected, 1e-3)
......@@ -142,12 +146,13 @@ def test_getitem():
assert SUNRGBD_dataset.CLASSES == ('bed', 'table')
import tempfile
tmp_file = tempfile.NamedTemporaryFile()
with open(tmp_file.name, 'w') as f:
f.write('bed\ntable\n')
with tempfile.TemporaryDirectory() as tmpdir:
path = tmpdir + 'classes.txt'
with open(path, 'w') as f:
f.write('bed\ntable\n')
SUNRGBD_dataset = SUNRGBDDataset(
root_path, ann_file, pipeline=None, classes=tmp_file.name)
root_path, ann_file, pipeline=None, classes=path)
assert SUNRGBD_dataset.CLASSES != original_classes
assert SUNRGBD_dataset.CLASSES == ['bed', 'table']
......
......@@ -27,7 +27,7 @@ def _generate_waymo_train_dataset_config():
points_loader=dict(
type='LoadPointsFromFile',
coord_type='LIDAR',
load_dim=5,
load_dim=6,
use_dim=[0, 1, 2, 3, 4],
file_client_args=file_client_args))
pipeline = [
......@@ -144,9 +144,12 @@ def test_evaluate():
# kitti protocol
metric = ['kitti']
ap_dict = waymo_dataset.evaluate([result], metric=metric)
assert np.isclose(ap_dict['KITTI/Overall_3D_easy'], 3.0303030303030307)
assert np.isclose(ap_dict['KITTI/Overall_3D_moderate'], 3.0303030303030307)
assert np.isclose(ap_dict['KITTI/Overall_3D_hard'], 3.0303030303030307)
assert np.isclose(ap_dict['KITTI/Overall_3D_AP11_easy'],
3.0303030303030307)
assert np.isclose(ap_dict['KITTI/Overall_3D_AP11_moderate'],
3.0303030303030307)
assert np.isclose(ap_dict['KITTI/Overall_3D_AP11_hard'],
3.0303030303030307)
# waymo protocol
metric = ['waymo']
......
......@@ -88,7 +88,7 @@ def test_object_sample():
gt_labels.append(CLASSES.index(cat))
else:
gt_labels.append(-1)
gt_labels = np.array(gt_labels, dtype=np.long)
gt_labels = np.array(gt_labels, dtype=np.int64)
points = LiDARPoints(points, points_dim=4)
input_dict = dict(
points=points, gt_bboxes_3d=gt_bboxes_3d, gt_labels_3d=gt_labels)
......@@ -176,7 +176,7 @@ def test_object_name_filter():
gt_labels.append(CLASSES.index(cat))
else:
gt_labels.append(-1)
gt_labels = np.array(gt_labels, dtype=np.long)
gt_labels = np.array(gt_labels, dtype=np.int64)
input_dict = dict(
gt_bboxes_3d=gt_bboxes_3d.clone(), gt_labels_3d=gt_labels.copy())
......@@ -200,9 +200,9 @@ def test_point_shuffle():
points = np.fromfile('tests/data/scannet/points/scene0000_00.bin',
np.float32).reshape(-1, 6)
ins_mask = np.fromfile('tests/data/scannet/instance_mask/scene0000_00.bin',
np.long)
np.int64)
sem_mask = np.fromfile('tests/data/scannet/semantic_mask/scene0000_00.bin',
np.long)
np.int64)
points = DepthPoints(
points.copy(), points_dim=6, attribute_dims=dict(color=[3, 4, 5]))
......@@ -244,9 +244,9 @@ def test_points_range_filter():
points = np.fromfile('tests/data/scannet/points/scene0000_00.bin',
np.float32).reshape(-1, 6)
ins_mask = np.fromfile('tests/data/scannet/instance_mask/scene0000_00.bin',
np.long)
np.int64)
sem_mask = np.fromfile('tests/data/scannet/semantic_mask/scene0000_00.bin',
np.long)
np.int64)
points = DepthPoints(
points.copy(), points_dim=6, attribute_dims=dict(color=[3, 4, 5]))
......@@ -286,7 +286,7 @@ def test_object_range_filter():
[18.7314, -18.559, 20.6547, 6.4800, 8.6000, 3.9200, -1.0100],
[3.7314, 42.559, -0.6547, 6.4800, 8.6000, 2.9200, 3.0100]])
gt_bboxes_3d = LiDARInstance3DBoxes(bbox, origin=(0.5, 0.5, 0.5))
gt_labels_3d = np.array([0, 2, 1, 1, 2, 0], dtype=np.long)
gt_labels_3d = np.array([0, 2, 1, 1, 2, 0], dtype=np.int64)
input_dict = dict(
gt_bboxes_3d=gt_bboxes_3d.clone(), gt_labels_3d=gt_labels_3d.copy())
......
......@@ -61,10 +61,10 @@ def test_scannet_pipeline():
if info['annos']['gt_num'] != 0:
scannet_gt_bboxes_3d = info['annos']['gt_boxes_upright_depth'].astype(
np.float32)
scannet_gt_labels_3d = info['annos']['class'].astype(np.long)
scannet_gt_labels_3d = info['annos']['class'].astype(np.int64)
else:
scannet_gt_bboxes_3d = np.zeros((1, 6), dtype=np.float32)
scannet_gt_labels_3d = np.zeros((1, ), dtype=np.long)
scannet_gt_labels_3d = np.zeros((1, ), dtype=np.int64)
results['ann_info'] = dict()
results['ann_info']['pts_instance_mask_path'] = osp.join(
data_path, info['pts_instance_mask_path'])
......@@ -294,10 +294,10 @@ def test_sunrgbd_pipeline():
if info['annos']['gt_num'] != 0:
gt_bboxes_3d = info['annos']['gt_boxes_upright_depth'].astype(
np.float32)
gt_labels_3d = info['annos']['class'].astype(np.long)
gt_labels_3d = info['annos']['class'].astype(np.int64)
else:
gt_bboxes_3d = np.zeros((1, 7), dtype=np.float32)
gt_labels_3d = np.zeros((1, ), dtype=np.long)
gt_labels_3d = np.zeros((1, ), dtype=np.int64)
# prepare input of pipeline
results['ann_info'] = dict()
......
......@@ -82,7 +82,7 @@ def test_indoor_seg_sample():
scannet_points, points_dim=6, attribute_dims=dict(color=[3, 4, 5]))
scannet_pts_semantic_mask = np.fromfile(
'./tests/data/scannet/semantic_mask/scene0000_00.bin', dtype=np.long)
'./tests/data/scannet/semantic_mask/scene0000_00.bin', dtype=np.int64)
scannet_results['pts_semantic_mask'] = scannet_pts_semantic_mask
scannet_results = scannet_seg_class_mapping(scannet_results)
......@@ -174,7 +174,7 @@ def test_indoor_seg_sample():
s3dis_points, points_dim=6, attribute_dims=dict(color=[3, 4, 5]))
s3dis_pts_semantic_mask = np.fromfile(
'./tests/data/s3dis/semantic_mask/Area_1_office_2.bin', dtype=np.long)
'./tests/data/s3dis/semantic_mask/Area_1_office_2.bin', dtype=np.int64)
s3dis_results['pts_semantic_mask'] = s3dis_pts_semantic_mask
s3dis_results = s3dis_patch_sample_points(s3dis_results)
......
# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
import torch
from mmdet3d.core import instance_seg_eval
def test_instance_seg_eval():
valid_class_ids = (3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34,
36, 39)
class_labels = ('cabinet', 'bed', 'chair', 'sofa', 'table', 'door',
'window', 'bookshelf', 'picture', 'counter', 'desk',
'curtain', 'refrigerator', 'showercurtrain', 'toilet',
'sink', 'bathtub', 'garbagebin')
n_points_list = [3300, 3000]
gt_labels_list = [[0, 0, 0, 0, 0, 0, 14, 14, 2, 1],
[13, 13, 2, 1, 3, 3, 0, 0, 0]]
gt_instance_masks = []
gt_semantic_masks = []
pred_instance_masks = []
pred_instance_labels = []
pred_instance_scores = []
for n_points, gt_labels in zip(n_points_list, gt_labels_list):
gt_instance_mask = np.ones(n_points, dtype=np.int) * -1
gt_semantic_mask = np.ones(n_points, dtype=np.int) * -1
pred_instance_mask = np.ones(n_points, dtype=np.int) * -1
labels = []
scores = []
for i, gt_label in enumerate(gt_labels):
begin = i * 300
end = begin + 300
gt_instance_mask[begin:end] = i
gt_semantic_mask[begin:end] = gt_label
pred_instance_mask[begin:end] = i
labels.append(gt_label)
scores.append(.99)
gt_instance_masks.append(torch.tensor(gt_instance_mask))
gt_semantic_masks.append(torch.tensor(gt_semantic_mask))
pred_instance_masks.append(torch.tensor(pred_instance_mask))
pred_instance_labels.append(torch.tensor(labels))
pred_instance_scores.append(torch.tensor(scores))
ret_value = instance_seg_eval(
gt_semantic_masks=gt_semantic_masks,
gt_instance_masks=gt_instance_masks,
pred_instance_masks=pred_instance_masks,
pred_instance_labels=pred_instance_labels,
pred_instance_scores=pred_instance_scores,
valid_class_ids=valid_class_ids,
class_labels=class_labels)
for label in [
'cabinet', 'bed', 'chair', 'sofa', 'showercurtrain', 'toilet'
]:
metrics = ret_value['classes'][label]
assert metrics['ap'] == 1.0
assert metrics['ap50%'] == 1.0
assert metrics['ap25%'] == 1.0
pred_instance_masks[1][2240:2700] = -1
pred_instance_masks[0][2700:3000] = 8
pred_instance_labels[0][9] = 2
ret_value = instance_seg_eval(
gt_semantic_masks=gt_semantic_masks,
gt_instance_masks=gt_instance_masks,
pred_instance_masks=pred_instance_masks,
pred_instance_labels=pred_instance_labels,
pred_instance_scores=pred_instance_scores,
valid_class_ids=valid_class_ids,
class_labels=class_labels)
assert abs(ret_value['classes']['cabinet']['ap50%'] - 0.72916) < 0.01
assert abs(ret_value['classes']['cabinet']['ap25%'] - 0.88888) < 0.01
assert abs(ret_value['classes']['bed']['ap50%'] - 0.5) < 0.01
assert abs(ret_value['classes']['bed']['ap25%'] - 0.5) < 0.01
assert abs(ret_value['classes']['chair']['ap50%'] - 0.375) < 0.01
assert abs(ret_value['classes']['chair']['ap25%'] - 1.0) < 0.01
......@@ -12,11 +12,11 @@ def test_dgcnn_gf_module():
self = DGCNNGFModule(
mlp_channels=[18, 64, 64],
num_sample=20,
knn_mod='D-KNN',
knn_mode='D-KNN',
radius=None,
norm_cfg=dict(type='BN2d'),
act_cfg=dict(type='ReLU'),
pool_mod='max').cuda()
pool_mode='max').cuda()
assert self.mlps[0].layer0.conv.in_channels == 18
assert self.mlps[0].layer0.conv.out_channels == 64
......@@ -36,11 +36,11 @@ def test_dgcnn_gf_module():
self = DGCNNGFModule(
mlp_channels=[6, 64, 64],
num_sample=20,
knn_mod='F-KNN',
knn_mode='F-KNN',
radius=None,
norm_cfg=dict(type='BN2d'),
act_cfg=dict(type='ReLU'),
pool_mod='max').cuda()
pool_mode='max').cuda()
# test forward
new_points = self(xyz)
......@@ -50,11 +50,11 @@ def test_dgcnn_gf_module():
self = DGCNNGFModule(
mlp_channels=[6, 64, 64],
num_sample=20,
knn_mod='F-KNN',
knn_mode='F-KNN',
radius=0.2,
norm_cfg=dict(type='BN2d'),
act_cfg=dict(type='ReLU'),
pool_mod='max').cuda()
pool_mode='max').cuda()
def test_dgcnn_fa_module():
......
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