├── .gitignore ├── model └── car_model │ └── netR_36.pth ├── pointnet2 ├── _ext-src │ ├── include │ │ ├── ball_query.h │ │ ├── group_points.h │ │ ├── ball_query_score.h │ │ ├── sampling.h │ │ ├── interpolate.h │ │ ├── utils.h │ │ └── cuda_utils.h │ └── src │ │ ├── bindings.cpp │ │ ├── ball_query.cpp │ │ ├── ball_query_score.cpp │ │ ├── ball_query_gpu.cu │ │ ├── ball_query_score_gpu.cu │ │ ├── group_points.cpp │ │ ├── group_points_gpu.cu │ │ ├── sampling.cpp │ │ ├── interpolate.cpp │ │ ├── interpolate_gpu.cu │ │ └── sampling_gpu.cu ├── models │ ├── __init__.py │ └── pointnet_tracking.py ├── utils │ ├── __init__.py │ ├── linalg_utils.py │ ├── pointnet2_modules.py │ └── pointnet2_utils.py └── __init__.py ├── requirements.txt ├── setup.py ├── README.md ├── metrics.py ├── searchspace.py ├── test_tracking.py ├── data_classes.py ├── train_tracking.py ├── Dataset.py └── kitty_utils.py /.gitignore: -------------------------------------------------------------------------------- 1 | pointnet2 2 | model 3 | build 4 | *.pyc -------------------------------------------------------------------------------- /model/car_model/netR_36.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HaozheQi/P2B/HEAD/model/car_model/netR_36.pth -------------------------------------------------------------------------------- /pointnet2/_ext-src/include/ball_query.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | 4 | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, 5 | const int nsample); 6 | -------------------------------------------------------------------------------- /pointnet2/models/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import ( 2 | division, 3 | absolute_import, 4 | with_statement, 5 | print_function, 6 | unicode_literals, 7 | ) 8 | from .pointnet_tracking import Pointnet_Tracking -------------------------------------------------------------------------------- /pointnet2/_ext-src/include/group_points.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | 4 | at::Tensor group_points(at::Tensor points, at::Tensor idx); 5 | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); 6 | -------------------------------------------------------------------------------- /pointnet2/_ext-src/include/ball_query_score.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | 4 | at::Tensor ball_query_score(at::Tensor new_xyz, at::Tensor xyz, at::Tensor score, const float radius, 5 | const int nsample); -------------------------------------------------------------------------------- /pointnet2/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import ( 2 | division, 3 | absolute_import, 4 | with_statement, 5 | print_function, 6 | unicode_literals, 7 | ) 8 | from . import pointnet2_utils 9 | from . import pointnet2_modules 10 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | git+https://github.com/v-wewei/etw_pytorch_utils.git@v1.1.1#egg=etw_pytorch_utils 2 | h5py 3 | numpy 4 | torch==1.4 5 | torchvision==0.5 6 | enum34 7 | future 8 | pandas 9 | shapely 10 | matplotlib 11 | pomegranate 12 | ipykernel 13 | jupyter 14 | imageio 15 | pyquaternion 16 | -------------------------------------------------------------------------------- /pointnet2/_ext-src/include/sampling.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | 4 | at::Tensor gather_points(at::Tensor points, at::Tensor idx); 5 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); 6 | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples); 7 | -------------------------------------------------------------------------------- /pointnet2/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import ( 2 | division, 3 | absolute_import, 4 | with_statement, 5 | print_function, 6 | unicode_literals, 7 | ) 8 | 9 | __version__ = "2.1.1" 10 | 11 | try: 12 | __POINTNET2_SETUP__ 13 | except NameError: 14 | __POINTNET2_SETUP__ = False 15 | 16 | if not __POINTNET2_SETUP__: 17 | from pointnet2 import utils 18 | from pointnet2 import models 19 | -------------------------------------------------------------------------------- /pointnet2/_ext-src/include/interpolate.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | 6 | std::vector three_nn(at::Tensor unknowns, at::Tensor knows); 7 | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, 8 | at::Tensor weight); 9 | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, 10 | at::Tensor weight, const int m); 11 | -------------------------------------------------------------------------------- /pointnet2/_ext-src/src/bindings.cpp: -------------------------------------------------------------------------------- 1 | #include "ball_query.h" 2 | #include "group_points.h" 3 | #include "interpolate.h" 4 | #include "sampling.h" 5 | #include "ball_query_score.h" 6 | 7 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 8 | m.def("gather_points", &gather_points); 9 | m.def("gather_points_grad", &gather_points_grad); 10 | m.def("furthest_point_sampling", &furthest_point_sampling); 11 | 12 | m.def("three_nn", &three_nn); 13 | m.def("three_interpolate", &three_interpolate); 14 | m.def("three_interpolate_grad", &three_interpolate_grad); 15 | 16 | m.def("ball_query", &ball_query); 17 | m.def("ball_query_score", &ball_query_score); 18 | 19 | m.def("group_points", &group_points); 20 | m.def("group_points_grad", &group_points_grad); 21 | } 22 | -------------------------------------------------------------------------------- /pointnet2/_ext-src/include/utils.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | #include 4 | 5 | #define CHECK_CUDA(x) \ 6 | do { \ 7 | AT_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor"); \ 8 | } while (0) 9 | 10 | #define CHECK_CONTIGUOUS(x) \ 11 | do { \ 12 | AT_CHECK(x.is_contiguous(), #x " must be a contiguous tensor"); \ 13 | } while (0) 14 | 15 | #define CHECK_IS_INT(x) \ 16 | do { \ 17 | AT_CHECK(x.scalar_type() == at::ScalarType::Int, \ 18 | #x " must be an int tensor"); \ 19 | } while (0) 20 | 21 | #define CHECK_IS_FLOAT(x) \ 22 | do { \ 23 | AT_CHECK(x.scalar_type() == at::ScalarType::Float, \ 24 | #x " must be a float tensor"); \ 25 | } while (0) 26 | -------------------------------------------------------------------------------- /pointnet2/_ext-src/src/ball_query.cpp: -------------------------------------------------------------------------------- 1 | #include "ball_query.h" 2 | #include "utils.h" 3 | 4 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 5 | int nsample, const float *new_xyz, 6 | const float *xyz, int *idx); 7 | 8 | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, 9 | const int nsample) { 10 | CHECK_CONTIGUOUS(new_xyz); 11 | CHECK_CONTIGUOUS(xyz); 12 | CHECK_IS_FLOAT(new_xyz); 13 | CHECK_IS_FLOAT(xyz); 14 | 15 | if (new_xyz.type().is_cuda()) { 16 | CHECK_CUDA(xyz); 17 | } 18 | 19 | at::Tensor idx = 20 | torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample}, 21 | at::device(new_xyz.device()).dtype(at::ScalarType::Int)); 22 | 23 | if (new_xyz.type().is_cuda()) { 24 | query_ball_point_kernel_wrapper(xyz.size(0), xyz.size(1), new_xyz.size(1), 25 | radius, nsample, new_xyz.data(), 26 | xyz.data(), idx.data()); 27 | } else { 28 | AT_CHECK(false, "CPU not supported"); 29 | } 30 | 31 | return idx; 32 | } 33 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, absolute_import, with_statement, print_function 2 | from setuptools import setup, find_packages 3 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 4 | import glob 5 | 6 | try: 7 | import builtins 8 | except: 9 | import __builtin__ as builtins 10 | 11 | builtins.__POINTNET2_SETUP__ = True 12 | import pointnet2 13 | 14 | _ext_src_root = "pointnet2/_ext-src" 15 | _ext_sources = glob.glob("{}/src/*.cpp".format(_ext_src_root)) + glob.glob( 16 | "{}/src/*.cu".format(_ext_src_root) 17 | ) 18 | _ext_headers = glob.glob("{}/include/*".format(_ext_src_root)) 19 | 20 | requirements = ["etw_pytorch_utils==1.1.1", "h5py", "pprint", "enum34", "future"] 21 | 22 | setup( 23 | name="pointnet2", 24 | version=pointnet2.__version__, 25 | author="Erik Wijmans", 26 | packages=find_packages(), 27 | install_requires=requirements, 28 | ext_modules=[ 29 | CUDAExtension( 30 | name="pointnet2._ext", 31 | sources=_ext_sources, 32 | extra_compile_args={ 33 | "cxx": ["-O2", "-I{}".format("{}/include".format(_ext_src_root))], 34 | "nvcc": ["-O2", "-I{}".format("{}/include".format(_ext_src_root))], 35 | }, 36 | ) 37 | ], 38 | cmdclass={"build_ext": BuildExtension}, 39 | ) 40 | -------------------------------------------------------------------------------- /pointnet2/_ext-src/src/ball_query_score.cpp: -------------------------------------------------------------------------------- 1 | #include "ball_query_score.h" 2 | #include "utils.h" 3 | 4 | void query_ball_point_score_kernel_wrapper(int b, int n, int m, float radius, 5 | int nsample, const float *new_xyz, 6 | const float *xyz, const float *score, float *unique_score); 7 | 8 | 9 | at::Tensor ball_query_score(at::Tensor new_xyz, at::Tensor xyz, at::Tensor score, const float radius, 10 | const int nsample) { 11 | CHECK_CONTIGUOUS(new_xyz); 12 | CHECK_CONTIGUOUS(xyz); 13 | CHECK_CONTIGUOUS(score); 14 | CHECK_IS_FLOAT(new_xyz); 15 | CHECK_IS_FLOAT(xyz); 16 | CHECK_IS_FLOAT(score); 17 | 18 | if (new_xyz.type().is_cuda()) { 19 | CHECK_CUDA(xyz); 20 | CHECK_CUDA(score); 21 | } 22 | 23 | at::Tensor unique_score = 24 | torch::zeros({new_xyz.size(0), new_xyz.size(1)}, 25 | at::device(new_xyz.device()).dtype(at::ScalarType::Float)); 26 | 27 | if (new_xyz.type().is_cuda()) { 28 | query_ball_point_score_kernel_wrapper(xyz.size(0), xyz.size(1), new_xyz.size(1), 29 | radius, nsample, new_xyz.data(), 30 | xyz.data(), score.data(), unique_score.data()); 31 | } else { 32 | AT_CHECK(false, "CPU not supported"); 33 | } 34 | 35 | return unique_score; 36 | } 37 | -------------------------------------------------------------------------------- /pointnet2/_ext-src/include/cuda_utils.h: -------------------------------------------------------------------------------- 1 | #ifndef _CUDA_UTILS_H 2 | #define _CUDA_UTILS_H 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #include 9 | #include 10 | 11 | #include 12 | 13 | #define TOTAL_THREADS 512 14 | 15 | inline int opt_n_threads(int work_size) { 16 | const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); 17 | 18 | return max(min(1 << pow_2, TOTAL_THREADS), 1); 19 | } 20 | 21 | inline dim3 opt_block_config(int x, int y) { 22 | const int x_threads = opt_n_threads(x); 23 | const int y_threads = 24 | max(min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1); 25 | dim3 block_config(x_threads, y_threads, 1); 26 | 27 | return block_config; 28 | } 29 | 30 | #define CUDA_CHECK_ERRORS() \ 31 | do { \ 32 | cudaError_t err = cudaGetLastError(); \ 33 | if (cudaSuccess != err) { \ 34 | fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n", \ 35 | cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \ 36 | __FILE__); \ 37 | exit(-1); \ 38 | } \ 39 | } while (0) 40 | 41 | #endif 42 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # P2B: Point-to-Box Network for 3D Object Tracking in Point Clouds 2 | 3 | ## Introduction 4 | 5 | This repository is released for P2B in our [CVPR 2020 paper (oral)](http://arxiv.org/abs/2005.13888). Here we include our P2B model (PyTorch) and code for data preparation, training and testing on KITTI tracking dataset. 6 | 7 | ## Preliminary 8 | 9 | * Install ``python 3.6``. 10 | 11 | * Install dependencies. 12 | ``` 13 | pip install -r requirements.txt 14 | ``` 15 | 16 | * Build `_ext` module. 17 | ``` 18 | python setup.py build_ext --inplace 19 | ``` 20 | 21 | * Download the dataset from [KITTI Tracking](http://www.cvlibs.net/datasets/kitti/eval_tracking.php). 22 | 23 | Download [velodyne](http://www.cvlibs.net/download.php?file=data_tracking_velodyne.zip), [calib](http://www.cvlibs.net/download.php?file=data_tracking_calib.zip) and [label_02](http://www.cvlibs.net/download.php?file=data_tracking_label_2.zip) in the dataset and place them under the same parent folder. 24 | 25 | ## Evaluation 26 | 27 | Train a new P2B model on KITTI data: 28 | ``` 29 | python train_tracking.py --data_dir= 30 | ``` 31 | 32 | Test a new P2B model on KITTI data: 33 | ``` 34 | python test_tracking.py --data_dir= 35 | ``` 36 | 37 | Please refer to the code for setting of other optional arguments, including data split, training and testing parameters, etc. 38 | 39 | ## Acknowledgements 40 | 41 | Thank Giancola for his implementation of [SC3D](https://github.com/SilvioGiancola/ShapeCompletion3DTracking). 42 | Thank Erik Wijmans for his implementation of [PointNet++](https://github.com/erikwijmans/Pointnet2_PyTorch) in PyTorch. 43 | Thank Charles R. Qi for his implementation of [Votenet](https://github.com/facebookresearch/votenet). 44 | They help and inspire this work. 45 | -------------------------------------------------------------------------------- /pointnet2/_ext-src/src/ball_query_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include "cuda_utils.h" 6 | 7 | // input: new_xyz(b, m, 3) xyz(b, n, 3) 8 | // output: idx(b, m, nsample) 9 | __global__ void query_ball_point_kernel(int b, int n, int m, float radius, 10 | int nsample, 11 | const float *__restrict__ new_xyz, 12 | const float *__restrict__ xyz, 13 | int *__restrict__ idx) { 14 | int batch_index = blockIdx.x; 15 | xyz += batch_index * n * 3; 16 | new_xyz += batch_index * m * 3; 17 | idx += m * nsample * batch_index; 18 | 19 | int index = threadIdx.x; 20 | int stride = blockDim.x; 21 | 22 | float radius2 = radius * radius; 23 | for (int j = index; j < m; j += stride) { 24 | float new_x = new_xyz[j * 3 + 0]; 25 | float new_y = new_xyz[j * 3 + 1]; 26 | float new_z = new_xyz[j * 3 + 2]; 27 | for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) { 28 | float x = xyz[k * 3 + 0]; 29 | float y = xyz[k * 3 + 1]; 30 | float z = xyz[k * 3 + 2]; 31 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + 32 | (new_z - z) * (new_z - z); 33 | if (d2 < radius2) { 34 | if (cnt == 0) { 35 | for (int l = 0; l < nsample; ++l) { 36 | idx[j * nsample + l] = k; 37 | } 38 | } 39 | idx[j * nsample + cnt] = k; 40 | ++cnt; 41 | } 42 | } 43 | } 44 | } 45 | 46 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 47 | int nsample, const float *new_xyz, 48 | const float *xyz, int *idx) { 49 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 50 | query_ball_point_kernel<<>>( 51 | b, n, m, radius, nsample, new_xyz, xyz, idx); 52 | 53 | CUDA_CHECK_ERRORS(); 54 | } 55 | -------------------------------------------------------------------------------- /pointnet2/_ext-src/src/ball_query_score_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include "cuda_utils.h" 6 | 7 | // input: new_xyz(b, m, 3) xyz(b, n, 3) 8 | // output: idx(b, m, nsample) 9 | __global__ void query_ball_point_score_kernel(int b, int n, int m, float radius, 10 | int nsample, 11 | const float *__restrict__ new_xyz, 12 | const float *__restrict__ xyz, 13 | const float *__restrict__ score, 14 | float *__restrict__ unique_score) { 15 | int batch_index = blockIdx.x; 16 | xyz += batch_index * n * 3; 17 | new_xyz += batch_index * m * 3; 18 | score += batch_index * n; 19 | unique_score += m * batch_index; 20 | 21 | int index = threadIdx.x; 22 | int stride = blockDim.x; 23 | 24 | float radius2 = radius * radius; 25 | for (int j = index; j < m; j += stride) { 26 | float new_x = new_xyz[j * 3 + 0]; 27 | float new_y = new_xyz[j * 3 + 1]; 28 | float new_z = new_xyz[j * 3 + 2]; 29 | for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) { 30 | float x = xyz[k * 3 + 0]; 31 | float y = xyz[k * 3 + 1]; 32 | float z = xyz[k * 3 + 2]; 33 | float s = score[k]; 34 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + 35 | (new_z - z) * (new_z - z); 36 | if (d2 < radius2) { 37 | unique_score[j] += s; 38 | ++cnt; 39 | } 40 | } 41 | } 42 | } 43 | 44 | void query_ball_point_score_kernel_wrapper(int b, int n, int m, float radius, 45 | int nsample, const float *new_xyz, 46 | const float *xyz, const float *score, float *unique_score) { 47 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 48 | query_ball_point_score_kernel<<>>( 49 | b, n, m, radius, nsample, new_xyz, xyz, score, unique_score); 50 | 51 | CUDA_CHECK_ERRORS(); 52 | } 53 | -------------------------------------------------------------------------------- /pointnet2/_ext-src/src/group_points.cpp: -------------------------------------------------------------------------------- 1 | #include "group_points.h" 2 | #include "utils.h" 3 | 4 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 5 | const float *points, const int *idx, 6 | float *out); 7 | 8 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 9 | int nsample, const float *grad_out, 10 | const int *idx, float *grad_points); 11 | 12 | at::Tensor group_points(at::Tensor points, at::Tensor idx) { 13 | CHECK_CONTIGUOUS(points); 14 | CHECK_CONTIGUOUS(idx); 15 | CHECK_IS_FLOAT(points); 16 | CHECK_IS_INT(idx); 17 | 18 | if (points.type().is_cuda()) { 19 | CHECK_CUDA(idx); 20 | } 21 | 22 | at::Tensor output = 23 | torch::zeros({points.size(0), points.size(1), idx.size(1), idx.size(2)}, 24 | at::device(points.device()).dtype(at::ScalarType::Float)); 25 | 26 | if (points.type().is_cuda()) { 27 | group_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 28 | idx.size(1), idx.size(2), points.data(), 29 | idx.data(), output.data()); 30 | } else { 31 | AT_CHECK(false, "CPU not supported"); 32 | } 33 | 34 | return output; 35 | } 36 | 37 | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n) { 38 | CHECK_CONTIGUOUS(grad_out); 39 | CHECK_CONTIGUOUS(idx); 40 | CHECK_IS_FLOAT(grad_out); 41 | CHECK_IS_INT(idx); 42 | 43 | if (grad_out.type().is_cuda()) { 44 | CHECK_CUDA(idx); 45 | } 46 | 47 | at::Tensor output = 48 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 49 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 50 | 51 | if (grad_out.type().is_cuda()) { 52 | group_points_grad_kernel_wrapper( 53 | grad_out.size(0), grad_out.size(1), n, idx.size(1), idx.size(2), 54 | grad_out.data(), idx.data(), output.data()); 55 | } else { 56 | AT_CHECK(false, "CPU not supported"); 57 | } 58 | 59 | return output; 60 | } 61 | -------------------------------------------------------------------------------- /pointnet2/utils/linalg_utils.py: -------------------------------------------------------------------------------- 1 | from __future__ import ( 2 | division, 3 | absolute_import, 4 | with_statement, 5 | print_function, 6 | unicode_literals, 7 | ) 8 | import torch 9 | from enum import Enum 10 | import numpy as np 11 | 12 | PDist2Order = Enum("PDist2Order", "d_first d_second") 13 | 14 | 15 | def pdist2(X, Z=None, order=PDist2Order.d_second): 16 | # type: (torch.Tensor, torch.Tensor, PDist2Order) -> torch.Tensor 17 | r""" Calculates the pairwise distance between X and Z 18 | 19 | D[b, i, j] = l2 distance X[b, i] and Z[b, j] 20 | 21 | Parameters 22 | --------- 23 | X : torch.Tensor 24 | X is a (B, N, d) tensor. There are B batches, and N vectors of dimension d 25 | Z: torch.Tensor 26 | Z is a (B, M, d) tensor. If Z is None, then Z = X 27 | 28 | Returns 29 | ------- 30 | torch.Tensor 31 | Distance matrix is size (B, N, M) 32 | """ 33 | 34 | if order == PDist2Order.d_second: 35 | if X.dim() == 2: 36 | X = X.unsqueeze(0) 37 | if Z is None: 38 | Z = X 39 | G = np.matmul(X, Z.transpose(-2, -1)) 40 | S = (X * X).sum(-1, keepdim=True) 41 | R = S.transpose(-2, -1) 42 | else: 43 | if Z.dim() == 2: 44 | Z = Z.unsqueeze(0) 45 | G = np.matmul(X, Z.transpose(-2, -1)) 46 | S = (X * X).sum(-1, keepdim=True) 47 | R = (Z * Z).sum(-1, keepdim=True).transpose(-2, -1) 48 | else: 49 | if X.dim() == 2: 50 | X = X.unsqueeze(0) 51 | if Z is None: 52 | Z = X 53 | G = np.matmul(X.transpose(-2, -1), Z) 54 | R = (X * X).sum(-2, keepdim=True) 55 | S = R.transpose(-2, -1) 56 | else: 57 | if Z.dim() == 2: 58 | Z = Z.unsqueeze(0) 59 | G = np.matmul(X.transpose(-2, -1), Z) 60 | S = (X * X).sum(-2, keepdim=True).transpose(-2, -1) 61 | R = (Z * Z).sum(-2, keepdim=True) 62 | 63 | return torch.abs(R + S - 2 * G).squeeze(0) 64 | 65 | 66 | def pdist2_slow(X, Z=None): 67 | if Z is None: 68 | Z = X 69 | D = torch.zeros(X.size(0), X.size(2), Z.size(2)) 70 | 71 | for b in range(D.size(0)): 72 | for i in range(D.size(1)): 73 | for j in range(D.size(2)): 74 | D[b, i, j] = torch.dist(X[b, :, i], Z[b, :, j]) 75 | return D 76 | 77 | 78 | if __name__ == "__main__": 79 | X = torch.randn(2, 3, 5) 80 | Z = torch.randn(2, 3, 3) 81 | 82 | print(pdist2(X, order=PDist2Order.d_first)) 83 | print(pdist2_slow(X)) 84 | print(torch.dist(pdist2(X, order=PDist2Order.d_first), pdist2_slow(X))) 85 | -------------------------------------------------------------------------------- /pointnet2/_ext-src/src/group_points_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "cuda_utils.h" 5 | 6 | // input: points(b, c, n) idx(b, npoints, nsample) 7 | // output: out(b, c, npoints, nsample) 8 | __global__ void group_points_kernel(int b, int c, int n, int npoints, 9 | int nsample, 10 | const float *__restrict__ points, 11 | const int *__restrict__ idx, 12 | float *__restrict__ out) { 13 | int batch_index = blockIdx.x; 14 | points += batch_index * n * c; 15 | idx += batch_index * npoints * nsample; 16 | out += batch_index * npoints * nsample * c; 17 | 18 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 19 | const int stride = blockDim.y * blockDim.x; 20 | for (int i = index; i < c * npoints; i += stride) { 21 | const int l = i / npoints; 22 | const int j = i % npoints; 23 | for (int k = 0; k < nsample; ++k) { 24 | int ii = idx[j * nsample + k]; 25 | out[(l * npoints + j) * nsample + k] = points[l * n + ii]; 26 | } 27 | } 28 | } 29 | 30 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 31 | const float *points, const int *idx, 32 | float *out) { 33 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 34 | 35 | group_points_kernel<<>>( 36 | b, c, n, npoints, nsample, points, idx, out); 37 | 38 | CUDA_CHECK_ERRORS(); 39 | } 40 | 41 | // input: grad_out(b, c, npoints, nsample), idx(b, npoints, nsample) 42 | // output: grad_points(b, c, n) 43 | __global__ void group_points_grad_kernel(int b, int c, int n, int npoints, 44 | int nsample, 45 | const float *__restrict__ grad_out, 46 | const int *__restrict__ idx, 47 | float *__restrict__ grad_points) { 48 | int batch_index = blockIdx.x; 49 | grad_out += batch_index * npoints * nsample * c; 50 | idx += batch_index * npoints * nsample; 51 | grad_points += batch_index * n * c; 52 | 53 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 54 | const int stride = blockDim.y * blockDim.x; 55 | for (int i = index; i < c * npoints; i += stride) { 56 | const int l = i / npoints; 57 | const int j = i % npoints; 58 | for (int k = 0; k < nsample; ++k) { 59 | int ii = idx[j * nsample + k]; 60 | atomicAdd(grad_points + l * n + ii, 61 | grad_out[(l * npoints + j) * nsample + k]); 62 | } 63 | } 64 | } 65 | 66 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 67 | int nsample, const float *grad_out, 68 | const int *idx, float *grad_points) { 69 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 70 | 71 | group_points_grad_kernel<<>>( 72 | b, c, n, npoints, nsample, grad_out, idx, grad_points); 73 | 74 | CUDA_CHECK_ERRORS(); 75 | } 76 | -------------------------------------------------------------------------------- /pointnet2/_ext-src/src/sampling.cpp: -------------------------------------------------------------------------------- 1 | #include "sampling.h" 2 | #include "utils.h" 3 | 4 | void gather_points_kernel_wrapper(int b, int c, int n, int npoints, 5 | const float *points, const int *idx, 6 | float *out); 7 | void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 8 | const float *grad_out, const int *idx, 9 | float *grad_points); 10 | 11 | void furthest_point_sampling_kernel_wrapper(int b, int n, int m, 12 | const float *dataset, float *temp, 13 | int *idxs); 14 | 15 | at::Tensor gather_points(at::Tensor points, at::Tensor idx) { 16 | CHECK_CONTIGUOUS(points); 17 | CHECK_CONTIGUOUS(idx); 18 | CHECK_IS_FLOAT(points); 19 | CHECK_IS_INT(idx); 20 | 21 | if (points.type().is_cuda()) { 22 | CHECK_CUDA(idx); 23 | } 24 | 25 | at::Tensor output = 26 | torch::zeros({points.size(0), points.size(1), idx.size(1)}, 27 | at::device(points.device()).dtype(at::ScalarType::Float)); 28 | 29 | if (points.type().is_cuda()) { 30 | gather_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 31 | idx.size(1), points.data(), 32 | idx.data(), output.data()); 33 | } else { 34 | AT_CHECK(false, "CPU not supported"); 35 | } 36 | 37 | return output; 38 | } 39 | 40 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, 41 | const int n) { 42 | CHECK_CONTIGUOUS(grad_out); 43 | CHECK_CONTIGUOUS(idx); 44 | CHECK_IS_FLOAT(grad_out); 45 | CHECK_IS_INT(idx); 46 | 47 | if (grad_out.type().is_cuda()) { 48 | CHECK_CUDA(idx); 49 | } 50 | 51 | at::Tensor output = 52 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 53 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 54 | 55 | if (grad_out.type().is_cuda()) { 56 | gather_points_grad_kernel_wrapper(grad_out.size(0), grad_out.size(1), n, 57 | idx.size(1), grad_out.data(), 58 | idx.data(), output.data()); 59 | } else { 60 | AT_CHECK(false, "CPU not supported"); 61 | } 62 | 63 | return output; 64 | } 65 | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples) { 66 | CHECK_CONTIGUOUS(points); 67 | CHECK_IS_FLOAT(points); 68 | 69 | at::Tensor output = 70 | torch::zeros({points.size(0), nsamples}, 71 | at::device(points.device()).dtype(at::ScalarType::Int)); 72 | 73 | at::Tensor tmp = 74 | torch::full({points.size(0), points.size(1)}, 1e10, 75 | at::device(points.device()).dtype(at::ScalarType::Float)); 76 | 77 | if (points.type().is_cuda()) { 78 | furthest_point_sampling_kernel_wrapper( 79 | points.size(0), points.size(1), nsamples, points.data(), 80 | tmp.data(), output.data()); 81 | } else { 82 | AT_CHECK(false, "CPU not supported"); 83 | } 84 | 85 | return output; 86 | } 87 | -------------------------------------------------------------------------------- /pointnet2/_ext-src/src/interpolate.cpp: -------------------------------------------------------------------------------- 1 | #include "interpolate.h" 2 | #include "utils.h" 3 | 4 | void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, 5 | const float *known, float *dist2, int *idx); 6 | void three_interpolate_kernel_wrapper(int b, int c, int m, int n, 7 | const float *points, const int *idx, 8 | const float *weight, float *out); 9 | void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m, 10 | const float *grad_out, 11 | const int *idx, const float *weight, 12 | float *grad_points); 13 | 14 | std::vector three_nn(at::Tensor unknowns, at::Tensor knows) { 15 | CHECK_CONTIGUOUS(unknowns); 16 | CHECK_CONTIGUOUS(knows); 17 | CHECK_IS_FLOAT(unknowns); 18 | CHECK_IS_FLOAT(knows); 19 | 20 | if (unknowns.type().is_cuda()) { 21 | CHECK_CUDA(knows); 22 | } 23 | 24 | at::Tensor idx = 25 | torch::zeros({unknowns.size(0), unknowns.size(1), 3}, 26 | at::device(unknowns.device()).dtype(at::ScalarType::Int)); 27 | at::Tensor dist2 = 28 | torch::zeros({unknowns.size(0), unknowns.size(1), 3}, 29 | at::device(unknowns.device()).dtype(at::ScalarType::Float)); 30 | 31 | if (unknowns.type().is_cuda()) { 32 | three_nn_kernel_wrapper(unknowns.size(0), unknowns.size(1), knows.size(1), 33 | unknowns.data(), knows.data(), 34 | dist2.data(), idx.data()); 35 | } else { 36 | AT_CHECK(false, "CPU not supported"); 37 | } 38 | 39 | return {dist2, idx}; 40 | } 41 | 42 | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, 43 | at::Tensor weight) { 44 | CHECK_CONTIGUOUS(points); 45 | CHECK_CONTIGUOUS(idx); 46 | CHECK_CONTIGUOUS(weight); 47 | CHECK_IS_FLOAT(points); 48 | CHECK_IS_INT(idx); 49 | CHECK_IS_FLOAT(weight); 50 | 51 | if (points.type().is_cuda()) { 52 | CHECK_CUDA(idx); 53 | CHECK_CUDA(weight); 54 | } 55 | 56 | at::Tensor output = 57 | torch::zeros({points.size(0), points.size(1), idx.size(1)}, 58 | at::device(points.device()).dtype(at::ScalarType::Float)); 59 | 60 | if (points.type().is_cuda()) { 61 | three_interpolate_kernel_wrapper( 62 | points.size(0), points.size(1), points.size(2), idx.size(1), 63 | points.data(), idx.data(), weight.data(), 64 | output.data()); 65 | } else { 66 | AT_CHECK(false, "CPU not supported"); 67 | } 68 | 69 | return output; 70 | } 71 | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, 72 | at::Tensor weight, const int m) { 73 | CHECK_CONTIGUOUS(grad_out); 74 | CHECK_CONTIGUOUS(idx); 75 | CHECK_CONTIGUOUS(weight); 76 | CHECK_IS_FLOAT(grad_out); 77 | CHECK_IS_INT(idx); 78 | CHECK_IS_FLOAT(weight); 79 | 80 | if (grad_out.type().is_cuda()) { 81 | CHECK_CUDA(idx); 82 | CHECK_CUDA(weight); 83 | } 84 | 85 | at::Tensor output = 86 | torch::zeros({grad_out.size(0), grad_out.size(1), m}, 87 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 88 | 89 | if (grad_out.type().is_cuda()) { 90 | three_interpolate_kernel_wrapper( 91 | grad_out.size(0), grad_out.size(1), grad_out.size(2), m, 92 | grad_out.data(), idx.data(), weight.data(), 93 | output.data()); 94 | } else { 95 | AT_CHECK(false, "CPU not supported"); 96 | } 97 | 98 | return output; 99 | } 100 | -------------------------------------------------------------------------------- /metrics.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from shapely.geometry import Polygon 3 | 4 | 5 | class AverageMeter(object): 6 | """Computes and stores the average and current value""" 7 | 8 | def __init__(self): 9 | self.reset() 10 | 11 | def reset(self): 12 | self.val = 0 13 | self.avg = 0 14 | self.sum = 0 15 | self.count = 0 16 | 17 | def update(self, val, n=1): 18 | self.val = val 19 | self.sum += val * n 20 | self.count += n 21 | self.avg = self.sum / self.count 22 | 23 | 24 | def estimateAccuracy(box_a, box_b, dim=3): 25 | if dim == 3: 26 | return np.linalg.norm(box_a.center - box_b.center, ord=2) 27 | elif dim == 2: 28 | return np.linalg.norm( 29 | box_a.center[[0, 2]] - box_b.center[[0, 2]], ord=2) 30 | 31 | 32 | def fromBoxToPoly(box): 33 | return Polygon(tuple(box.corners()[[0, 2]].T[[0, 1, 5, 4]])) 34 | 35 | 36 | def estimateOverlap(box_a, box_b, dim=2): 37 | # if box_a == box_b: 38 | # return 1.0 39 | 40 | Poly_anno = fromBoxToPoly(box_a) 41 | Poly_subm = fromBoxToPoly(box_b) 42 | 43 | box_inter = Poly_anno.intersection(Poly_subm) 44 | box_union = Poly_anno.union(Poly_subm) 45 | if dim == 2: 46 | return box_inter.area / box_union.area 47 | 48 | else: 49 | 50 | ymax = min(box_a.center[1], box_b.center[1]) 51 | ymin = max(box_a.center[1] - box_a.wlh[2], 52 | box_b.center[1] - box_b.wlh[2]) 53 | 54 | inter_vol = box_inter.area * max(0, ymax - ymin) 55 | anno_vol = box_a.wlh[0] * box_a.wlh[1] * box_a.wlh[2] 56 | subm_vol = box_b.wlh[0] * box_b.wlh[1] * box_b.wlh[2] 57 | 58 | overlap = inter_vol * 1.0 / (anno_vol + subm_vol - inter_vol) 59 | 60 | return overlap 61 | 62 | 63 | class Success(object): 64 | """Computes and stores the Success""" 65 | 66 | def __init__(self, n=21, max_overlap=1): 67 | self.max_overlap = max_overlap 68 | self.Xaxis = np.linspace(0, self.max_overlap, n) 69 | self.reset() 70 | 71 | def reset(self): 72 | self.overlaps = [] 73 | 74 | def add_overlap(self, val): 75 | self.overlaps.append(val) 76 | 77 | @property 78 | def count(self): 79 | return len(self.overlaps) 80 | 81 | @property 82 | def value(self): 83 | succ = [ 84 | np.sum(i >= thres 85 | for i in self.overlaps).astype(float) / self.count 86 | for thres in self.Xaxis 87 | ] 88 | return np.array(succ) 89 | 90 | @property 91 | def average(self): 92 | if len(self.overlaps) == 0: 93 | return 0 94 | return np.trapz(self.value, x=self.Xaxis) * 100 / self.max_overlap 95 | 96 | 97 | class Precision(object): 98 | """Computes and stores the Precision""" 99 | 100 | def __init__(self, n=21, max_accuracy=2): 101 | self.max_accuracy = max_accuracy 102 | self.Xaxis = np.linspace(0, self.max_accuracy, n) 103 | self.reset() 104 | 105 | def reset(self): 106 | self.accuracies = [] 107 | 108 | def add_accuracy(self, val): 109 | self.accuracies.append(val) 110 | 111 | @property 112 | def count(self): 113 | return len(self.accuracies) 114 | 115 | @property 116 | def value(self): 117 | prec = [ 118 | np.sum(i <= thres 119 | for i in self.accuracies).astype(float) / self.count 120 | for thres in self.Xaxis 121 | ] 122 | return np.array(prec) 123 | 124 | @property 125 | def average(self): 126 | if len(self.accuracies) == 0: 127 | return 0 128 | return np.trapz(self.value, x=self.Xaxis) * 100 / self.max_accuracy 129 | 130 | -------------------------------------------------------------------------------- /pointnet2/_ext-src/src/interpolate_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include "cuda_utils.h" 6 | 7 | // input: unknown(b, n, 3) known(b, m, 3) 8 | // output: dist2(b, n, 3), idx(b, n, 3) 9 | __global__ void three_nn_kernel(int b, int n, int m, 10 | const float *__restrict__ unknown, 11 | const float *__restrict__ known, 12 | float *__restrict__ dist2, 13 | int *__restrict__ idx) { 14 | int batch_index = blockIdx.x; 15 | unknown += batch_index * n * 3; 16 | known += batch_index * m * 3; 17 | dist2 += batch_index * n * 3; 18 | idx += batch_index * n * 3; 19 | 20 | int index = threadIdx.x; 21 | int stride = blockDim.x; 22 | for (int j = index; j < n; j += stride) { 23 | float ux = unknown[j * 3 + 0]; 24 | float uy = unknown[j * 3 + 1]; 25 | float uz = unknown[j * 3 + 2]; 26 | 27 | double best1 = 1e40, best2 = 1e40, best3 = 1e40; 28 | int besti1 = 0, besti2 = 0, besti3 = 0; 29 | for (int k = 0; k < m; ++k) { 30 | float x = known[k * 3 + 0]; 31 | float y = known[k * 3 + 1]; 32 | float z = known[k * 3 + 2]; 33 | float d = (ux - x) * (ux - x) + (uy - y) * (uy - y) + (uz - z) * (uz - z); 34 | if (d < best1) { 35 | best3 = best2; 36 | besti3 = besti2; 37 | best2 = best1; 38 | besti2 = besti1; 39 | best1 = d; 40 | besti1 = k; 41 | } else if (d < best2) { 42 | best3 = best2; 43 | besti3 = besti2; 44 | best2 = d; 45 | besti2 = k; 46 | } else if (d < best3) { 47 | best3 = d; 48 | besti3 = k; 49 | } 50 | } 51 | dist2[j * 3 + 0] = best1; 52 | dist2[j * 3 + 1] = best2; 53 | dist2[j * 3 + 2] = best3; 54 | 55 | idx[j * 3 + 0] = besti1; 56 | idx[j * 3 + 1] = besti2; 57 | idx[j * 3 + 2] = besti3; 58 | } 59 | } 60 | 61 | void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, 62 | const float *known, float *dist2, int *idx) { 63 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 64 | three_nn_kernel<<>>(b, n, m, unknown, known, 65 | dist2, idx); 66 | 67 | CUDA_CHECK_ERRORS(); 68 | } 69 | 70 | // input: points(b, c, m), idx(b, n, 3), weight(b, n, 3) 71 | // output: out(b, c, n) 72 | __global__ void three_interpolate_kernel(int b, int c, int m, int n, 73 | const float *__restrict__ points, 74 | const int *__restrict__ idx, 75 | const float *__restrict__ weight, 76 | float *__restrict__ out) { 77 | int batch_index = blockIdx.x; 78 | points += batch_index * m * c; 79 | 80 | idx += batch_index * n * 3; 81 | weight += batch_index * n * 3; 82 | 83 | out += batch_index * n * c; 84 | 85 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 86 | const int stride = blockDim.y * blockDim.x; 87 | for (int i = index; i < c * n; i += stride) { 88 | const int l = i / n; 89 | const int j = i % n; 90 | float w1 = weight[j * 3 + 0]; 91 | float w2 = weight[j * 3 + 1]; 92 | float w3 = weight[j * 3 + 2]; 93 | 94 | int i1 = idx[j * 3 + 0]; 95 | int i2 = idx[j * 3 + 1]; 96 | int i3 = idx[j * 3 + 2]; 97 | 98 | out[i] = points[l * m + i1] * w1 + points[l * m + i2] * w2 + 99 | points[l * m + i3] * w3; 100 | } 101 | } 102 | 103 | void three_interpolate_kernel_wrapper(int b, int c, int m, int n, 104 | const float *points, const int *idx, 105 | const float *weight, float *out) { 106 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 107 | three_interpolate_kernel<<>>( 108 | b, c, m, n, points, idx, weight, out); 109 | 110 | CUDA_CHECK_ERRORS(); 111 | } 112 | 113 | // input: grad_out(b, c, n), idx(b, n, 3), weight(b, n, 3) 114 | // output: grad_points(b, c, m) 115 | 116 | __global__ void three_interpolate_grad_kernel( 117 | int b, int c, int n, int m, const float *__restrict__ grad_out, 118 | const int *__restrict__ idx, const float *__restrict__ weight, 119 | float *__restrict__ grad_points) { 120 | int batch_index = blockIdx.x; 121 | grad_out += batch_index * n * c; 122 | idx += batch_index * n * 3; 123 | weight += batch_index * n * 3; 124 | grad_points += batch_index * m * c; 125 | 126 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 127 | const int stride = blockDim.y * blockDim.x; 128 | for (int i = index; i < c * n; i += stride) { 129 | const int l = i / n; 130 | const int j = i % n; 131 | float w1 = weight[j * 3 + 0]; 132 | float w2 = weight[j * 3 + 1]; 133 | float w3 = weight[j * 3 + 2]; 134 | 135 | int i1 = idx[j * 3 + 0]; 136 | int i2 = idx[j * 3 + 1]; 137 | int i3 = idx[j * 3 + 2]; 138 | 139 | atomicAdd(grad_points + l * m + i1, grad_out[i] * w1); 140 | atomicAdd(grad_points + l * m + i2, grad_out[i] * w2); 141 | atomicAdd(grad_points + l * m + i3, grad_out[i] * w3); 142 | } 143 | } 144 | 145 | void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m, 146 | const float *grad_out, 147 | const int *idx, const float *weight, 148 | float *grad_points) { 149 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 150 | three_interpolate_grad_kernel<<>>( 151 | b, c, n, m, grad_out, idx, weight, grad_points); 152 | 153 | CUDA_CHECK_ERRORS(); 154 | } 155 | -------------------------------------------------------------------------------- /searchspace.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from pomegranate import MultivariateGaussianDistribution, GeneralMixtureModel 3 | import logging 4 | 5 | 6 | class SearchSpace(object): 7 | 8 | def reset(self): 9 | raise NotImplementedError 10 | 11 | def sample(self): 12 | raise NotImplementedError 13 | 14 | def addData(self, data, score): 15 | return 16 | 17 | 18 | class ExhaustiveSearch(SearchSpace): 19 | 20 | def __init__(self, 21 | search_space=[[-3.0, 3.0], [-3.0, 3.0], [-10.0, 10.0]], 22 | search_dims=[7, 7, 3]): 23 | 24 | x_space = np.linspace( 25 | search_space[0][0], search_space[0][1], 26 | search_dims[0]) 27 | 28 | y_space = np.linspace( 29 | search_space[1][0], search_space[1][1], 30 | search_dims[1]) 31 | 32 | a_space = np.linspace( 33 | search_space[2][0], search_space[2][1], 34 | search_dims[2]) 35 | 36 | X, Y, A = np.meshgrid(x_space, y_space, a_space) # create mesh grid 37 | 38 | self.search_grid = np.array([X.flatten(), Y.flatten(), A.flatten()]).T 39 | 40 | self.reset() 41 | 42 | def reset(self): 43 | return 44 | 45 | def sample(self, n=0): 46 | return self.search_grid 47 | 48 | 49 | class ParticleFiltering(SearchSpace): 50 | def __init__(self, bnd=[1, 1, 10]): 51 | self.bnd = bnd 52 | self.reset() 53 | 54 | def sample(self, n=10): 55 | samples = [] 56 | for i in range(n): 57 | if len(self.data) > 0: 58 | i_mean = np.random.choice( 59 | list(range(len(self.data))), 60 | p=self.score / np.linalg.norm(self.score, ord=1)) 61 | sample = np.random.multivariate_normal( 62 | mean=self.data[i_mean], cov=np.diag(np.array(self.bnd))) 63 | else: 64 | sample = np.random.multivariate_normal( 65 | mean=np.zeros(len(self.bnd)), 66 | cov=np.diag(np.array(self.bnd) * 3)) 67 | 68 | samples.append(sample) 69 | return np.array(samples) 70 | 71 | def addData(self, data, score): 72 | score = score.clip(min=1e-5) # prevent sum=0 in case of bad scores 73 | self.data = data 74 | self.score = score 75 | 76 | def reset(self): 77 | if len(self.bnd) == 2: 78 | self.data = np.array([[], []]).T 79 | else: 80 | self.data = np.array([[], [], []]).T 81 | self.score = np.ones(np.shape(self.data)[0]) 82 | self.score = self.score / np.linalg.norm(self.score, ord=1) 83 | 84 | 85 | class KalmanFiltering(SearchSpace): 86 | def __init__(self, bnd=[1, 1, 10]): 87 | self.bnd = bnd 88 | self.reset() 89 | 90 | def sample(self, n=10): 91 | return np.random.multivariate_normal(self.mean, self.cov, size=n) 92 | 93 | def addData(self, data, score): 94 | score = score.clip(min=1e-5) # prevent sum=0 in case of bad scores 95 | self.data = np.concatenate((self.data, data)) 96 | self.score = np.concatenate((self.score, score)) 97 | self.mean = np.average(self.data, weights=self.score, axis=0) 98 | self.cov = np.cov(self.data.T, ddof=0, aweights=self.score) 99 | 100 | def reset(self): 101 | self.mean = np.zeros(len(self.bnd)) 102 | self.cov = np.diag(self.bnd) 103 | if len(self.bnd) == 2: 104 | self.data = np.array([[], []]).T 105 | else: 106 | self.data = np.array([[], [], []]).T 107 | self.score = np.array([]) 108 | 109 | 110 | class GaussianMixtureModel(SearchSpace): 111 | 112 | def __init__(self, n_comp=5, dim=3): 113 | self.dim = dim 114 | self.reset(n_comp) 115 | 116 | def sample(self, n=10): 117 | try: 118 | X1 = np.stack(self.model.sample(int(np.round(0.8 * n)))) 119 | if self.dim == 2: 120 | mean = np.mean(X1, axis=0) 121 | std = np.diag([1.0, 1.0]) 122 | gmm = MultivariateGaussianDistribution(mean, std) 123 | X2 = np.stack(gmm.sample(int(np.round(0.1 * n)))) 124 | 125 | mean = np.mean(X1, axis=0) 126 | std = np.diag([1e-3, 1e-3]) 127 | gmm = MultivariateGaussianDistribution(mean, std) 128 | X3 = np.stack(gmm.sample(int(np.round(0.1 * n)))) 129 | 130 | else: 131 | mean = np.mean(X1, axis=0) 132 | std = np.diag([1.0, 1.0, 1e-3]) 133 | gmm = MultivariateGaussianDistribution(mean, std) 134 | X2 = np.stack(gmm.sample(int(np.round(0.1 * n)))) 135 | 136 | mean = np.mean(X1, axis=0) 137 | std = np.diag([1e-3, 1e-3, 10.0]) 138 | gmm = MultivariateGaussianDistribution(mean, std) 139 | X3 = np.stack(gmm.sample(int(np.round(0.1 * n)))) 140 | 141 | X = np.concatenate((X1, X2, X3)) 142 | 143 | except ValueError: 144 | print("exception caught on sampling") 145 | if self.dim == 2: 146 | mean = np.zeros(self.dim) 147 | std = np.diag([1.0, 1.0]) 148 | gmm = MultivariateGaussianDistribution(mean, std) 149 | X = gmm.sample(int(n)) 150 | else: 151 | mean = np.zeros(self.dim) 152 | std = np.diag([1.0, 1.0, 5.0]) 153 | gmm = MultivariateGaussianDistribution(mean, std) 154 | X = gmm.sample(int(n)) 155 | return X 156 | 157 | def addData(self, data, score): 158 | score = score.clip(min=1e-5) 159 | self.data = data 160 | self.score = score 161 | 162 | score_normed = self.score / np.linalg.norm(self.score, ord=1) 163 | try: 164 | model = GeneralMixtureModel.from_samples( 165 | MultivariateGaussianDistribution, 166 | n_components=self.n_comp, 167 | X=self.data, 168 | weights=score_normed) 169 | self.model = model 170 | except: 171 | logging.info("catched an exception") 172 | 173 | def reset(self, n_comp=5): 174 | self.n_comp = n_comp 175 | 176 | if self.dim == 2: 177 | self.data = np.array([[], []]).T 178 | else: 179 | self.data = np.array([[], [], []]).T 180 | self.score = np.ones(np.shape(self.data)[0]) 181 | self.score = self.score / np.linalg.norm(self.score, ord=1) 182 | if self.dim == 2: 183 | self.model = MultivariateGaussianDistribution( 184 | np.zeros(self.dim), np.diag([1.0, 1.0])) 185 | else: 186 | self.model = MultivariateGaussianDistribution( 187 | np.zeros(self.dim), np.diag([1.0, 1.0, 5.0])) 188 | -------------------------------------------------------------------------------- /pointnet2/models/pointnet_tracking.py: -------------------------------------------------------------------------------- 1 | from __future__ import ( 2 | division, 3 | absolute_import, 4 | with_statement, 5 | print_function, 6 | unicode_literals, 7 | ) 8 | import torch 9 | import torch.nn as nn 10 | import etw_pytorch_utils as pt_utils 11 | from collections import namedtuple 12 | import torch.nn.functional as F 13 | 14 | from pointnet2.utils.pointnet2_modules import PointnetSAModule, PointnetFPModule, PointnetProposalModule 15 | 16 | 17 | 18 | class Pointnet_Backbone(nn.Module): 19 | r""" 20 | PointNet2 with single-scale grouping 21 | Semantic segmentation network that uses feature propogation layers 22 | 23 | Parameters 24 | ---------- 25 | num_classes: int 26 | Number of semantics classes to predict over -- size of softmax classifier that run for each point 27 | input_channels: int = 6 28 | Number of input channels in the feature descriptor for each point. If the point cloud is Nx9, this 29 | value should be 6 as in an Nx9 point cloud, 3 of the channels are xyz, and 6 are feature descriptors 30 | use_xyz: bool = True 31 | Whether or not to use the xyz position of a point as a feature 32 | """ 33 | 34 | def __init__(self, input_channels=3, use_xyz=True): 35 | super(Pointnet_Backbone, self).__init__() 36 | 37 | self.SA_modules = nn.ModuleList() 38 | self.SA_modules.append( 39 | PointnetSAModule( 40 | radius=0.3, 41 | nsample=32, 42 | mlp=[input_channels, 64, 64, 128], 43 | use_xyz=use_xyz, 44 | ) 45 | ) 46 | self.SA_modules.append( 47 | PointnetSAModule( 48 | radius=0.5, 49 | nsample=32, 50 | mlp=[128, 128, 128, 256], 51 | use_xyz=use_xyz, 52 | ) 53 | ) 54 | self.SA_modules.append( 55 | PointnetSAModule( 56 | radius=0.7, 57 | nsample=32, 58 | mlp=[256, 256, 256, 256], 59 | use_xyz=use_xyz, 60 | ) 61 | ) 62 | self.cov_final = nn.Conv1d(256, 256, kernel_size=1) 63 | 64 | 65 | def _break_up_pc(self, pc): 66 | xyz = pc[..., 0:3].contiguous() 67 | features = pc[..., 3:].transpose(1, 2).contiguous() if pc.size(-1) > 3 else None 68 | 69 | return xyz, features 70 | 71 | def forward(self, pointcloud, numpoints): 72 | # type: (Pointnet2SSG, torch.cuda.FloatTensor) -> pt_utils.Seq 73 | r""" 74 | Forward pass of the network 75 | 76 | Parameters 77 | ---------- 78 | pointcloud: Variable(torch.cuda.FloatTensor) 79 | (B, N, 3 + input_channels) tensor 80 | Point cloud to run predicts on 81 | Each point in the point-cloud MUST 82 | be formated as (x, y, z, features...) 83 | """ 84 | xyz, features = self._break_up_pc(pointcloud) 85 | 86 | l_xyz, l_features = [xyz], [features] 87 | for i in range(len(self.SA_modules)): 88 | li_xyz, li_features = self.SA_modules[i](l_xyz[i], l_features[i], numpoints[i]) 89 | l_xyz.append(li_xyz) 90 | l_features.append(li_features) 91 | 92 | 93 | return l_xyz[-1], self.cov_final(l_features[-1]) 94 | 95 | 96 | 97 | 98 | class Pointnet_Tracking(nn.Module): 99 | r""" 100 | xorr the search and the template 101 | """ 102 | def __init__(self, input_channels=3, use_xyz=True, objective = False): 103 | super(Pointnet_Tracking, self).__init__() 104 | 105 | self.backbone_net = Pointnet_Backbone(input_channels, use_xyz) 106 | 107 | self.cosine = nn.CosineSimilarity(dim=1) 108 | 109 | self.mlp = pt_utils.SharedMLP([4+256,256,256,256], bn=True) 110 | 111 | self.FC_layer_cla = ( 112 | pt_utils.Seq(256) 113 | .conv1d(256, bn=True) 114 | .conv1d(256, bn=True) 115 | .conv1d(1, activation=None)) 116 | self.fea_layer = (pt_utils.Seq(256) 117 | .conv1d(256, bn=True) 118 | .conv1d(256, activation=None)) 119 | self.vote_layer = ( 120 | pt_utils.Seq(3+256) 121 | .conv1d(256, bn=True) 122 | .conv1d(256, bn=True) 123 | .conv1d(3+256, activation=None)) 124 | self.vote_aggregation = PointnetSAModule( 125 | radius=0.3, 126 | nsample=16, 127 | mlp=[1+256, 256, 256, 256], 128 | use_xyz=use_xyz) 129 | self.num_proposal = 64 130 | self.FC_proposal = ( 131 | pt_utils.Seq(256) 132 | .conv1d(256, bn=True) 133 | .conv1d(256, bn=True) 134 | .conv1d(3+1+1, activation=None)) 135 | 136 | def xcorr(self, x_label, x_object, template_xyz): 137 | 138 | B = x_object.size(0) 139 | f = x_object.size(1) 140 | n1 = x_object.size(2) 141 | n2 = x_label.size(2) 142 | final_out_cla = self.cosine(x_object.unsqueeze(-1).expand(B,f,n1,n2), x_label.unsqueeze(2).expand(B,f,n1,n2)) 143 | 144 | fusion_feature = torch.cat((final_out_cla.unsqueeze(1),template_xyz.transpose(1, 2).contiguous().unsqueeze(-1).expand(B,3,n1,n2)),dim = 1) 145 | 146 | fusion_feature = torch.cat((fusion_feature,x_object.unsqueeze(-1).expand(B,f,n1,n2)),dim = 1) 147 | 148 | fusion_feature = self.mlp(fusion_feature) 149 | 150 | fusion_feature = F.max_pool2d(fusion_feature, kernel_size=[fusion_feature.size(2), 1]) 151 | fusion_feature = fusion_feature.squeeze(2) 152 | fusion_feature = self.fea_layer(fusion_feature) 153 | 154 | return fusion_feature 155 | 156 | def forward(self, template, search): 157 | r""" 158 | template: B*512*3 or B*512*6 159 | search: B*1024*3 or B*1024*6 160 | """ 161 | template_xyz, template_feature = self.backbone_net(template, [256, 128, 64]) 162 | 163 | search_xyz, search_feature = self.backbone_net(search, [512, 256, 128]) 164 | 165 | fusion_feature = self.xcorr(search_feature, template_feature, template_xyz) 166 | 167 | estimation_cla = self.FC_layer_cla(fusion_feature).squeeze(1) 168 | 169 | score = estimation_cla.sigmoid() 170 | 171 | fusion_xyz_feature = torch.cat((search_xyz.transpose(1, 2).contiguous(),fusion_feature),dim = 1) 172 | 173 | offset = self.vote_layer(fusion_xyz_feature) 174 | vote = fusion_xyz_feature + offset 175 | vote_xyz = vote[:,0:3,:].transpose(1, 2).contiguous() 176 | vote_feature = vote[:,3:,:] 177 | 178 | vote_feature = torch.cat((score.unsqueeze(1),vote_feature),dim = 1) 179 | 180 | center_xyzs, proposal_features = self.vote_aggregation(vote_xyz, vote_feature, self.num_proposal) 181 | 182 | proposal_offsets = self.FC_proposal(proposal_features) 183 | 184 | estimation_boxs = torch.cat((proposal_offsets[:,0:3,:]+center_xyzs.transpose(1, 2).contiguous(),proposal_offsets[:,3:5,:]),dim=1) 185 | 186 | return estimation_cla, vote_xyz, estimation_boxs.transpose(1, 2).contiguous(), center_xyzs -------------------------------------------------------------------------------- /pointnet2/_ext-src/src/sampling_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "cuda_utils.h" 5 | 6 | // input: points(b, c, n) idx(b, m) 7 | // output: out(b, c, m) 8 | __global__ void gather_points_kernel(int b, int c, int n, int m, 9 | const float *__restrict__ points, 10 | const int *__restrict__ idx, 11 | float *__restrict__ out) { 12 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 13 | for (int l = blockIdx.y; l < c; l += gridDim.y) { 14 | for (int j = threadIdx.x; j < m; j += blockDim.x) { 15 | int a = idx[i * m + j]; 16 | out[(i * c + l) * m + j] = points[(i * c + l) * n + a]; 17 | } 18 | } 19 | } 20 | } 21 | 22 | void gather_points_kernel_wrapper(int b, int c, int n, int npoints, 23 | const float *points, const int *idx, 24 | float *out) { 25 | gather_points_kernel<<>>(b, c, n, npoints, 27 | points, idx, out); 28 | 29 | CUDA_CHECK_ERRORS(); 30 | } 31 | 32 | // input: grad_out(b, c, m) idx(b, m) 33 | // output: grad_points(b, c, n) 34 | __global__ void gather_points_grad_kernel(int b, int c, int n, int m, 35 | const float *__restrict__ grad_out, 36 | const int *__restrict__ idx, 37 | float *__restrict__ grad_points) { 38 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 39 | for (int l = blockIdx.y; l < c; l += gridDim.y) { 40 | for (int j = threadIdx.x; j < m; j += blockDim.x) { 41 | int a = idx[i * m + j]; 42 | atomicAdd(grad_points + (i * c + l) * n + a, 43 | grad_out[(i * c + l) * m + j]); 44 | } 45 | } 46 | } 47 | } 48 | 49 | void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 50 | const float *grad_out, const int *idx, 51 | float *grad_points) { 52 | gather_points_grad_kernel<<>>( 54 | b, c, n, npoints, grad_out, idx, grad_points); 55 | 56 | CUDA_CHECK_ERRORS(); 57 | } 58 | 59 | __device__ void __update(float *__restrict__ dists, int *__restrict__ dists_i, 60 | int idx1, int idx2) { 61 | const float v1 = dists[idx1], v2 = dists[idx2]; 62 | const int i1 = dists_i[idx1], i2 = dists_i[idx2]; 63 | dists[idx1] = max(v1, v2); 64 | dists_i[idx1] = v2 > v1 ? i2 : i1; 65 | } 66 | 67 | // Input dataset: (b, n, 3), tmp: (b, n) 68 | // Ouput idxs (b, m) 69 | template 70 | __global__ void furthest_point_sampling_kernel( 71 | int b, int n, int m, const float *__restrict__ dataset, 72 | float *__restrict__ temp, int *__restrict__ idxs) { 73 | if (m <= 0) return; 74 | __shared__ float dists[block_size]; 75 | __shared__ int dists_i[block_size]; 76 | 77 | int batch_index = blockIdx.x; 78 | dataset += batch_index * n * 3; 79 | temp += batch_index * n; 80 | idxs += batch_index * m; 81 | 82 | int tid = threadIdx.x; 83 | const int stride = block_size; 84 | 85 | int old = 0; 86 | if (threadIdx.x == 0) idxs[0] = old; 87 | 88 | __syncthreads(); 89 | for (int j = 1; j < m; j++) { 90 | int besti = 0; 91 | float best = -1; 92 | float x1 = dataset[old * 3 + 0]; 93 | float y1 = dataset[old * 3 + 1]; 94 | float z1 = dataset[old * 3 + 2]; 95 | for (int k = tid; k < n; k += stride) { 96 | float x2, y2, z2; 97 | x2 = dataset[k * 3 + 0]; 98 | y2 = dataset[k * 3 + 1]; 99 | z2 = dataset[k * 3 + 2]; 100 | float mag = (x2 * x2) + (y2 * y2) + (z2 * z2); 101 | if (mag <= 1e-3) continue; 102 | 103 | float d = 104 | (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + (z2 - z1) * (z2 - z1); 105 | 106 | float d2 = min(d, temp[k]); 107 | temp[k] = d2; 108 | besti = d2 > best ? k : besti; 109 | best = d2 > best ? d2 : best; 110 | } 111 | dists[tid] = best; 112 | dists_i[tid] = besti; 113 | __syncthreads(); 114 | 115 | if (block_size >= 512) { 116 | if (tid < 256) { 117 | __update(dists, dists_i, tid, tid + 256); 118 | } 119 | __syncthreads(); 120 | } 121 | if (block_size >= 256) { 122 | if (tid < 128) { 123 | __update(dists, dists_i, tid, tid + 128); 124 | } 125 | __syncthreads(); 126 | } 127 | if (block_size >= 128) { 128 | if (tid < 64) { 129 | __update(dists, dists_i, tid, tid + 64); 130 | } 131 | __syncthreads(); 132 | } 133 | if (block_size >= 64) { 134 | if (tid < 32) { 135 | __update(dists, dists_i, tid, tid + 32); 136 | } 137 | __syncthreads(); 138 | } 139 | if (block_size >= 32) { 140 | if (tid < 16) { 141 | __update(dists, dists_i, tid, tid + 16); 142 | } 143 | __syncthreads(); 144 | } 145 | if (block_size >= 16) { 146 | if (tid < 8) { 147 | __update(dists, dists_i, tid, tid + 8); 148 | } 149 | __syncthreads(); 150 | } 151 | if (block_size >= 8) { 152 | if (tid < 4) { 153 | __update(dists, dists_i, tid, tid + 4); 154 | } 155 | __syncthreads(); 156 | } 157 | if (block_size >= 4) { 158 | if (tid < 2) { 159 | __update(dists, dists_i, tid, tid + 2); 160 | } 161 | __syncthreads(); 162 | } 163 | if (block_size >= 2) { 164 | if (tid < 1) { 165 | __update(dists, dists_i, tid, tid + 1); 166 | } 167 | __syncthreads(); 168 | } 169 | 170 | old = dists_i[0]; 171 | if (tid == 0) idxs[j] = old; 172 | } 173 | } 174 | 175 | void furthest_point_sampling_kernel_wrapper(int b, int n, int m, 176 | const float *dataset, float *temp, 177 | int *idxs) { 178 | unsigned int n_threads = opt_n_threads(n); 179 | 180 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 181 | 182 | switch (n_threads) { 183 | case 512: 184 | furthest_point_sampling_kernel<512> 185 | <<>>(b, n, m, dataset, temp, idxs); 186 | break; 187 | case 256: 188 | furthest_point_sampling_kernel<256> 189 | <<>>(b, n, m, dataset, temp, idxs); 190 | break; 191 | case 128: 192 | furthest_point_sampling_kernel<128> 193 | <<>>(b, n, m, dataset, temp, idxs); 194 | break; 195 | case 64: 196 | furthest_point_sampling_kernel<64> 197 | <<>>(b, n, m, dataset, temp, idxs); 198 | break; 199 | case 32: 200 | furthest_point_sampling_kernel<32> 201 | <<>>(b, n, m, dataset, temp, idxs); 202 | break; 203 | case 16: 204 | furthest_point_sampling_kernel<16> 205 | <<>>(b, n, m, dataset, temp, idxs); 206 | break; 207 | case 8: 208 | furthest_point_sampling_kernel<8> 209 | <<>>(b, n, m, dataset, temp, idxs); 210 | break; 211 | case 4: 212 | furthest_point_sampling_kernel<4> 213 | <<>>(b, n, m, dataset, temp, idxs); 214 | break; 215 | case 2: 216 | furthest_point_sampling_kernel<2> 217 | <<>>(b, n, m, dataset, temp, idxs); 218 | break; 219 | case 1: 220 | furthest_point_sampling_kernel<1> 221 | <<>>(b, n, m, dataset, temp, idxs); 222 | break; 223 | default: 224 | furthest_point_sampling_kernel<512> 225 | <<>>(b, n, m, dataset, temp, idxs); 226 | } 227 | 228 | CUDA_CHECK_ERRORS(); 229 | } 230 | -------------------------------------------------------------------------------- /test_tracking.py: -------------------------------------------------------------------------------- 1 | import time 2 | import os 3 | import logging 4 | import argparse 5 | import random 6 | 7 | import numpy as np 8 | from tqdm import tqdm 9 | 10 | import torch 11 | 12 | import kitty_utils as utils 13 | import copy 14 | from datetime import datetime 15 | 16 | 17 | from metrics import AverageMeter, Success, Precision 18 | from metrics import estimateOverlap, estimateAccuracy 19 | from data_classes import PointCloud 20 | from Dataset import SiameseTest 21 | 22 | import torch.nn.functional as F 23 | from torch.autograd import Variable 24 | 25 | from pointnet2.models import Pointnet_Tracking 26 | 27 | def test(loader,model,epoch=-1,shape_aggregation="",reference_BB="",model_fusion="pointcloud",max_iter=-1,IoU_Space=3): 28 | 29 | batch_time = AverageMeter() 30 | data_time = AverageMeter() 31 | 32 | Success_main = Success() 33 | Precision_main = Precision() 34 | Success_batch = Success() 35 | Precision_batch = Precision() 36 | 37 | # switch to evaluate mode 38 | model.eval() 39 | end = time.time() 40 | 41 | dataset = loader.dataset 42 | batch_num = 0 43 | 44 | with tqdm(enumerate(loader), total=len(loader.dataset.list_of_anno)) as t: 45 | for batch in loader: 46 | batch_num = batch_num+1 47 | # measure data loading time 48 | data_time.update((time.time() - end)) 49 | for PCs, BBs, list_of_anno in batch: # tracklet 50 | results_BBs = [] 51 | 52 | for i, _ in enumerate(PCs): 53 | this_anno = list_of_anno[i] 54 | this_BB = BBs[i] 55 | this_PC = PCs[i] 56 | gt_boxs = [] 57 | result_boxs = [] 58 | 59 | # INITIAL FRAME 60 | if i == 0: 61 | box = BBs[i] 62 | results_BBs.append(box) 63 | model_PC = utils.getModel([this_PC], [this_BB], offset=dataset.offset_BB, scale=dataset.scale_BB) 64 | 65 | else: 66 | previous_BB = BBs[i - 1] 67 | 68 | # DEFINE REFERENCE BB 69 | if ("previous_result".upper() in reference_BB.upper()): 70 | ref_BB = results_BBs[-1] 71 | elif ("previous_gt".upper() in reference_BB.upper()): 72 | ref_BB = previous_BB 73 | # ref_BB = utils.getOffsetBB(this_BB,np.array([-1,1,1])) 74 | elif ("current_gt".upper() in reference_BB.upper()): 75 | ref_BB = this_BB 76 | 77 | candidate_PC,candidate_label,candidate_reg, new_ref_box, new_this_box = utils.cropAndCenterPC_label_test( 78 | this_PC, 79 | ref_BB,this_BB, 80 | offset=dataset.offset_BB, 81 | scale=dataset.scale_BB) 82 | 83 | candidate_PCs,candidate_labels,candidate_reg = utils.regularizePCwithlabel(candidate_PC, candidate_label,candidate_reg,dataset.input_size,istrain=False) 84 | 85 | candidate_PCs_torch = candidate_PCs.unsqueeze(0).cuda() 86 | 87 | # AGGREGATION: IO vs ONLY0 vs ONLYI vs ALL 88 | if ("firstandprevious".upper() in shape_aggregation.upper()): 89 | model_PC = utils.getModel([PCs[0], PCs[i-1]], [results_BBs[0],results_BBs[i-1]],offset=dataset.offset_BB,scale=dataset.scale_BB) 90 | elif ("first".upper() in shape_aggregation.upper()): 91 | model_PC = utils.getModel([PCs[0]], [results_BBs[0]],offset=dataset.offset_BB,scale=dataset.scale_BB) 92 | elif ("previous".upper() in shape_aggregation.upper()): 93 | model_PC = utils.getModel([PCs[i-1]], [results_BBs[i-1]],offset=dataset.offset_BB,scale=dataset.scale_BB) 94 | elif ("all".upper() in shape_aggregation.upper()): 95 | model_PC = utils.getModel(PCs[:i],results_BBs,offset=dataset.offset_BB,scale=dataset.scale_BB) 96 | else: 97 | model_PC = utils.getModel(PCs[:i],results_BBs,offset=dataset.offset_BB,scale=dataset.scale_BB) 98 | 99 | model_PC_torch = utils.regularizePC(model_PC, dataset.input_size,istrain=False).unsqueeze(0) 100 | model_PC_torch = Variable(model_PC_torch, requires_grad=False).cuda() 101 | candidate_PCs_torch = Variable(candidate_PCs_torch, requires_grad=False).cuda() 102 | 103 | estimation_cla, estimation_reg, estimation_box, center_xyz = model(model_PC_torch, candidate_PCs_torch) 104 | estimation_boxs_cpu = estimation_box.squeeze(0).detach().cpu().numpy() 105 | box_idx = estimation_boxs_cpu[:,4].argmax() 106 | estimation_box_cpu = estimation_boxs_cpu[box_idx,0:4] 107 | 108 | box = utils.getOffsetBB(ref_BB,estimation_box_cpu) 109 | results_BBs.append(box) 110 | 111 | # estimate overlap/accuracy fro current sample 112 | this_overlap = estimateOverlap(BBs[i], results_BBs[-1], dim=IoU_Space) 113 | this_accuracy = estimateAccuracy(BBs[i], results_BBs[-1], dim=IoU_Space) 114 | 115 | Success_main.add_overlap(this_overlap) 116 | Precision_main.add_accuracy(this_accuracy) 117 | Success_batch.add_overlap(this_overlap) 118 | Precision_batch.add_accuracy(this_accuracy) 119 | 120 | # measure elapsed time 121 | batch_time.update(time.time() - end) 122 | end = time.time() 123 | 124 | t.update(1) 125 | 126 | if Success_main.count >= max_iter and max_iter >= 0: 127 | return Success_main.average, Precision_main.average 128 | 129 | 130 | t.set_description('Test {}: '.format(epoch)+ 131 | 'Time {:.3f}s '.format(batch_time.avg)+ 132 | '(it:{:.3f}s) '.format(batch_time.val)+ 133 | 'Data:{:.3f}s '.format(data_time.avg)+ 134 | '(it:{:.3f}s), '.format(data_time.val)+ 135 | 'Succ/Prec:'+ 136 | '{:.1f}/'.format(Success_main.average)+ 137 | '{:.1f}'.format(Precision_main.average)) 138 | logging.info('batch {}'.format(batch_num)+'Succ/Prec:'+ 139 | '{:.1f}/'.format(Success_batch.average)+ 140 | '{:.1f}'.format(Precision_batch.average)) 141 | Success_batch.reset() 142 | Precision_batch.reset() 143 | 144 | return Success_main.average, Precision_main.average 145 | 146 | 147 | 148 | if __name__ == '__main__': 149 | parser = argparse.ArgumentParser() 150 | parser.add_argument('--ngpu', type=int, default=2, help='# GPUs') 151 | parser.add_argument('--save_root_dir', type=str, default='./model/car_model/', help='output folder') 152 | parser.add_argument('--data_dir', type=str, default = './data/training', help='dataset path') 153 | parser.add_argument('--model', type=str, default = 'netR_36.pth', help='model name for training resume') 154 | parser.add_argument('--category_name', type=str, default = 'Car', help='Object to Track (Car/Pedetrian/Van/Cyclist)') 155 | parser.add_argument('--shape_aggregation',required=False,type=str,default="firstandprevious",help='Aggregation of shapes (first/previous/firstandprevious/all)') 156 | parser.add_argument('--reference_BB',required=False,type=str,default="previous_result",help='previous_result/previous_gt/current_gt') 157 | parser.add_argument('--model_fusion',required=False,type=str,default="pointcloud",help='early or late fusion (pointcloud/latent/space)') 158 | parser.add_argument('--IoU_Space',required=False,type=int,default=3,help='IoUBox vs IoUBEV (2 vs 3)') 159 | args = parser.parse_args() 160 | print (args) 161 | 162 | os.environ["CUDA_VISIBLE_DEVICES"] = '1,2' 163 | 164 | logging.basicConfig(format='%(asctime)s %(message)s', datefmt='%Y/%m/%d %H:%M:%S', \ 165 | filename=os.path.join(args.save_root_dir, datetime.now().strftime('%Y-%m-%d %H-%M-%S.log')), level=logging.INFO) 166 | logging.info('======================================================') 167 | 168 | args.manualSeed = 1 169 | random.seed(args.manualSeed) 170 | torch.manual_seed(args.manualSeed) 171 | 172 | netR = Pointnet_Tracking(input_channels=0, use_xyz=True) 173 | if args.ngpu > 1: 174 | netR = torch.nn.DataParallel(netR, range(args.ngpu)) 175 | if args.model != '': 176 | netR.load_state_dict(torch.load(os.path.join(args.save_root_dir, args.model))) 177 | netR.cuda() 178 | print(netR) 179 | torch.cuda.synchronize() 180 | # Car/Pedestrian/Van/Cyclist 181 | dataset_Test = SiameseTest( 182 | input_size=1024, 183 | path= args.data_dir, 184 | split='Test', 185 | category_name=args.category_name, 186 | offset_BB=0, 187 | scale_BB=1.25) 188 | 189 | test_loader = torch.utils.data.DataLoader( 190 | dataset_Test, 191 | collate_fn=lambda x: x, 192 | batch_size=1, 193 | shuffle=False, 194 | num_workers=1, 195 | pin_memory=True) 196 | 197 | Success_run = AverageMeter() 198 | Precision_run = AverageMeter() 199 | 200 | if dataset_Test.isTiny(): 201 | max_epoch = 2 202 | else: 203 | max_epoch = 1 204 | 205 | for epoch in range(max_epoch): 206 | Succ, Prec = test( 207 | test_loader, 208 | netR, 209 | epoch=epoch + 1, 210 | shape_aggregation=args.shape_aggregation, 211 | reference_BB=args.reference_BB, 212 | model_fusion=args.model_fusion, 213 | IoU_Space=args.IoU_Space) 214 | Success_run.update(Succ) 215 | Precision_run.update(Prec) 216 | logging.info("mean Succ/Prec {}/{}".format(Success_run.avg,Precision_run.avg)) 217 | -------------------------------------------------------------------------------- /data_classes.py: -------------------------------------------------------------------------------- 1 | # nuScenes dev-kit. 2 | # Code written by Oscar Beijbom, 2018. 3 | # Licensed under the Creative Commons [see licence.txt] 4 | 5 | #from __future__ import annotations 6 | import torch 7 | import numpy as np 8 | from pyquaternion import Quaternion 9 | 10 | 11 | class PointCloud: 12 | 13 | def __init__(self, points): 14 | """ 15 | Class for manipulating and viewing point clouds. 16 | :param points: . Input point cloud matrix. 17 | """ 18 | self.points = points 19 | if self.points.shape[0] > 3: 20 | self.points = self.points[0:3, :] 21 | 22 | @staticmethod 23 | def load_pcd_bin(file_name): 24 | """ 25 | Loads from binary format. Data is stored as (x, y, z, intensity, ring index). 26 | :param file_name: . 27 | :return: . Point cloud matrix (x, y, z, intensity). 28 | """ 29 | scan = np.fromfile(file_name, dtype=np.float32) 30 | points = scan.reshape((-1, 5))[:, :4] 31 | return points.T 32 | 33 | @classmethod 34 | def from_file(cls, file_name): 35 | """ 36 | Instantiate from a .pcl, .pdc, .npy, or .bin file. 37 | :param file_name: . Path of the pointcloud file on disk. 38 | :return: . 39 | """ 40 | 41 | if file_name.endswith('.bin'): 42 | points = cls.load_pcd_bin(file_name) 43 | elif file_name.endswith('.npy'): 44 | points = np.load(file_name) 45 | else: 46 | raise ValueError('Unsupported filetype {}'.format(file_name)) 47 | 48 | return cls(points) 49 | 50 | def nbr_points(self): 51 | """ 52 | Returns the number of points. 53 | :return: . Number of points. 54 | """ 55 | return self.points.shape[1] 56 | 57 | def subsample(self, ratio): 58 | """ 59 | Sub-samples the pointcloud. 60 | :param ratio: . Fraction to keep. 61 | :return: . 62 | """ 63 | selected_ind = np.random.choice(np.arange(0, self.nbr_points()), 64 | size=int(self.nbr_points() * ratio)) 65 | self.points = self.points[:, selected_ind] 66 | 67 | def remove_close(self, radius): 68 | """ 69 | Removes point too close within a certain radius from origin. 70 | :param radius: . 71 | :return: . 72 | """ 73 | 74 | x_filt = np.abs(self.points[0, :]) < radius 75 | y_filt = np.abs(self.points[1, :]) < radius 76 | not_close = np.logical_not(np.logical_and(x_filt, y_filt)) 77 | self.points = self.points[:, not_close] 78 | 79 | def translate(self, x): 80 | """ 81 | Applies a translation to the point cloud. 82 | :param x: . Translation in x, y, z. 83 | :return: . 84 | """ 85 | for i in range(3): 86 | self.points[i, :] = self.points[i, :] + x[i] 87 | 88 | def rotate(self, rot_matrix): 89 | """ 90 | Applies a rotation. 91 | :param rot_matrix: . Rotation matrix. 92 | :return: . 93 | """ 94 | self.points[:3, :] = np.dot(rot_matrix, self.points[:3, :]) 95 | 96 | def transform(self, transf_matrix): 97 | """ 98 | Applies a homogeneous transform. 99 | :param transf_matrix: . Homogenous transformation matrix. 100 | :return: . 101 | """ 102 | self.points[:3, :] = transf_matrix.dot( 103 | np.vstack((self.points[:3, :], np.ones(self.nbr_points()))))[:3, :] 104 | 105 | def convertToPytorch(self): 106 | """ 107 | Helper from pytorch. 108 | :return: Pytorch array of points. 109 | """ 110 | return torch.from_numpy(self.points) 111 | 112 | @staticmethod 113 | def fromPytorch(cls, pytorchTensor): 114 | """ 115 | Loads from binary format. Data is stored as (x, y, z, intensity, ring index). 116 | :param pyttorchTensor: . 117 | :return: . Point cloud matrix (x, y, z, intensity). 118 | """ 119 | points = pytorchTensor.numpy() 120 | # points = points.reshape((-1, 5))[:, :4] 121 | return cls(points) 122 | 123 | def normalize(self, wlh): 124 | normalizer = [wlh[1], wlh[0], wlh[2]] 125 | self.points = self.points / np.atleast_2d(normalizer).T 126 | 127 | 128 | class Box: 129 | """ Simple data class representing a 3d box including, label, score and velocity. """ 130 | 131 | def __init__(self, center, size, orientation, label=np.nan, score=np.nan, velocity=(np.nan, np.nan, np.nan), 132 | name=None): 133 | """ 134 | :param center: [: 3]. Center of box given as x, y, z. 135 | :param size: [: 3]. Size of box in width, length, height. 136 | :param orientation: . Box orientation. 137 | :param label: . Integer label, optional. 138 | :param score: . Classification score, optional. 139 | :param velocity: [: 3]. Box velocity in x, y, z direction. 140 | :param name: . Box name, optional. Can be used e.g. for denote category name. 141 | """ 142 | assert not np.any(np.isnan(center)) 143 | assert not np.any(np.isnan(size)) 144 | assert len(center) == 3 145 | assert len(size) == 3 146 | # assert type(orientation) == Quaternion 147 | 148 | self.center = np.array(center) 149 | self.wlh = np.array(size) 150 | self.orientation = orientation 151 | self.label = int(label) if not np.isnan(label) else label 152 | self.score = float(score) if not np.isnan(score) else score 153 | self.velocity = np.array(velocity) 154 | self.name = name 155 | 156 | def __eq__(self, other): 157 | center = np.allclose(self.center, other.center) 158 | wlh = np.allclose(self.wlh, other.wlh) 159 | orientation = np.allclose(self.orientation.elements, other.orientation.elements) 160 | label = (self.label == other.label) or (np.isnan(self.label) and np.isnan(other.label)) 161 | score = (self.score == other.score) or (np.isnan(self.score) and np.isnan(other.score)) 162 | vel = (np.allclose(self.velocity, other.velocity) or 163 | (np.all(np.isnan(self.velocity)) and np.all(np.isnan(other.velocity)))) 164 | 165 | return center and wlh and orientation and label and score and vel 166 | 167 | def __repr__(self): 168 | repr_str = 'label: {}, score: {:.2f}, xyz: [{:.2f}, {:.2f}, {:.2f}], wlh: [{:.2f}, {:.2f}, {:.2f}], ' \ 169 | 'rot axis: [{:.2f}, {:.2f}, {:.2f}], ang(degrees): {:.2f}, ang(rad): {:.2f}, ' \ 170 | 'vel: {:.2f}, {:.2f}, {:.2f}, name: {}' 171 | 172 | return repr_str.format(self.label, self.score, self.center[0], self.center[1], self.center[2], self.wlh[0], 173 | self.wlh[1], self.wlh[2], self.orientation.axis[0], self.orientation.axis[1], 174 | self.orientation.axis[2], self.orientation.degrees, self.orientation.radians, 175 | self.velocity[0], self.velocity[1], self.velocity[2], self.name) 176 | 177 | def encode(self): 178 | """ 179 | Encodes the box instance to a JSON-friendly vector representation. 180 | :return: [: 16]. List of floats encoding the box. 181 | """ 182 | return self.center.tolist() + self.wlh.tolist() + self.orientation.elements.tolist() + [self.label] + [self.score] + self.velocity.tolist() + [self.name] 183 | 184 | @classmethod 185 | def decode(cls, data): 186 | """ 187 | Instantiates a Box instance from encoded vector representation. 188 | :param data: [: 16]. Output from encode. 189 | :return: . 190 | """ 191 | return Box(data[0:3], data[3:6], Quaternion(data[6:10]), label=data[10], score=data[11], velocity=data[12:15], 192 | name=data[15]) 193 | 194 | @property 195 | def rotation_matrix(self): 196 | """ 197 | Return a rotation matrix. 198 | :return: . 199 | """ 200 | return self.orientation.rotation_matrix 201 | 202 | def translate(self, x): 203 | """ 204 | Applies a translation. 205 | :param x: . Translation in x, y, z direction. 206 | :return: . 207 | """ 208 | self.center += x 209 | 210 | def rotate(self, quaternion): 211 | """ 212 | Rotates box. 213 | :param quaternion: . Rotation to apply. 214 | :return: . 215 | """ 216 | self.center = np.dot(quaternion.rotation_matrix, self.center) 217 | self.orientation = quaternion * self.orientation 218 | self.velocity = np.dot(quaternion.rotation_matrix, self.velocity) 219 | 220 | def transform(self, transf_matrix): 221 | transformed = np.dot(transf_matrix[0:3,0:4].T, self.center) 222 | self.center = transformed[0:3]/transformed[3] 223 | self.orientation = self.orientation* Quaternion(matrix = transf_matrix[0:3,0:3]) 224 | self.velocity = np.dot(transf_matrix[0:3,0:3], self.velocity) 225 | 226 | def corners(self, wlh_factor=1.0): 227 | """ 228 | Returns the bounding box corners. 229 | :param wlh_factor: . Multiply w, l, h by a factor to inflate or deflate the box. 230 | :return: . First four corners are the ones facing forward. 231 | The last four are the ones facing backwards. 232 | """ 233 | w, l, h = self.wlh * wlh_factor 234 | 235 | # 3D bounding box corners. (Convention: x points forward, y to the left, z up.) 236 | x_corners = l / 2 * np.array([1, 1, 1, 1, -1, -1, -1, -1]) 237 | y_corners = w / 2 * np.array([1, -1, -1, 1, 1, -1, -1, 1]) 238 | z_corners = h / 2 * np.array([1, 1, -1, -1, 1, 1, -1, -1]) 239 | corners = np.vstack((x_corners, y_corners, z_corners)) 240 | 241 | # Rotate 242 | corners = np.dot(self.orientation.rotation_matrix, corners) 243 | 244 | # Translate 245 | x, y, z = self.center 246 | corners[0, :] = corners[0, :] + x 247 | corners[1, :] = corners[1, :] + y 248 | corners[2, :] = corners[2, :] + z 249 | 250 | return corners 251 | 252 | def bottom_corners(self): 253 | """ 254 | Returns the four bottom corners. 255 | :return: . Bottom corners. First two face forward, last two face backwards. 256 | """ 257 | return self.corners()[:, [2, 3, 7, 6]] 258 | -------------------------------------------------------------------------------- /pointnet2/utils/pointnet2_modules.py: -------------------------------------------------------------------------------- 1 | from __future__ import ( 2 | division, 3 | absolute_import, 4 | with_statement, 5 | print_function, 6 | unicode_literals, 7 | ) 8 | import torch 9 | import torch.nn as nn 10 | import torch.nn.functional as F 11 | import etw_pytorch_utils as pt_utils 12 | 13 | from pointnet2.utils import pointnet2_utils 14 | 15 | if False: 16 | # Workaround for type hints without depending on the `typing` module 17 | from typing import * 18 | 19 | 20 | class _PointnetSAModuleBase(nn.Module): 21 | def __init__(self): 22 | super(_PointnetSAModuleBase, self).__init__() 23 | self.groupers = None 24 | self.mlps = None 25 | 26 | def forward(self, xyz, features, npoint): 27 | # type: (_PointnetSAModuleBase, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor] 28 | r""" 29 | Parameters 30 | ---------- 31 | xyz : torch.Tensor 32 | (B, N, 3) tensor of the xyz coordinates of the features 33 | features : torch.Tensor 34 | (B, N, C) tensor of the descriptors of the the features 35 | 36 | Returns 37 | ------- 38 | new_xyz : torch.Tensor 39 | (B, npoint, 3) tensor of the new features' xyz 40 | new_features : torch.Tensor 41 | (B, \sum_k(mlps[k][-1]), npoint) tensor of the new_features descriptors 42 | """ 43 | 44 | self.npoint = npoint 45 | new_features_list = [] 46 | 47 | xyz_flipped = xyz.transpose(1, 2).contiguous() 48 | 49 | new_xyz = ( 50 | pointnet2_utils.gather_operation( 51 | #xyz_flipped, pointnet2_utils.furthest_point_sample(xyz, self.npoint) 52 | xyz_flipped,torch.arange(self.npoint).repeat(xyz.size(0),1).int().cuda() 53 | ) 54 | .transpose(1, 2) 55 | .contiguous() 56 | ) 57 | 58 | for i in range(len(self.groupers)): 59 | new_features = self.groupers[i]( 60 | xyz, new_xyz, features 61 | ) # (B, C, npoint, nsample) 62 | 63 | new_features = self.mlps[i](new_features) # (B, mlp[-1], npoint, nsample) 64 | new_features = F.max_pool2d( 65 | new_features, kernel_size=[1, new_features.size(3)] 66 | ) # (B, mlp[-1], npoint, 1) 67 | new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint) 68 | 69 | new_features_list.append(new_features) 70 | 71 | return new_xyz, torch.cat(new_features_list, dim=1) 72 | 73 | 74 | class PointnetSAModuleMSG(_PointnetSAModuleBase): 75 | r"""Pointnet set abstrction layer with multiscale grouping 76 | 77 | Parameters 78 | ---------- 79 | npoint : int 80 | Number of features 81 | radii : list of float32 82 | list of radii to group with 83 | nsamples : list of int32 84 | Number of samples in each ball query 85 | mlps : list of list of int32 86 | Spec of the pointnet before the global max_pool for each scale 87 | bn : bool 88 | Use batchnorm 89 | """ 90 | 91 | def __init__(self, radii, nsamples, mlps, bn=True, use_xyz=True, vote = False): 92 | # type: (PointnetSAModuleMSG, int, List[float], List[int], List[List[int]], bool, bool) -> None 93 | super(PointnetSAModuleMSG, self).__init__() 94 | 95 | assert len(radii) == len(nsamples) == len(mlps) 96 | 97 | self.groupers = nn.ModuleList() 98 | self.mlps = nn.ModuleList() 99 | for i in range(len(radii)): 100 | radius = radii[i] 101 | nsample = nsamples[i] 102 | if vote is False: 103 | self.groupers.append( 104 | pointnet2_utils.QueryAndGroup(radius, nsample, use_xyz=use_xyz)) 105 | else: 106 | self.groupers.append( 107 | pointnet2_utils.QueryAndGroup_score(radius, nsample, use_xyz=use_xyz)) 108 | 109 | mlp_spec = mlps[i] 110 | if use_xyz: 111 | mlp_spec[0] += 3 112 | 113 | self.mlps.append(pt_utils.SharedMLP(mlp_spec, bn=bn)) 114 | 115 | 116 | class PointnetSAModule(PointnetSAModuleMSG): 117 | r"""Pointnet set abstrction layer 118 | 119 | Parameters 120 | ---------- 121 | npoint : int 122 | Number of features 123 | radius : float 124 | Radius of ball 125 | nsample : int 126 | Number of samples in the ball query 127 | mlp : list 128 | Spec of the pointnet before the global max_pool 129 | bn : bool 130 | Use batchnorm 131 | """ 132 | 133 | def __init__( 134 | self, mlp, radius=None, nsample=None, bn=True, use_xyz=True 135 | ): 136 | # type: (PointnetSAModule, List[int], int, float, int, bool, bool) -> None 137 | super(PointnetSAModule, self).__init__( 138 | mlps=[mlp], 139 | radii=[radius], 140 | nsamples=[nsample], 141 | bn=bn, 142 | use_xyz=use_xyz, 143 | ) 144 | 145 | 146 | class PointnetProposalModule(PointnetSAModuleMSG): 147 | r"""Pointnet set abstrction layer 148 | 149 | Parameters 150 | ---------- 151 | npoint : int 152 | Number of features 153 | radius : float 154 | Radius of ball 155 | nsample : int 156 | Number of samples in the ball query 157 | mlp : list 158 | Spec of the pointnet before the global max_pool 159 | bn : bool 160 | Use batchnorm 161 | """ 162 | 163 | def __init__( 164 | self, mlp, radius=None, nsample=None, bn=True, use_xyz=True, vote = True 165 | ): 166 | # type: (PointnetSAModule, List[int], int, float, int, bool, bool) -> None 167 | super(PointnetProposalModule, self).__init__( 168 | mlps=[mlp], 169 | radii=[radius], 170 | nsamples=[nsample], 171 | bn=bn, 172 | use_xyz=use_xyz, 173 | vote = vote 174 | ) 175 | 176 | 177 | def forward(self, xyz, features, npoint,score): 178 | # type: (_PointnetSAModuleBase, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor] 179 | r""" 180 | Parameters 181 | ---------- 182 | xyz : torch.Tensor 183 | (B, N, 3) tensor of the xyz coordinates of the features 184 | features : torch.Tensor 185 | (B, N, C) tensor of the descriptors of the the features 186 | 187 | Returns 188 | ------- 189 | new_xyz : torch.Tensor 190 | (B, npoint, 3) tensor of the new features' xyz 191 | new_features : torch.Tensor 192 | (B, \sum_k(mlps[k][-1]), npoint) tensor of the new_features descriptors 193 | """ 194 | 195 | self.npoint = npoint 196 | new_features_list = [] 197 | 198 | xyz_flipped = xyz.transpose(1, 2).contiguous() 199 | 200 | new_xyz = ( 201 | pointnet2_utils.gather_operation( 202 | #xyz_flipped, pointnet2_utils.furthest_point_sample(xyz, self.npoint) 203 | xyz_flipped,torch.arange(self.npoint).repeat(xyz.size(0),1).int().cuda() 204 | ) 205 | .transpose(1, 2) 206 | .contiguous() 207 | ) 208 | 209 | for i in range(len(self.groupers)): 210 | new_features,score_id = self.groupers[i]( 211 | xyz, new_xyz, score, features 212 | ) # (B, C, npoint, nsample) 213 | #score_id = new_features[:,3,:,:].sum(dim = 2).argmax(dim = 1) 214 | 215 | #B 216 | #new_features_cpu = new_features.squeeze(0).detach().cpu().numpy() 217 | #np.savetxt('vote4.txt',new_features_cpu[0:4,i,:]) 218 | idx = torch.arange(new_features.size(0)) 219 | new_features = new_features[idx,:,score_id,:] 220 | #B*C*nsample 221 | new_features = new_features.unsqueeze(2) 222 | #B*C*1*nsample 223 | new_xyz = new_xyz[idx,score_id,:] 224 | #B*3 225 | 226 | new_features = self.mlps[i](new_features) # (B, mlp[-1], npoint, nsample) 227 | new_features = F.max_pool2d( 228 | new_features, kernel_size=[1, new_features.size(3)] 229 | ) # (B, mlp[-1], npoint, 1) 230 | new_features = new_features.squeeze(-1).squeeze(-1) # (B, mlp[-1]) 231 | 232 | new_features_list.append(new_features) 233 | 234 | return new_xyz, torch.cat(new_features_list, dim=1) 235 | 236 | 237 | class PointnetFPModule(nn.Module): 238 | r"""Propigates the features of one set to another 239 | 240 | Parameters 241 | ---------- 242 | mlp : list 243 | Pointnet module parameters 244 | bn : bool 245 | Use batchnorm 246 | """ 247 | 248 | def __init__(self, mlp, bn=True): 249 | # type: (PointnetFPModule, List[int], bool) -> None 250 | super(PointnetFPModule, self).__init__() 251 | self.mlp = pt_utils.SharedMLP(mlp, bn=bn) 252 | 253 | def forward(self, unknown, known, unknow_feats, known_feats): 254 | # type: (PointnetFPModule, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor) -> torch.Tensor 255 | r""" 256 | Parameters 257 | ---------- 258 | unknown : torch.Tensor 259 | (B, n, 3) tensor of the xyz positions of the unknown features 260 | known : torch.Tensor 261 | (B, m, 3) tensor of the xyz positions of the known features 262 | unknow_feats : torch.Tensor 263 | (B, C1, n) tensor of the features to be propigated to 264 | known_feats : torch.Tensor 265 | (B, C2, m) tensor of features to be propigated 266 | 267 | Returns 268 | ------- 269 | new_features : torch.Tensor 270 | (B, mlp[-1], n) tensor of the features of the unknown features 271 | """ 272 | 273 | if known is not None: 274 | dist, idx = pointnet2_utils.three_nn(unknown, known) 275 | dist_recip = 1.0 / (dist + 1e-8) 276 | norm = torch.sum(dist_recip, dim=2, keepdim=True) 277 | weight = dist_recip / norm 278 | 279 | interpolated_feats = pointnet2_utils.three_interpolate( 280 | known_feats, idx, weight 281 | ) 282 | else: 283 | interpolated_feats = known_feats.expand( 284 | *(known_feats.size()[0:2] + [unknown.size(1)]) 285 | ) 286 | 287 | if unknow_feats is not None: 288 | new_features = torch.cat( 289 | [interpolated_feats, unknow_feats], dim=1 290 | ) # (B, C2 + C1, n) 291 | else: 292 | new_features = interpolated_feats 293 | 294 | new_features = new_features.unsqueeze(-1) 295 | new_features = self.mlp(new_features) 296 | 297 | return new_features.squeeze(-1) 298 | 299 | 300 | if __name__ == "__main__": 301 | from torch.autograd import Variable 302 | 303 | torch.manual_seed(1) 304 | torch.cuda.manual_seed_all(1) 305 | xyz = Variable(torch.randn(2, 9, 3).cuda(), requires_grad=True) 306 | xyz_feats = Variable(torch.randn(2, 9, 6).cuda(), requires_grad=True) 307 | 308 | test_module = PointnetSAModuleMSG( 309 | npoint=2, radii=[5.0, 10.0], nsamples=[6, 3], mlps=[[9, 3], [9, 6]] 310 | ) 311 | test_module.cuda() 312 | print(test_module(xyz, xyz_feats)) 313 | 314 | # test_module = PointnetFPModule(mlp=[6, 6]) 315 | # test_module.cuda() 316 | # from torch.autograd import gradcheck 317 | # inputs = (xyz, xyz, None, xyz_feats) 318 | # test = gradcheck(test_module, inputs, eps=1e-6, atol=1e-4) 319 | # print(test) 320 | 321 | for _ in range(1): 322 | _, new_features = test_module(xyz, xyz_feats) 323 | new_features.backward(torch.cuda.FloatTensor(*new_features.size()).fill_(1)) 324 | print(new_features) 325 | print(xyz.grad) 326 | -------------------------------------------------------------------------------- /train_tracking.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import random 4 | import time 5 | import logging 6 | import pdb 7 | from tqdm import tqdm 8 | import numpy as np 9 | import scipy.io as sio 10 | 11 | import torch 12 | import torch.nn as nn 13 | import torch.nn.parallel 14 | import torch.backends.cudnn as cudnn 15 | import torch.optim as optim 16 | import torch.optim.lr_scheduler as lr_scheduler 17 | import torch.utils.data 18 | import torch.nn.functional as F 19 | from torch.autograd import Variable 20 | 21 | from Dataset import SiameseTrain 22 | from pointnet2.models import Pointnet_Tracking 23 | 24 | parser = argparse.ArgumentParser() 25 | parser.add_argument('--batchSize', type=int, default=48, help='input batch size') 26 | parser.add_argument('--workers', type=int, default=0, help='number of data loading workers') 27 | parser.add_argument('--nepoch', type=int, default=60, help='number of epochs to train for') 28 | parser.add_argument('--ngpu', type=int, default=2, help='# GPUs') 29 | parser.add_argument('--learning_rate', type=float, default=0.001, help='learning rate at t=0') 30 | parser.add_argument('--input_feature_num', type=int, default = 0, help='number of input point features') 31 | parser.add_argument('--data_dir', type=str, default = './data/training', help='dataset path') 32 | parser.add_argument('--category_name', type=str, default = 'Car', help='Object to Track (Car/Pedetrian/Van/Cyclist)') 33 | parser.add_argument('--save_root_dir', type=str, default='results', help='output folder') 34 | parser.add_argument('--model', type=str, default = '', help='model name for training resume') 35 | parser.add_argument('--optimizer', type=str, default = '', help='optimizer name for training resume') 36 | 37 | opt = parser.parse_args() 38 | print (opt) 39 | 40 | #torch.cuda.set_device(opt.main_gpu) 41 | os.environ["CUDA_VISIBLE_DEVICES"] = '0,1' 42 | 43 | opt.manualSeed = 1 44 | random.seed(opt.manualSeed) 45 | torch.manual_seed(opt.manualSeed) 46 | 47 | save_dir = opt.save_root_dir 48 | 49 | try: 50 | os.makedirs(save_dir) 51 | except OSError: 52 | pass 53 | 54 | logging.basicConfig(format='%(asctime)s %(message)s', datefmt='%Y/%m/%d %H:%M:%S', \ 55 | filename=os.path.join(save_dir, 'train.log'), level=logging.INFO) 56 | logging.info('======================================================') 57 | 58 | # 1. Load data 59 | train_data = SiameseTrain( 60 | input_size=1024, 61 | path= opt.data_dir, 62 | split='Train', 63 | category_name=opt.category_name, 64 | offset_BB=0, 65 | scale_BB=1.25) 66 | 67 | train_dataloader = torch.utils.data.DataLoader( 68 | train_data, 69 | batch_size=opt.batchSize, 70 | shuffle=True, 71 | num_workers=int(opt.workers), 72 | pin_memory=True) 73 | 74 | test_data = SiameseTrain( 75 | input_size=1024, 76 | path=opt.data_dir, 77 | split='Valid', 78 | category_name=opt.category_name, 79 | offset_BB=0, 80 | scale_BB=1.25) 81 | 82 | test_dataloader = torch.utils.data.DataLoader( 83 | test_data, 84 | batch_size=int(opt.batchSize/2), 85 | shuffle=False, 86 | num_workers=int(opt.workers), 87 | pin_memory=True) 88 | 89 | 90 | print('#Train data:', len(train_data), '#Test data:', len(test_data)) 91 | print (opt) 92 | 93 | # 2. Define model, loss and optimizer 94 | netR = Pointnet_Tracking(input_channels=opt.input_feature_num, use_xyz=True) 95 | if opt.ngpu > 1: 96 | netR = torch.nn.DataParallel(netR, range(opt.ngpu)) 97 | if opt.model != '': 98 | netR.load_state_dict(torch.load(os.path.join(save_dir, opt.model))) 99 | 100 | netR.cuda() 101 | print(netR) 102 | 103 | criterion_cla = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([1.0])).cuda() 104 | criterion_reg = nn.SmoothL1Loss(reduction='none').cuda() 105 | criterion_objective = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([2.0]), reduction='none').cuda() 106 | criterion_box = nn.SmoothL1Loss(reduction='none').cuda() 107 | optimizer = optim.Adam(netR.parameters(), lr=opt.learning_rate, betas = (0.5, 0.999), eps=1e-06) 108 | if opt.optimizer != '': 109 | optimizer.load_state_dict(torch.load(os.path.join(save_dir, opt.optimizer))) 110 | scheduler = lr_scheduler.StepLR(optimizer, step_size=12, gamma=0.2) 111 | 112 | # 3. Training and testing 113 | for epoch in range(opt.nepoch): 114 | scheduler.step(epoch) 115 | print('======>>>>> Online epoch: #%d, lr=%f <<<<<======' %(epoch, scheduler.get_lr()[0])) 116 | # # 3.1 switch to train mode 117 | torch.cuda.synchronize() 118 | netR.train() 119 | train_mse = 0.0 120 | timer = time.time() 121 | 122 | batch_correct = 0.0 123 | batch_cla_loss = 0.0 124 | batch_reg_loss = 0.0 125 | batch_box_loss = 0.0 126 | batch_num = 0.0 127 | batch_iou = 0.0 128 | batch_true_correct = 0.0 129 | for i, data in enumerate(tqdm(train_dataloader, 0)): 130 | if len(data[0]) == 1: 131 | continue 132 | torch.cuda.synchronize() 133 | # 3.1.1 load inputs and targets 134 | label_point_set, label_cla, label_reg, object_point_set = data 135 | label_cla = Variable(label_cla, requires_grad=False).cuda() 136 | label_reg = Variable(label_reg, requires_grad=False).cuda() 137 | object_point_set = Variable(object_point_set, requires_grad=False).cuda() 138 | label_point_set = Variable(label_point_set, requires_grad=False).cuda() 139 | 140 | # 3.1.2 compute output 141 | optimizer.zero_grad() 142 | estimation_cla, estimation_reg,estimation_box,center_xyz = netR(object_point_set, label_point_set) 143 | loss_cla = criterion_cla(estimation_cla, label_cla) 144 | loss_reg = criterion_reg(estimation_reg, label_reg[:,:,0:3]) 145 | loss_reg = (loss_reg.mean(2) * label_cla).sum()/(label_cla.sum()+1e-06) 146 | 147 | dist = torch.sum((center_xyz - label_reg[:,0:64,0:3])**2, dim=-1) 148 | dist = torch.sqrt(dist+1e-6) 149 | B = dist.size(0) 150 | K = 64 151 | objectness_label = torch.zeros((B,K), dtype=torch.float).cuda() 152 | objectness_mask = torch.zeros((B,K)).cuda() 153 | objectness_label[dist<0.3] = 1 154 | objectness_mask[dist<0.3] = 1 155 | objectness_mask[dist>0.6] = 1 156 | loss_objective = criterion_objective(estimation_box[:,:,4], objectness_label) 157 | loss_objective = torch.sum(loss_objective * objectness_mask)/(torch.sum(objectness_mask)+1e-6) 158 | loss_box = criterion_box(estimation_box[:,:,0:4],label_reg[:,0:64,:]) 159 | loss_box = (loss_box.mean(2) * objectness_label).sum()/(objectness_label.sum()+1e-06) 160 | 161 | loss = 0.2*loss_cla + loss_reg + 1.5*loss_objective + 0.2*loss_box 162 | 163 | # 3.1.3 compute gradient and do SGD step 164 | loss.backward() 165 | optimizer.step() 166 | torch.cuda.synchronize() 167 | 168 | # 3.1.4 update training error 169 | estimation_cla_cpu = estimation_cla.sigmoid().detach().cpu().numpy() 170 | label_cla_cpu = label_cla.detach().cpu().numpy() 171 | correct = float(np.sum((estimation_cla_cpu[0:len(label_point_set),:] > 0.4) == label_cla_cpu[0:len(label_point_set),:])) / 128.0 172 | true_correct = float(np.sum((np.float32(estimation_cla_cpu[0:len(label_point_set),:] > 0.4) + label_cla_cpu[0:len(label_point_set),:]) == 2)/(np.sum(label_cla_cpu[0:len(label_point_set),:]))) 173 | 174 | train_mse = train_mse + loss.data*len(label_point_set) 175 | batch_correct += correct 176 | batch_cla_loss += loss_cla.data 177 | batch_reg_loss += loss_reg.data 178 | batch_box_loss += loss_box.data 179 | batch_num += len(label_point_set) 180 | batch_true_correct += true_correct 181 | if (i+1)%20 == 0: 182 | print('\n ---- batch: %03d ----' % (i+1)) 183 | print('cla_loss: %f, reg_loss: %f, box_loss: %f' % (batch_cla_loss/20,batch_reg_loss/20,batch_box_loss/20)) 184 | print('accuracy: %f' % (batch_correct / float(batch_num))) 185 | print('true accuracy: %f' % (batch_true_correct / 20)) 186 | batch_correct = 0.0 187 | batch_cla_loss = 0.0 188 | batch_reg_loss = 0.0 189 | batch_box_loss = 0.0 190 | batch_num = 0.0 191 | batch_true_correct = 0.0 192 | 193 | # time taken 194 | train_mse = train_mse/len(train_data) 195 | torch.cuda.synchronize() 196 | timer = time.time() - timer 197 | timer = timer / len(train_data) 198 | print('==> time to learn 1 sample = %f (ms)' %(timer*1000)) 199 | 200 | torch.save(netR.state_dict(), '%s/netR_%d.pth' % (save_dir, epoch)) 201 | torch.save(optimizer.state_dict(), '%s/optimizer_%d.pth' % (save_dir, epoch)) 202 | 203 | # 3.2 switch to evaluate mode 204 | torch.cuda.synchronize() 205 | netR.eval() 206 | test_cla_loss = 0.0 207 | test_reg_loss = 0.0 208 | test_box_loss = 0.0 209 | test_correct = 0.0 210 | test_true_correct = 0.0 211 | timer = time.time() 212 | for i, data in enumerate(tqdm(test_dataloader, 0)): 213 | torch.cuda.synchronize() 214 | # 3.2.1 load inputs and targets 215 | label_point_set, label_cla, label_reg, object_point_set = data 216 | label_cla = Variable(label_cla, requires_grad=False).cuda() 217 | label_reg = Variable(label_reg, requires_grad=False).cuda() 218 | object_point_set = Variable(object_point_set, requires_grad=False).cuda() 219 | label_point_set = Variable(label_point_set, requires_grad=False).cuda() 220 | 221 | # 3.2.2 compute output 222 | estimation_cla, estimation_reg,estimation_box,center_xyz = netR(object_point_set, label_point_set) 223 | loss_cla = criterion_cla(estimation_cla, label_cla) 224 | loss_reg = criterion_reg(estimation_reg, label_reg[:,:,0:3]) 225 | loss_reg = (loss_reg.mean(2) * label_cla).sum()/(label_cla.sum()+1e-06) 226 | 227 | dist = torch.sum((center_xyz - label_reg[:,0:64,0:3])**2, dim=-1) 228 | dist = torch.sqrt(dist+1e-6) 229 | B = dist.size(0) 230 | K = 64 231 | objectness_label = torch.zeros((B,K), dtype=torch.float).cuda() 232 | objectness_mask = torch.zeros((B,K)).cuda() 233 | objectness_label[dist<0.3] = 1 234 | objectness_mask[dist<0.3] = 1 235 | objectness_mask[dist>0.6] = 1 236 | loss_objective = criterion_objective(estimation_box[:,:,4], objectness_label) 237 | loss_objective = torch.sum(loss_objective * objectness_mask)/(torch.sum(objectness_mask)+1e-6) 238 | loss_box = criterion_box(estimation_box[:,:,0:4],label_reg[:,0:64,:]) 239 | loss_box = (loss_box.mean(2) * objectness_label).sum()/(objectness_label.sum()+1e-06) 240 | loss = 0.2*loss_cla + loss_reg + 1.5*loss_objective + 0.2*loss_box 241 | 242 | torch.cuda.synchronize() 243 | test_cla_loss = test_cla_loss + loss_cla.data*len(label_point_set) 244 | test_reg_loss = test_reg_loss + loss_reg.data*len(label_point_set) 245 | test_box_loss = test_box_loss + loss_box.data*len(label_point_set) 246 | estimation_cla_cpu = estimation_cla.sigmoid().detach().cpu().numpy() 247 | label_cla_cpu = label_cla.detach().cpu().numpy() 248 | correct = float(np.sum((estimation_cla_cpu[0:len(label_point_set),:] > 0.4) == label_cla_cpu[0:len(label_point_set),:])) / 128.0 249 | true_correct = float(np.sum((np.float32(estimation_cla_cpu[0:len(label_point_set),:] > 0.4) + label_cla_cpu[0:len(label_point_set),:]) == 2)/(np.sum(label_cla_cpu[0:len(label_point_set),:]))) 250 | test_correct += correct 251 | test_true_correct += true_correct*len(label_point_set) 252 | 253 | # time taken 254 | torch.cuda.synchronize() 255 | timer = time.time() - timer 256 | timer = timer / len(test_data) 257 | print('==> time to learn 1 sample = %f (ms)' %(timer*1000)) 258 | # print mse 259 | test_cla_loss = test_cla_loss / len(test_data) 260 | test_reg_loss = test_reg_loss / len(test_data) 261 | test_box_loss = test_box_loss / len(test_data) 262 | print('cla_loss: %f, reg_loss: %f, box_loss: %f, #test_data = %d' %(test_cla_loss,test_reg_loss,test_box_loss, len(test_data))) 263 | test_correct = test_correct / len(test_data) 264 | print('mean-correct of 1 sample: %f, #test_data = %d' %(test_correct, len(test_data))) 265 | test_true_correct = test_true_correct / len(test_data) 266 | print('true correct of 1 sample: %f' %(test_true_correct)) 267 | # log 268 | logging.info('Epoch#%d: train error=%e, test error=%e,%e,%e, test correct=%e, %e, lr = %f' %(epoch, train_mse, test_cla_loss,test_reg_loss,test_box_loss, test_correct, test_true_correct, scheduler.get_lr()[0])) 269 | -------------------------------------------------------------------------------- /Dataset.py: -------------------------------------------------------------------------------- 1 | from torch.utils.data import Dataset 2 | from data_classes import PointCloud, Box 3 | from pyquaternion import Quaternion 4 | import numpy as np 5 | import pandas as pd 6 | import os 7 | import torch 8 | from tqdm import tqdm 9 | import kitty_utils as utils 10 | from kitty_utils import getModel 11 | from searchspace import KalmanFiltering 12 | import logging 13 | 14 | 15 | class kittiDataset(): 16 | 17 | def __init__(self, path): 18 | self.KITTI_Folder = path 19 | self.KITTI_velo = os.path.join(self.KITTI_Folder, "velodyne") 20 | self.KITTI_label = os.path.join(self.KITTI_Folder, "label_02") 21 | 22 | def getSceneID(self, split): 23 | if "TRAIN" in split.upper(): # Training SET 24 | if "TINY" in split.upper(): 25 | sceneID = [0] 26 | else: 27 | sceneID = list(range(0, 17)) 28 | elif "VALID" in split.upper(): # Validation Set 29 | if "TINY" in split.upper(): 30 | sceneID = [18] 31 | else: 32 | sceneID = list(range(17, 19)) 33 | elif "TEST" in split.upper(): # Testing Set 34 | if "TINY" in split.upper(): 35 | sceneID = [19] 36 | else: 37 | sceneID = list(range(19, 21)) 38 | 39 | else: # Full Dataset 40 | sceneID = list(range(21)) 41 | return sceneID 42 | 43 | def getBBandPC(self, anno): 44 | calib_path = os.path.join(self.KITTI_Folder, 'calib', 45 | anno['scene'] + ".txt") 46 | calib = self.read_calib_file(calib_path) 47 | transf_mat = np.vstack((calib["Tr_velo_cam"], np.array([0, 0, 0, 1]))) 48 | PC, box = self.getPCandBBfromPandas(anno, transf_mat) 49 | return PC, box 50 | 51 | def getListOfAnno(self, sceneID, category_name="Car"): 52 | list_of_scene = [ 53 | path for path in os.listdir(self.KITTI_velo) 54 | if os.path.isdir(os.path.join(self.KITTI_velo, path)) and 55 | int(path) in sceneID 56 | ] 57 | # print(self.list_of_scene) 58 | list_of_tracklet_anno = [] 59 | for scene in list_of_scene: 60 | 61 | label_file = os.path.join(self.KITTI_label, scene + ".txt") 62 | df = pd.read_csv( 63 | label_file, 64 | sep=' ', 65 | names=[ 66 | "frame", "track_id", "type", "truncated", "occluded", 67 | "alpha", "bbox_left", "bbox_top", "bbox_right", 68 | "bbox_bottom", "height", "width", "length", "x", "y", "z", 69 | "rotation_y" 70 | ]) 71 | df = df[df["type"] == category_name] 72 | df.insert(loc=0, column="scene", value=scene) 73 | for track_id in df.track_id.unique(): 74 | df_tracklet = df[df["track_id"] == track_id] 75 | df_tracklet = df_tracklet.reset_index(drop=True) 76 | tracklet_anno = [anno for index, anno in df_tracklet.iterrows()] 77 | list_of_tracklet_anno.append(tracklet_anno) 78 | 79 | return list_of_tracklet_anno 80 | 81 | def getPCandBBfromPandas(self, box, calib): 82 | center = [box["x"], box["y"] - box["height"] / 2, box["z"]] 83 | size = [box["width"], box["length"], box["height"]] 84 | orientation = Quaternion( 85 | axis=[0, 1, 0], radians=box["rotation_y"]) * Quaternion( 86 | axis=[1, 0, 0], radians=np.pi / 2) 87 | BB = Box(center, size, orientation) 88 | 89 | try: 90 | # VELODYNE PointCloud 91 | velodyne_path = os.path.join(self.KITTI_velo, box["scene"], 92 | '{:06}.bin'.format(box["frame"])) 93 | PC = PointCloud( 94 | np.fromfile(velodyne_path, dtype=np.float32).reshape(-1, 4).T) 95 | PC.transform(calib) 96 | except : 97 | # in case the Point cloud is missing 98 | # (0001/[000177-000180].bin) 99 | PC = PointCloud(np.array([[0, 0, 0]]).T) 100 | 101 | return PC, BB 102 | 103 | def read_calib_file(self, filepath): 104 | """Read in a calibration file and parse into a dictionary.""" 105 | data = {} 106 | with open(filepath, 'r') as f: 107 | for line in f.readlines(): 108 | values = line.split() 109 | # The only non-float values in these files are dates, which 110 | # we don't care about anyway 111 | try: 112 | data[values[0]] = np.array( 113 | [float(x) for x in values[1:]]).reshape(3, 4) 114 | except ValueError: 115 | pass 116 | return data 117 | 118 | 119 | class SiameseDataset(Dataset): 120 | 121 | def __init__(self, 122 | input_size, 123 | path, 124 | split, 125 | category_name="Car", 126 | regress="GAUSSIAN", 127 | offset_BB=0, 128 | scale_BB=1.0): 129 | 130 | self.dataset = kittiDataset(path=path) 131 | 132 | self.input_size = input_size 133 | self.split = split 134 | self.sceneID = self.dataset.getSceneID(split=split) 135 | self.getBBandPC = self.dataset.getBBandPC 136 | 137 | self.category_name = category_name 138 | self.regress = regress 139 | 140 | self.list_of_tracklet_anno = self.dataset.getListOfAnno( 141 | self.sceneID, category_name) 142 | self.list_of_anno = [ 143 | anno for tracklet_anno in self.list_of_tracklet_anno 144 | for anno in tracklet_anno 145 | ] 146 | 147 | def isTiny(self): 148 | return ("TINY" in self.split.upper()) 149 | 150 | def __getitem__(self, index): 151 | return self.getitem(index) 152 | 153 | 154 | class SiameseTrain(SiameseDataset): 155 | 156 | def __init__(self, 157 | input_size, 158 | path, 159 | split="", 160 | category_name="Car", 161 | regress="GAUSSIAN", 162 | sigma_Gaussian=1, 163 | offset_BB=0, 164 | scale_BB=1.0): 165 | super(SiameseTrain,self).__init__( 166 | input_size=input_size, 167 | path=path, 168 | split=split, 169 | category_name=category_name, 170 | regress=regress, 171 | offset_BB=offset_BB, 172 | scale_BB=scale_BB) 173 | 174 | self.sigma_Gaussian = sigma_Gaussian 175 | self.offset_BB = offset_BB 176 | self.scale_BB = scale_BB 177 | 178 | self.num_candidates_perframe = 4 179 | 180 | logging.info("preloading PC...") 181 | self.list_of_PCs = [None] * len(self.list_of_anno) 182 | self.list_of_BBs = [None] * len(self.list_of_anno) 183 | for index in tqdm(range(len(self.list_of_anno))): 184 | anno = self.list_of_anno[index] 185 | PC, box = self.getBBandPC(anno) 186 | new_PC = utils.cropPC(PC, box, offset=10) 187 | 188 | self.list_of_PCs[index] = new_PC 189 | self.list_of_BBs[index] = box 190 | logging.info("PC preloaded!") 191 | 192 | logging.info("preloading Model..") 193 | self.model_PC = [None] * len(self.list_of_tracklet_anno) 194 | for i in tqdm(range(len(self.list_of_tracklet_anno))): 195 | list_of_anno = self.list_of_tracklet_anno[i] 196 | PCs = [] 197 | BBs = [] 198 | cnt = 0 199 | for anno in list_of_anno: 200 | this_PC, this_BB = self.getBBandPC(anno) 201 | PCs.append(this_PC) 202 | BBs.append(this_BB) 203 | anno["model_idx"] = i 204 | anno["relative_idx"] = cnt 205 | cnt += 1 206 | 207 | self.model_PC[i] = getModel( 208 | PCs, BBs, offset=self.offset_BB, scale=self.scale_BB) 209 | 210 | logging.info("Model preloaded!") 211 | 212 | def __getitem__(self, index): 213 | return self.getitem(index) 214 | 215 | def getPCandBBfromIndex(self, anno_idx): 216 | this_PC = self.list_of_PCs[anno_idx] 217 | this_BB = self.list_of_BBs[anno_idx] 218 | return this_PC, this_BB 219 | 220 | def getitem(self, index): 221 | anno_idx = self.getAnnotationIndex(index) 222 | sample_idx = self.getSearchSpaceIndex(index) 223 | 224 | if sample_idx == 0: 225 | sample_offsets = np.zeros(3) 226 | else: 227 | gaussian = KalmanFiltering(bnd=[1, 1, 5]) 228 | sample_offsets = gaussian.sample(1)[0] 229 | 230 | this_anno = self.list_of_anno[anno_idx] 231 | 232 | this_PC, this_BB = self.getPCandBBfromIndex(anno_idx) 233 | sample_BB = utils.getOffsetBB(this_BB, sample_offsets) 234 | 235 | # sample_PC = utils.cropAndCenterPC( 236 | # this_PC, sample_BB, offset=self.offset_BB, scale=self.scale_BB) 237 | sample_PC, sample_label, sample_reg = utils.cropAndCenterPC_label( 238 | this_PC,sample_BB, this_BB, sample_offsets, offset=self.offset_BB, scale=self.scale_BB) 239 | if sample_PC.nbr_points() <= 20: 240 | return self.getitem(np.random.randint(0, self.__len__())) 241 | # sample_PC = utils.regularizePC(sample_PC, self.input_size)[0] 242 | sample_PC, sample_label, sample_reg = utils.regularizePCwithlabel(sample_PC,sample_label,sample_reg,self.input_size) 243 | 244 | if this_anno["relative_idx"] == 0: 245 | prev_idx = 0 246 | fir_idx = 0 247 | else: 248 | prev_idx = anno_idx - 1 249 | fir_idx = anno_idx - this_anno["relative_idx"] 250 | gt_PC_pre, gt_BB_pre = self.getPCandBBfromIndex(prev_idx) 251 | gt_PC_fir, gt_BB_fir = self.getPCandBBfromIndex(fir_idx) 252 | 253 | if sample_idx == 0: 254 | samplegt_offsets = np.zeros(3) 255 | else: 256 | samplegt_offsets = np.random.uniform(low=-0.3, high=0.3, size=3) 257 | samplegt_offsets[2] = samplegt_offsets[2]*5.0 258 | gt_BB_pre = utils.getOffsetBB(gt_BB_pre, samplegt_offsets) 259 | 260 | gt_PC = getModel([gt_PC_fir,gt_PC_pre], [gt_BB_fir,gt_BB_pre], offset=self.offset_BB, scale=self.scale_BB) 261 | 262 | if gt_PC.nbr_points() <= 20: 263 | return self.getitem(np.random.randint(0, self.__len__())) 264 | gt_PC = utils.regularizePC(gt_PC,self.input_size) 265 | 266 | return sample_PC, sample_label, sample_reg, gt_PC 267 | 268 | def __len__(self): 269 | nb_anno = len(self.list_of_anno) 270 | return nb_anno * self.num_candidates_perframe 271 | 272 | def getAnnotationIndex(self, index): 273 | return int(index / (self.num_candidates_perframe)) 274 | 275 | def getSearchSpaceIndex(self, index): 276 | return int(index % self.num_candidates_perframe) 277 | 278 | 279 | class SiameseTest(SiameseDataset): 280 | 281 | def __init__(self, 282 | input_size, 283 | path, 284 | split="", 285 | category_name="Car", 286 | regress="GAUSSIAN", 287 | offset_BB=0, 288 | scale_BB=1.0): 289 | super(SiameseTest,self).__init__( 290 | input_size=input_size, 291 | path=path, 292 | split=split, 293 | category_name=category_name, 294 | regress=regress, 295 | offset_BB=offset_BB, 296 | scale_BB=scale_BB) 297 | self.split = split 298 | self.offset_BB = offset_BB 299 | self.scale_BB = scale_BB 300 | 301 | def getitem(self, index): 302 | list_of_anno = self.list_of_tracklet_anno[index] 303 | PCs = [] 304 | BBs = [] 305 | for anno in list_of_anno: 306 | this_PC, this_BB = self.getBBandPC(anno) 307 | PCs.append(this_PC) 308 | BBs.append(this_BB) 309 | return PCs, BBs, list_of_anno 310 | 311 | def __len__(self): 312 | return len(self.list_of_tracklet_anno) 313 | 314 | 315 | if __name__ == '__main__': 316 | 317 | dataset_Training = SiameseTrain( 318 | input_size=2048, 319 | path='/data/qihaozhe/Kitty_data/training', 320 | split='Tiny_Train', 321 | category_name='Car', 322 | offset_BB=0, 323 | scale_BB=1.15) 324 | aa = dataset_Training.getitem(201) 325 | aa = dataset_Training.getitem(30) 326 | aa = dataset_Training.getitem(100) 327 | aa = dataset_Training.getitem(120) 328 | aa = dataset_Training.getitem(200) 329 | asdf = aa[2].numpy() -------------------------------------------------------------------------------- /pointnet2/utils/pointnet2_utils.py: -------------------------------------------------------------------------------- 1 | from __future__ import ( 2 | division, 3 | absolute_import, 4 | with_statement, 5 | print_function, 6 | unicode_literals, 7 | ) 8 | import torch 9 | from torch.autograd import Function 10 | import torch.nn as nn 11 | import etw_pytorch_utils as pt_utils 12 | import sys 13 | 14 | try: 15 | import builtins 16 | except: 17 | import __builtin__ as builtins 18 | 19 | try: 20 | import pointnet2._ext as _ext 21 | except ImportError: 22 | if not getattr(builtins, "__POINTNET2_SETUP__", False): 23 | raise ImportError( 24 | "Could not import _ext module.\n" 25 | "Please see the setup instructions in the README: " 26 | "https://github.com/erikwijmans/Pointnet2_PyTorch/blob/master/README.rst" 27 | ) 28 | 29 | if False: 30 | # Workaround for type hints without depending on the `typing` module 31 | from typing import * 32 | 33 | 34 | class RandomDropout(nn.Module): 35 | def __init__(self, p=0.5, inplace=False): 36 | super(RandomDropout, self).__init__() 37 | self.p = p 38 | self.inplace = inplace 39 | 40 | def forward(self, X): 41 | theta = torch.Tensor(1).uniform_(0, self.p)[0] 42 | return pt_utils.feature_dropout_no_scaling(X, theta, self.train, self.inplace) 43 | 44 | 45 | class FurthestPointSampling(Function): 46 | @staticmethod 47 | def forward(ctx, xyz, npoint): 48 | # type: (Any, torch.Tensor, int) -> torch.Tensor 49 | r""" 50 | Uses iterative furthest point sampling to select a set of npoint features that have the largest 51 | minimum distance 52 | 53 | Parameters 54 | ---------- 55 | xyz : torch.Tensor 56 | (B, N, 3) tensor where N > npoint 57 | npoint : int32 58 | number of features in the sampled set 59 | 60 | Returns 61 | ------- 62 | torch.Tensor 63 | (B, npoint) tensor containing the set 64 | """ 65 | return _ext.furthest_point_sampling(xyz, npoint) 66 | 67 | @staticmethod 68 | def backward(xyz, a=None): 69 | return None, None 70 | 71 | 72 | furthest_point_sample = FurthestPointSampling.apply 73 | 74 | 75 | class GatherOperation(Function): 76 | @staticmethod 77 | def forward(ctx, features, idx): 78 | # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor 79 | r""" 80 | 81 | Parameters 82 | ---------- 83 | features : torch.Tensor 84 | (B, C, N) tensor 85 | 86 | idx : torch.Tensor 87 | (B, npoint) tensor of the features to gather 88 | 89 | Returns 90 | ------- 91 | torch.Tensor 92 | (B, C, npoint) tensor 93 | """ 94 | 95 | _, C, N = features.size() 96 | 97 | ctx.for_backwards = (idx, C, N) 98 | 99 | return _ext.gather_points(features, idx) 100 | 101 | @staticmethod 102 | def backward(ctx, grad_out): 103 | idx, C, N = ctx.for_backwards 104 | 105 | grad_features = _ext.gather_points_grad(grad_out.contiguous(), idx, N) 106 | return grad_features, None 107 | 108 | 109 | gather_operation = GatherOperation.apply 110 | 111 | 112 | class ThreeNN(Function): 113 | @staticmethod 114 | def forward(ctx, unknown, known): 115 | # type: (Any, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor] 116 | r""" 117 | Find the three nearest neighbors of unknown in known 118 | Parameters 119 | ---------- 120 | unknown : torch.Tensor 121 | (B, n, 3) tensor of known features 122 | known : torch.Tensor 123 | (B, m, 3) tensor of unknown features 124 | 125 | Returns 126 | ------- 127 | dist : torch.Tensor 128 | (B, n, 3) l2 distance to the three nearest neighbors 129 | idx : torch.Tensor 130 | (B, n, 3) index of 3 nearest neighbors 131 | """ 132 | dist2, idx = _ext.three_nn(unknown, known) 133 | 134 | return torch.sqrt(dist2), idx 135 | 136 | @staticmethod 137 | def backward(ctx, a=None, b=None): 138 | return None, None 139 | 140 | 141 | three_nn = ThreeNN.apply 142 | 143 | 144 | class ThreeInterpolate(Function): 145 | @staticmethod 146 | def forward(ctx, features, idx, weight): 147 | # type(Any, torch.Tensor, torch.Tensor, torch.Tensor) -> Torch.Tensor 148 | r""" 149 | Performs weight linear interpolation on 3 features 150 | Parameters 151 | ---------- 152 | features : torch.Tensor 153 | (B, c, m) Features descriptors to be interpolated from 154 | idx : torch.Tensor 155 | (B, n, 3) three nearest neighbors of the target features in features 156 | weight : torch.Tensor 157 | (B, n, 3) weights 158 | 159 | Returns 160 | ------- 161 | torch.Tensor 162 | (B, c, n) tensor of the interpolated features 163 | """ 164 | B, c, m = features.size() 165 | n = idx.size(1) 166 | 167 | ctx.three_interpolate_for_backward = (idx, weight, m) 168 | 169 | return _ext.three_interpolate(features, idx, weight) 170 | 171 | @staticmethod 172 | def backward(ctx, grad_out): 173 | # type: (Any, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor] 174 | r""" 175 | Parameters 176 | ---------- 177 | grad_out : torch.Tensor 178 | (B, c, n) tensor with gradients of ouputs 179 | 180 | Returns 181 | ------- 182 | grad_features : torch.Tensor 183 | (B, c, m) tensor with gradients of features 184 | 185 | None 186 | 187 | None 188 | """ 189 | idx, weight, m = ctx.three_interpolate_for_backward 190 | 191 | grad_features = _ext.three_interpolate_grad( 192 | grad_out.contiguous(), idx, weight, m 193 | ) 194 | 195 | return grad_features, None, None 196 | 197 | 198 | three_interpolate = ThreeInterpolate.apply 199 | 200 | 201 | class GroupingOperation(Function): 202 | @staticmethod 203 | def forward(ctx, features, idx): 204 | # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor 205 | r""" 206 | 207 | Parameters 208 | ---------- 209 | features : torch.Tensor 210 | (B, C, N) tensor of features to group 211 | idx : torch.Tensor 212 | (B, npoint, nsample) tensor containing the indicies of features to group with 213 | 214 | Returns 215 | ------- 216 | torch.Tensor 217 | (B, C, npoint, nsample) tensor 218 | """ 219 | B, nfeatures, nsample = idx.size() 220 | _, C, N = features.size() 221 | 222 | ctx.for_backwards = (idx, N) 223 | 224 | return _ext.group_points(features, idx) 225 | 226 | @staticmethod 227 | def backward(ctx, grad_out): 228 | # type: (Any, torch.tensor) -> Tuple[torch.Tensor, torch.Tensor] 229 | r""" 230 | 231 | Parameters 232 | ---------- 233 | grad_out : torch.Tensor 234 | (B, C, npoint, nsample) tensor of the gradients of the output from forward 235 | 236 | Returns 237 | ------- 238 | torch.Tensor 239 | (B, C, N) gradient of the features 240 | None 241 | """ 242 | idx, N = ctx.for_backwards 243 | 244 | grad_features = _ext.group_points_grad(grad_out.contiguous(), idx, N) 245 | 246 | return grad_features, None 247 | 248 | 249 | grouping_operation = GroupingOperation.apply 250 | 251 | 252 | class BallQuery(Function): 253 | @staticmethod 254 | def forward(ctx, radius, nsample, xyz, new_xyz): 255 | # type: (Any, float, int, torch.Tensor, torch.Tensor) -> torch.Tensor 256 | r""" 257 | 258 | Parameters 259 | ---------- 260 | radius : float 261 | radius of the balls 262 | nsample : int 263 | maximum number of features in the balls 264 | xyz : torch.Tensor 265 | (B, N, 3) xyz coordinates of the features 266 | new_xyz : torch.Tensor 267 | (B, npoint, 3) centers of the ball query 268 | 269 | Returns 270 | ------- 271 | torch.Tensor 272 | (B, npoint, nsample) tensor with the indicies of the features that form the query balls 273 | """ 274 | return _ext.ball_query(new_xyz, xyz, radius, nsample) 275 | 276 | @staticmethod 277 | def backward(ctx, a=None): 278 | return None, None, None, None 279 | 280 | 281 | ball_query = BallQuery.apply 282 | 283 | class BallQuery_score(Function): 284 | @staticmethod 285 | def forward(ctx, radius, nsample, xyz, new_xyz, score): 286 | # type: (Any, float, int, torch.Tensor, torch.Tensor) -> torch.Tensor 287 | r""" 288 | 289 | Parameters 290 | ---------- 291 | radius : float 292 | radius of the balls 293 | nsample : int 294 | maximum number of features in the balls 295 | xyz : torch.Tensor 296 | (B, N, 3) xyz coordinates of the features 297 | new_xyz : torch.Tensor 298 | (B, npoint, 3) centers of the ball query 299 | 300 | Returns 301 | ------- 302 | torch.Tensor 303 | (B, npoint, nsample) tensor with the indicies of the features that form the query balls 304 | """ 305 | return _ext.ball_query_score(new_xyz, xyz, score, radius, nsample) 306 | 307 | @staticmethod 308 | def backward(ctx, a=None): 309 | return None, None, None, None, None 310 | 311 | 312 | ball_query_score = BallQuery_score.apply 313 | 314 | 315 | class QueryAndGroup(nn.Module): 316 | r""" 317 | Groups with a ball query of radius 318 | 319 | Parameters 320 | --------- 321 | radius : float32 322 | Radius of ball 323 | nsample : int32 324 | Maximum number of features to gather in the ball 325 | """ 326 | 327 | def __init__(self, radius, nsample, use_xyz=True): 328 | # type: (QueryAndGroup, float, int, bool) -> None 329 | super(QueryAndGroup, self).__init__() 330 | self.radius, self.nsample, self.use_xyz = radius, nsample, use_xyz 331 | 332 | def forward(self, xyz, new_xyz, features=None): 333 | # type: (QueryAndGroup, torch.Tensor. torch.Tensor, torch.Tensor) -> Tuple[Torch.Tensor] 334 | r""" 335 | Parameters 336 | ---------- 337 | xyz : torch.Tensor 338 | xyz coordinates of the features (B, N, 3) 339 | new_xyz : torch.Tensor 340 | centriods (B, npoint, 3) 341 | features : torch.Tensor 342 | Descriptors of the features (B, C, N) 343 | 344 | Returns 345 | ------- 346 | new_features : torch.Tensor 347 | (B, 3 + C, npoint, nsample) tensor 348 | """ 349 | 350 | idx = ball_query(self.radius, self.nsample, xyz, new_xyz) 351 | xyz_trans = xyz.transpose(1, 2).contiguous() 352 | grouped_xyz = grouping_operation(xyz_trans, idx) # (B, 3, npoint, nsample) 353 | grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1) 354 | 355 | if features is not None: 356 | grouped_features = grouping_operation(features, idx) 357 | if self.use_xyz: 358 | new_features = torch.cat( 359 | [grouped_xyz, grouped_features], dim=1 360 | ) # (B, C + 3, npoint, nsample) 361 | else: 362 | new_features = grouped_features 363 | else: 364 | assert ( 365 | self.use_xyz 366 | ), "Cannot have not features and not use xyz as a feature!" 367 | new_features = grouped_xyz 368 | 369 | return new_features 370 | 371 | class QueryAndGroup_score(nn.Module): 372 | r""" 373 | Groups with a ball query of radius 374 | 375 | Parameters 376 | --------- 377 | radius : float32 378 | Radius of ball 379 | nsample : int32 380 | Maximum number of features to gather in the ball 381 | """ 382 | 383 | def __init__(self, radius, nsample, use_xyz=True): 384 | # type: (QueryAndGroup, float, int, bool) -> None 385 | super(QueryAndGroup_score, self).__init__() 386 | self.radius, self.nsample, self.use_xyz = radius, nsample, use_xyz 387 | 388 | def forward(self, xyz, new_xyz,score, features=None): 389 | # type: (QueryAndGroup, torch.Tensor. torch.Tensor, torch.Tensor) -> Tuple[Torch.Tensor] 390 | r""" 391 | Parameters 392 | ---------- 393 | xyz : torch.Tensor 394 | xyz coordinates of the features (B, N, 3) 395 | new_xyz : torch.Tensor 396 | centriods (B, npoint, 3) 397 | features : torch.Tensor 398 | Descriptors of the features (B, C, N) 399 | 400 | Returns 401 | ------- 402 | new_features : torch.Tensor 403 | (B, 3 + C, npoint, nsample) tensor 404 | """ 405 | 406 | idx = ball_query(self.radius, self.nsample, xyz, new_xyz) 407 | unique_score = ball_query_score(self.radius, self.nsample, xyz, new_xyz, score) 408 | # unique_score1 = torch.zeros((idx.shape[0], idx.shape[1])) 409 | # for i_batch in range(idx.shape[0]): 410 | # for i_region in range(idx.shape[1]): 411 | # unique_ind = torch.unique(idx[i_batch, i_region, :]).long() 412 | # unique_score1[i_batch, i_region] = score[i_batch,unique_ind].sum() 413 | # # score_cpu = score.detach().cpu().numpy() 414 | # # unique_ind_cpu = unique_ind.detach().cpu().numpy() 415 | # # aaa =1 416 | # unique_score_cpu = unique_score.detach().cpu().numpy() 417 | # unique_score1_cpu = unique_score1.detach().cpu().numpy() 418 | # print(unique_score_cpu) 419 | # print('............\n') 420 | # print(unique_score1_cpu) 421 | score_id = unique_score.argmax(dim = 1) 422 | 423 | xyz_trans = xyz.transpose(1, 2).contiguous() 424 | grouped_xyz = grouping_operation(xyz_trans, idx) # (B, 3, npoint, nsample) 425 | grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1) 426 | 427 | if features is not None: 428 | grouped_features = grouping_operation(features, idx) 429 | if self.use_xyz: 430 | new_features = torch.cat( 431 | [grouped_xyz, grouped_features], dim=1 432 | ) # (B, C + 3, npoint, nsample) 433 | else: 434 | new_features = grouped_features 435 | else: 436 | assert ( 437 | self.use_xyz 438 | ), "Cannot have not features and not use xyz as a feature!" 439 | new_features = grouped_xyz 440 | 441 | return new_features, score_id 442 | 443 | 444 | class GroupAll(nn.Module): 445 | r""" 446 | Groups all features 447 | 448 | Parameters 449 | --------- 450 | """ 451 | 452 | def __init__(self, use_xyz=True): 453 | # type: (GroupAll, bool) -> None 454 | super(GroupAll, self).__init__() 455 | self.use_xyz = use_xyz 456 | 457 | def forward(self, xyz, new_xyz, features=None): 458 | # type: (GroupAll, torch.Tensor, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor] 459 | r""" 460 | Parameters 461 | ---------- 462 | xyz : torch.Tensor 463 | xyz coordinates of the features (B, N, 3) 464 | new_xyz : torch.Tensor 465 | Ignored 466 | features : torch.Tensor 467 | Descriptors of the features (B, C, N) 468 | 469 | Returns 470 | ------- 471 | new_features : torch.Tensor 472 | (B, C + 3, 1, N) tensor 473 | """ 474 | 475 | grouped_xyz = xyz.transpose(1, 2).unsqueeze(2) 476 | if features is not None: 477 | grouped_features = features.unsqueeze(2) 478 | if self.use_xyz: 479 | new_features = torch.cat( 480 | [grouped_xyz, grouped_features], dim=1 481 | ) # (B, 3 + C, 1, N) 482 | else: 483 | new_features = grouped_features 484 | else: 485 | new_features = grouped_xyz 486 | 487 | return new_features 488 | -------------------------------------------------------------------------------- /kitty_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import os 3 | import copy 4 | import numpy as np 5 | from pyquaternion import Quaternion 6 | from data_classes import PointCloud 7 | from metrics import estimateOverlap 8 | from scipy.optimize import leastsq 9 | 10 | 11 | def distanceBB_Gaussian(box1, box2, sigma=1): 12 | off1 = np.array([ 13 | box1.center[0], box1.center[2], 14 | Quaternion(matrix=box1.rotation_matrix).degrees 15 | ]) 16 | off2 = np.array([ 17 | box2.center[0], box2.center[2], 18 | Quaternion(matrix=box2.rotation_matrix).degrees 19 | ]) 20 | dist = np.linalg.norm(off1 - off2) 21 | score = np.exp(-0.5 * (dist) / (sigma * sigma)) 22 | return score 23 | 24 | 25 | # IoU or Gaussian score map 26 | def getScoreGaussian(offset, sigma=1): 27 | coeffs = [1, 1, 1 / 5] 28 | dist = np.linalg.norm(np.multiply(offset, coeffs)) 29 | score = np.exp(-0.5 * (dist) / (sigma * sigma)) 30 | return torch.tensor([score]) 31 | 32 | 33 | def getScoreIoU(a, b): 34 | score = estimateOverlap(a, b) 35 | return torch.tensor([score]) 36 | 37 | 38 | def getScoreHingeIoU(a, b): 39 | score = estimateOverlap(a, b) 40 | if score < 0.5: 41 | score = 0.0 42 | return torch.tensor([score]) 43 | 44 | 45 | def getOffsetBB(box, offset): 46 | rot_quat = Quaternion(matrix=box.rotation_matrix) 47 | trans = np.array(box.center) 48 | 49 | new_box = copy.deepcopy(box) 50 | 51 | new_box.translate(-trans) 52 | new_box.rotate(rot_quat.inverse) 53 | 54 | # REMOVE TRANSfORM 55 | if len(offset) == 3: 56 | new_box.rotate( 57 | Quaternion(axis=[0, 0, 1], angle=offset[2] * np.pi / 180)) 58 | elif len(offset) == 4: 59 | new_box.rotate( 60 | Quaternion(axis=[0, 0, 1], angle=offset[3] * np.pi / 180)) 61 | if offset[0]>new_box.wlh[0]: 62 | offset[0] = np.random.uniform(-1,1) 63 | if offset[1]>min(new_box.wlh[1],2): 64 | offset[1] = np.random.uniform(-1,1) 65 | new_box.translate(np.array([offset[0], offset[1], 0])) 66 | 67 | # APPLY PREVIOUS TRANSFORMATION 68 | new_box.rotate(rot_quat) 69 | new_box.translate(trans) 70 | return new_box 71 | 72 | 73 | def voxelize(PC, dim_size=[48, 108, 48]): 74 | # PC = normalizePC(PC) 75 | if np.isscalar(dim_size): 76 | dim_size = [dim_size] * 3 77 | dim_size = np.atleast_2d(dim_size).T 78 | PC = (PC + 0.5) * dim_size 79 | # truncate to integers 80 | xyz = PC.astype(np.int) 81 | # discard voxels that fall outside dims 82 | valid_ix = ~np.any((xyz < 0) | (xyz >= dim_size), 0) 83 | xyz = xyz[:, valid_ix] 84 | out = np.zeros(dim_size.flatten(), dtype=np.float32) 85 | out[tuple(xyz)] = 1 86 | # print(out) 87 | return out 88 | 89 | 90 | def regularizePC2(input_size, PC,): 91 | return regularizePC(PC=PC, input_size=input_size) 92 | 93 | 94 | def regularizePC(PC,input_size,istrain=True): 95 | PC = np.array(PC.points, dtype=np.float32) 96 | if np.shape(PC)[1] > 2: 97 | if PC.shape[0] > 3: 98 | PC = PC[0:3, :] 99 | if PC.shape[1] != int(input_size/2): 100 | if not istrain: 101 | np.random.seed(1) 102 | new_pts_idx = np.random.randint( 103 | low=0, high=PC.shape[1], size=int(input_size/2), dtype=np.int64) 104 | PC = PC[:, new_pts_idx] 105 | PC = PC.reshape((3, int(input_size/2))).T 106 | 107 | else: 108 | PC = np.zeros((3, int(input_size/2))).T 109 | 110 | return torch.from_numpy(PC).float() 111 | 112 | def regularizePCwithlabel(PC,label,reg, input_size,istrain=True): 113 | PC = np.array(PC.points, dtype=np.float32) 114 | if np.shape(PC)[1] > 2: 115 | if PC.shape[0] > 3: 116 | PC = PC[0:3, :] 117 | if PC.shape[1] != input_size: 118 | if not istrain: 119 | np.random.seed(1) 120 | new_pts_idx = np.random.randint( 121 | low=0, high=PC.shape[1], size=input_size, dtype=np.int64) 122 | PC = PC[:, new_pts_idx] 123 | label = label[new_pts_idx] 124 | PC = PC.reshape((3, input_size)).T 125 | label = label[0:128] 126 | reg = np.tile(reg,[np.size(label),1]) 127 | 128 | else: 129 | PC = np.zeros((3, input_size)).T 130 | label = np.zeros(128) 131 | reg = np.tile(reg,[np.size(label),1]) 132 | 133 | return torch.from_numpy(PC).float(),torch.from_numpy(label).float(),torch.from_numpy(reg).float() 134 | 135 | 136 | def getModel(PCs, boxes, offset=0, scale=1.0, normalize=False): 137 | 138 | if len(PCs) == 0: 139 | return PointCloud(np.ones((3, 0))) 140 | points = np.ones((PCs[0].points.shape[0], 0)) 141 | 142 | for PC, box in zip(PCs, boxes): 143 | cropped_PC = cropAndCenterPC( 144 | PC, box, offset=offset, scale=scale, normalize=normalize) 145 | # try: 146 | if cropped_PC.points.shape[1] > 0: 147 | points = np.concatenate([points, cropped_PC.points], axis=1) 148 | 149 | PC = PointCloud(points) 150 | 151 | return PC 152 | 153 | 154 | def cropPC(PC, box, offset=0, scale=1.0): 155 | box_tmp = copy.deepcopy(box) 156 | box_tmp.wlh = box_tmp.wlh * scale 157 | maxi = np.max(box_tmp.corners(), 1) + offset 158 | mini = np.min(box_tmp.corners(), 1) - offset 159 | 160 | x_filt_max = PC.points[0, :] < maxi[0] 161 | x_filt_min = PC.points[0, :] > mini[0] 162 | y_filt_max = PC.points[1, :] < maxi[1] 163 | y_filt_min = PC.points[1, :] > mini[1] 164 | z_filt_max = PC.points[2, :] < maxi[2] 165 | z_filt_min = PC.points[2, :] > mini[2] 166 | 167 | close = np.logical_and(x_filt_min, x_filt_max) 168 | close = np.logical_and(close, y_filt_min) 169 | close = np.logical_and(close, y_filt_max) 170 | close = np.logical_and(close, z_filt_min) 171 | close = np.logical_and(close, z_filt_max) 172 | 173 | new_PC = PointCloud(PC.points[:, close]) 174 | return new_PC 175 | 176 | def getlabelPC(PC, box, offset=0, scale=1.0): 177 | box_tmp = copy.deepcopy(box) 178 | new_PC = PointCloud(PC.points.copy()) 179 | rot_mat = np.transpose(box_tmp.rotation_matrix) 180 | trans = -box_tmp.center 181 | 182 | # align data 183 | new_PC.translate(trans) 184 | box_tmp.translate(trans) 185 | new_PC.rotate((rot_mat)) 186 | box_tmp.rotate(Quaternion(matrix=(rot_mat))) 187 | 188 | box_tmp.wlh = box_tmp.wlh * scale 189 | maxi = np.max(box_tmp.corners(), 1) + offset 190 | mini = np.min(box_tmp.corners(), 1) - offset 191 | 192 | x_filt_max = new_PC.points[0, :] < maxi[0] 193 | x_filt_min = new_PC.points[0, :] > mini[0] 194 | y_filt_max = new_PC.points[1, :] < maxi[1] 195 | y_filt_min = new_PC.points[1, :] > mini[1] 196 | z_filt_max = new_PC.points[2, :] < maxi[2] 197 | z_filt_min = new_PC.points[2, :] > mini[2] 198 | 199 | close = np.logical_and(x_filt_min, x_filt_max) 200 | close = np.logical_and(close, y_filt_min) 201 | close = np.logical_and(close, y_filt_max) 202 | close = np.logical_and(close, z_filt_min) 203 | close = np.logical_and(close, z_filt_max) 204 | 205 | new_label = np.zeros(new_PC.points.shape[1]) 206 | new_label[close] = 1 207 | return new_label 208 | 209 | def cropPCwithlabel(PC, box,label, offset=0, scale=1.0): 210 | box_tmp = copy.deepcopy(box) 211 | box_tmp.wlh = box_tmp.wlh * scale 212 | maxi = np.max(box_tmp.corners(), 1) + offset 213 | mini = np.min(box_tmp.corners(), 1) - offset 214 | 215 | x_filt_max = PC.points[0, :] < maxi[0] 216 | x_filt_min = PC.points[0, :] > mini[0] 217 | y_filt_max = PC.points[1, :] < maxi[1] 218 | y_filt_min = PC.points[1, :] > mini[1] 219 | z_filt_max = PC.points[2, :] < maxi[2] 220 | z_filt_min = PC.points[2, :] > mini[2] 221 | 222 | close = np.logical_and(x_filt_min, x_filt_max) 223 | close = np.logical_and(close, y_filt_min) 224 | close = np.logical_and(close, y_filt_max) 225 | close = np.logical_and(close, z_filt_min) 226 | close = np.logical_and(close, z_filt_max) 227 | 228 | new_PC = PointCloud(PC.points[:, close]) 229 | new_label = label[close] 230 | return new_PC,new_label 231 | 232 | def weight_process(include,low,high): 233 | if includehigh: 236 | weight = 1 237 | else: 238 | weight = (include*2.0+3.0*high-5.0*low)/(5*(high-low)) 239 | return weight 240 | 241 | def func(a, x): 242 | k, b = a 243 | return k * x + b 244 | def dist(a, x, y): 245 | return func(a, x) - y 246 | 247 | def weight_process2(k): 248 | k = abs(k) 249 | if k>1: 250 | weight = 0.7 251 | else: 252 | weight = 1-0.3*k 253 | return weight 254 | 255 | 256 | 257 | def cropAndCenterPC(PC, box, offset=0, scale=1.0, normalize=False): 258 | 259 | new_PC = cropPC(PC, box, offset=2 * offset, scale=4 * scale) 260 | 261 | new_box = copy.deepcopy(box) 262 | 263 | rot_mat = np.transpose(new_box.rotation_matrix) 264 | trans = -new_box.center 265 | 266 | # align data 267 | new_PC.translate(trans) 268 | new_box.translate(trans) 269 | new_PC.rotate((rot_mat)) 270 | new_box.rotate(Quaternion(matrix=(rot_mat))) 271 | 272 | # crop around box 273 | new_PC = cropPC(new_PC, new_box, offset=offset, scale=scale) 274 | 275 | if normalize: 276 | new_PC.normalize(box.wlh) 277 | return new_PC 278 | 279 | def Centerbox(sample_box, gt_box): 280 | rot_mat = np.transpose(gt_box.rotation_matrix) 281 | trans = -gt_box.center 282 | 283 | new_box = copy.deepcopy(sample_box) 284 | new_box.translate(trans) 285 | new_box.rotate(Quaternion(matrix=(rot_mat))) 286 | 287 | return new_box 288 | 289 | 290 | def cropAndCenterPC_label(PC, sample_box, gt_box,sample_offsets, offset=0, scale=1.0, normalize=False): 291 | 292 | new_PC = cropPC(PC, sample_box, offset=2 * offset, scale=4 * scale) 293 | 294 | new_box = copy.deepcopy(sample_box) 295 | 296 | new_label = getlabelPC(new_PC, gt_box, offset=offset, scale=scale) 297 | new_box_gt = copy.deepcopy(gt_box) 298 | # new_box_gt2 = copy.deepcopy(gt_box) 299 | 300 | #rot_quat = Quaternion(matrix=new_box.rotation_matrix) 301 | rot_mat = np.transpose(new_box.rotation_matrix) 302 | trans = -new_box.center 303 | 304 | # align data 305 | new_PC.translate(trans) 306 | new_box.translate(trans) 307 | new_PC.rotate((rot_mat)) 308 | new_box.rotate(Quaternion(matrix=(rot_mat))) 309 | 310 | new_box_gt.translate(trans) 311 | new_box_gt.rotate(Quaternion(matrix=(rot_mat))) 312 | 313 | # crop around box 314 | new_PC, new_label = cropPCwithlabel(new_PC, new_box, new_label, offset=offset+2.0, scale=1 * scale) 315 | #new_PC, new_label = cropPCwithlabel(new_PC, new_box, new_label, offset=offset+0.6, scale=1 * scale) 316 | 317 | label_reg = [new_box_gt.center[0],new_box_gt.center[1],new_box_gt.center[2],-sample_offsets[2]] 318 | label_reg = np.array(label_reg) 319 | # label_reg = np.tile(label_reg,[np.size(new_label),1]) 320 | 321 | if normalize: 322 | new_PC.normalize(sample_box.wlh) 323 | return new_PC, new_label, label_reg 324 | 325 | def cropAndCenterPC_label_test_time(PC, sample_box, offset=0, scale=1.0): 326 | 327 | new_PC = cropPC(PC, sample_box, offset=2 * offset, scale=4 * scale) 328 | 329 | new_box = copy.deepcopy(sample_box) 330 | 331 | rot_quat = Quaternion(matrix=new_box.rotation_matrix) 332 | rot_mat = np.transpose(new_box.rotation_matrix) 333 | trans = -new_box.center 334 | 335 | # align data 336 | new_PC.translate(trans) 337 | new_box.translate(trans) 338 | new_PC.rotate((rot_mat)) 339 | new_box.rotate(Quaternion(matrix=(rot_mat))) 340 | 341 | # crop around box 342 | new_PC = cropPC(new_PC, new_box, offset=offset+2.0, scale=scale) 343 | 344 | return new_PC 345 | 346 | def cropAndCenterPC_label_test(PC, sample_box, gt_box, offset=0, scale=1.0, normalize=False): 347 | 348 | new_PC = cropPC(PC, sample_box, offset=2 * offset, scale=4 * scale) 349 | 350 | new_box = copy.deepcopy(sample_box) 351 | 352 | new_label = getlabelPC(new_PC, gt_box, offset=offset, scale=scale) 353 | new_box_gt = copy.deepcopy(gt_box) 354 | # new_box_gt2 = copy.deepcopy(gt_box) 355 | 356 | rot_quat = Quaternion(matrix=new_box.rotation_matrix) 357 | rot_mat = np.transpose(new_box.rotation_matrix) 358 | trans = -new_box.center 359 | 360 | # align data 361 | new_PC.translate(trans) 362 | new_box.translate(trans) 363 | new_PC.rotate((rot_mat)) 364 | new_box.rotate(Quaternion(matrix=(rot_mat))) 365 | 366 | new_box_gt.translate(trans) 367 | new_box_gt.rotate(Quaternion(matrix=(rot_mat))) 368 | # new_box_gt2.translate(trans) 369 | # new_box_gt2.rotate(rot_quat.inverse) 370 | 371 | # crop around box 372 | new_PC, new_label = cropPCwithlabel(new_PC, new_box, new_label, offset=offset+2.0, scale=1 * scale) 373 | #new_PC, new_label = cropPCwithlabel(new_PC, new_box, new_label, offset=offset+0.6, scale=1 * scale) 374 | 375 | label_reg = [new_box_gt.center[0],new_box_gt.center[1],new_box_gt.center[2]] 376 | label_reg = np.array(label_reg) 377 | # label_reg = (new_PC.points - np.tile(new_box_gt.center,[np.size(new_label),1]).T) * np.expand_dims(new_label, axis=0) 378 | # newoff = [new_box_gt.center[0],new_box_gt.center[1],new_box_gt.center[2]] 379 | # newoff = np.array(newoff) 380 | 381 | if normalize: 382 | new_PC.normalize(sample_box.wlh) 383 | return new_PC, new_label, label_reg, new_box, new_box_gt 384 | 385 | def distanceBB(box1, box2): 386 | 387 | eucl = np.linalg.norm(box1.center - box2.center) 388 | angl = Quaternion.distance( 389 | Quaternion(matrix=box1.rotation_matrix), 390 | Quaternion(matrix=box2.rotation_matrix)) 391 | return eucl + angl 392 | 393 | 394 | def generate_boxes(box, search_space=[[0, 0, 0]]): 395 | # Geenrate more candidate boxes based on prior and search space 396 | # Input : Prior position, search space and seaarch size 397 | # Output : List of boxes 398 | 399 | candidate_boxes = [getOffsetBB(box, offset) for offset in search_space] 400 | return candidate_boxes 401 | 402 | 403 | def getDataframeGT(anno): 404 | df = { 405 | "scene": anno["scene"], 406 | "frame": anno["frame"], 407 | "track_id": anno["track_id"], 408 | "type": anno["type"], 409 | "truncated": anno["truncated"], 410 | "occluded": anno["occluded"], 411 | "alpha": anno["alpha"], 412 | "bbox_left": anno["bbox_left"], 413 | "bbox_top": anno["bbox_top"], 414 | "bbox_right": anno["bbox_right"], 415 | "bbox_bottom": anno["bbox_bottom"], 416 | "height": anno["height"], 417 | "width": anno["width"], 418 | "length": anno["length"], 419 | "x": anno["x"], 420 | "y": anno["y"], 421 | "z": anno["z"], 422 | "rotation_y": anno["rotation_y"] 423 | } 424 | return df 425 | 426 | 427 | def getDataframe(anno, box, score): 428 | myquat = (box.orientation * Quaternion(axis=[1, 0, 0], radians=-np.pi / 2)) 429 | df = { 430 | "scene": anno["scene"], 431 | "frame": anno["frame"], 432 | "track_id": anno["track_id"], 433 | "type": anno["type"], 434 | "truncated": anno["truncated"], 435 | "occluded": anno["occluded"], 436 | "alpha": 0.0, 437 | "bbox_left": 0.0, 438 | "bbox_top": 0.0, 439 | "bbox_right": 0.0, 440 | "bbox_bottom": 0.0, 441 | "height": box.wlh[2], 442 | "width": box.wlh[0], 443 | "length": box.wlh[1], 444 | "x": box.center[0], 445 | "y": box.center[1] + box.wlh[2] / 2, 446 | "z": box.center[2], 447 | "rotation_y": 448 | np.sign(myquat.axis[1]) * myquat.radians, # this_anno["rotation_y"], # 449 | "score": score 450 | } 451 | return df 452 | 453 | 454 | def saveTrackingResults(df_3D, dataset_loader, export=None, epoch=-1): 455 | 456 | for i_scene, scene in enumerate(df_3D.scene.unique()): 457 | new_df_3D = df_3D[df_3D["scene"] == scene] 458 | new_df_3D = new_df_3D.drop(["scene"], axis=1) 459 | new_df_3D = new_df_3D.sort_values(by=['frame', 'track_id']) 460 | 461 | os.makedirs(os.path.join("results", export, "data"), exist_ok=True) 462 | if epoch == -1: 463 | path = os.path.join("results", export, "data", "{}.txt".format(scene)) 464 | else: 465 | path = os.path.join("results", export, "data", 466 | "{}_epoch{}.txt".format(scene,epoch)) 467 | 468 | new_df_3D.to_csv( 469 | path, sep=" ", header=False, index=False, float_format='%.6f') 470 | --------------------------------------------------------------------------------