From 26710a72538fe953628559ffa11a702198fa3d19 Mon Sep 17 00:00:00 2001 From: Shiming Date: Mon, 26 May 2025 19:52:55 +0000 Subject: [PATCH 01/10] added README for VoteFlow --- README.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/README.md b/README.md index ed2af54..5cc3f7f 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,11 @@ OpenSceneFlow is a codebase for point cloud scene flow estimation. It is also an official implementation of the following papers (sored by the time of publication): +- **VoteFlow:VoteFlow: Enforcing Local Rigidity in Self-Supervised Scene Flow** +*Yancong Lin\*, Shiming Wang\*, Liangliang Nan, Julian Kooij, Holger Caesar* +Conference on Computer Vision and Pattern Recognition (**CVPR**) 2025 +[ Backbone ] [ Self-Supervised ] - [ [arXiv](https://arxiv.org/abs/2503.22328) ] [ [Project](https://github.com/tudelft-iv/VoteFlow/)] + - **HiMo: High-Speed Objects Motion Compensation in Point Clouds** (SeFlow++) *Qingwen Zhang, Ajinkya Khoche, Yi Yang, Li Ling, Sina Sharif Mansouri, Olov Andersson, Patric Jensfelt* Preprint; Under review; 2025 From cab34623f3d37346d118290d66a5978a3b047e93 Mon Sep 17 00:00:00 2001 From: Shiming Date: Mon, 26 May 2025 20:42:05 +0000 Subject: [PATCH 02/10] adapted the voteflow code into OpenSceneFlow and test training needed --- .gitignore | 5 +- README.md | 14 + src/models/basic/voteflow_plugin/__init__.py | 3 + .../hough_transformation/__init__.py | 4 + .../cpp_im2ht/cuda_utils.h | 77 ++++ .../hough_transformation/cpp_im2ht/ht_cuda.cu | 116 ++++++ .../hough_transformation/cpp_im2ht/ht_cuda.h | 29 ++ .../cpp_im2ht/ht_setup.sh | 7 + .../hough_transformation/cpp_im2ht/im2ht.cpp | 64 +++ .../hough_transformation/cpp_im2ht/im2ht.h | 28 ++ .../cpp_im2ht/im2ht_cuda.cuh | 206 ++++++++++ .../hough_transformation/cpp_im2ht/setup.py | 37 ++ .../hough_transformation/ht_cuda.py | 63 +++ .../hough_transformation/im2ht.py | 118 ++++++ .../hough_transformation/test_backward.py | 73 ++++ .../hough_transformation/test_forward.py | 247 ++++++++++++ src/models/basic/voteflow_plugin/utils.py | 48 +++ .../basic/voteflow_plugin/voteflow_module.py | 198 +++++++++ src/models/voteflow.py | 376 ++++++++++++++++++ 19 files changed, 1712 insertions(+), 1 deletion(-) create mode 100644 src/models/basic/voteflow_plugin/__init__.py create mode 100755 src/models/basic/voteflow_plugin/hough_transformation/__init__.py create mode 100755 src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/cuda_utils.h create mode 100755 src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_cuda.cu create mode 100755 src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_cuda.h create mode 100755 src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_setup.sh create mode 100755 src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht.cpp create mode 100755 src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht.h create mode 100755 src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht_cuda.cuh create mode 100755 src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/setup.py create mode 100755 src/models/basic/voteflow_plugin/hough_transformation/ht_cuda.py create mode 100755 src/models/basic/voteflow_plugin/hough_transformation/im2ht.py create mode 100755 src/models/basic/voteflow_plugin/hough_transformation/test_backward.py create mode 100644 src/models/basic/voteflow_plugin/hough_transformation/test_forward.py create mode 100644 src/models/basic/voteflow_plugin/utils.py create mode 100644 src/models/basic/voteflow_plugin/voteflow_module.py create mode 100644 src/models/voteflow.py diff --git a/.gitignore b/.gitignore index 7eed5e2..744215b 100644 --- a/.gitignore +++ b/.gitignore @@ -25,4 +25,7 @@ cancel.sh # cuda build files *.egg-info build* -dist* \ No newline at end of file +dist* + +data +*__pycache__ \ No newline at end of file diff --git a/README.md b/README.md index 5cc3f7f..08d6bee 100644 --- a/README.md +++ b/README.md @@ -130,6 +130,13 @@ And free yourself from trainning, you can download the pretrained weight from [H mamba activate opensf ``` +### VoteFLow +Extra pakcges needed for VoteFlow, [pytorch3d](https://pytorch3d.org/) (prefer 0.7.7) and [torch-scatter](https://github.com/rusty1s/pytorch_scatter?tab=readme-ov-file) (prefer 2.1.2): +```bash +pip install pytorch3d +pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.0+cu117.html +``` + ### Flow4D Train Flow4D with the leaderboard submit config. [Runtime: Around 18 hours in 4x RTX 3090 GPUs.] @@ -316,6 +323,13 @@ And our excellent collaborators works contributed to this codebase also: journal={arXiv preprint arXiv:2501.17821}, year={2025} } + +@inproceedings{lin2025voteflow, + title={VoteFlow: Enforcing Local Rigidity in Self-Supervised Scene Flow}, + author={Lin, Yancong and Wang, Shiming and Nan, Liangliang and Kooij, Julian and Caesar, Holger}, + booktitle={CVPR}, + year={2025}, +} ``` Thank you for your support! ❤️ diff --git a/src/models/basic/voteflow_plugin/__init__.py b/src/models/basic/voteflow_plugin/__init__.py new file mode 100644 index 0000000..02e9652 --- /dev/null +++ b/src/models/basic/voteflow_plugin/__init__.py @@ -0,0 +1,3 @@ +from .hough_transformation import HT_CUDA +from .utils import calculate_unq_voxels, batched_masked_gather, pad_to_batch +from .voteflow_module import VolConvBN, Decoder, FastFlow3DUNet \ No newline at end of file diff --git a/src/models/basic/voteflow_plugin/hough_transformation/__init__.py b/src/models/basic/voteflow_plugin/hough_transformation/__init__.py new file mode 100755 index 0000000..be0a4d0 --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/__init__.py @@ -0,0 +1,4 @@ +from .ht_cuda import HT_CUDA +__all__ = [ + 'HT_CUDA', +] diff --git a/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/cuda_utils.h b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/cuda_utils.h new file mode 100755 index 0000000..2a231e1 --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/cuda_utils.h @@ -0,0 +1,77 @@ +// Copyright (c) Facebook, Inc. and its affiliates. +// +// This source code is licensed under the MIT license found in the +// LICENSE file in the root directory of this source tree. + +// #include +// #include +// #include +#include +#include + + +#define CUDA_KERNEL_LOOP(i, n) \ + for (int i = blockIdx.x * blockDim.x + threadIdx.x; \ + i < (n); \ + i += blockDim.x * gridDim.x) + +// CUDA: thread number configuration. +// Use 1024 threads per block, which requires cuda sm_2x or above, +// or fall back to attempt compatibility (best of luck to you). +//#if __CUDA_ARCH__ >= 200 +// const int CUDA_NUM_THREADS = 1024; +//#else +// const int CUDA_NUM_THREADS = 512; +//#endif +const int CUDA_NUM_THREADS = 1024; + +inline int GET_BLOCKS(const int N) +{ + return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS; +} + + +// #ifndef CUDA_NUM_THREADS +// #define CUDA_NUM_THREADS 1024 +// #endif + +// #define CUDA_KERNEL_LOOP(i, n) \ +// for (int i = blockIdx.x * blockDim.x + threadIdx.x; \ +// i < (n); \ +// i += blockDim.x * gridDim.x) + +// const int CUDA_NUM_THREADS = 1024; + +// inline int GET_BLOCKS(const int N) +// { +// return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS; +// } + + +// // C++ interface +// // NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. +// // #define CHECK_CUDA(x) AT_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") +// #define CHECK_CONTIGUOUS(x) AT_CHECK(x.is_contiguous(), #x " must be contiguous") +// #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) + +// #define CHECK_CUDA(x) \ +// do { \ +// AT_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor"); \ +// } while (0) + +// #define CHECK_CONTIGUOUS(x) \ +// do { \ +// AT_CHECK(x.is_contiguous(), #x " must be a contiguous tensor"); \ +// } while (0) + +// #define CHECK_IS_INT(x) \ +// do { \ +// AT_CHECK(x.scalar_type() == at::ScalarType::Int, \ +// #x " must be an int tensor"); \ +// } while (0) + +// #define CHECK_IS_FLOAT(x) \ +// do { \ +// AT_CHECK(x.scalar_type() == at::ScalarType::Float, \ +// #x " must be a float tensor"); \ +// } while (0) diff --git a/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_cuda.cu b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_cuda.cu new file mode 100755 index 0000000..095be69 --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_cuda.cu @@ -0,0 +1,116 @@ +#include +#include +#include +#include +#include +#include +#include "im2ht_cuda.cuh" + + +at::Tensor +ht_cuda_forward( + const at::Tensor &feats_src_dst, + const at::Tensor &voxels_src, + const at::Tensor &voxels_dst, + const at::Tensor &idxs_src, + const at::Tensor &idxs_dst, + const int h, + const int w, + const int d + ) +{ + AT_ASSERTM(feats_src_dst.is_contiguous(), "feats tensor has to be contiguous"); + AT_ASSERTM(voxels_src.is_contiguous(), "voxels tensor has to be contiguous"); + AT_ASSERTM(voxels_dst.is_contiguous(), "voxels tensor has to be contiguous"); + AT_ASSERTM(idxs_src.is_contiguous(), "idxs tensor has to be contiguous"); + AT_ASSERTM(idxs_dst.is_contiguous(), "idxs tensor has to be contiguous"); + + // AT_ASSERTM(input.type().is_cuda(), "input must be a CUDA tensor"); + + const int b = feats_src_dst.size(0); + const int l = feats_src_dst.size(1); // num of points + const int c = feats_src_dst.size(-1); // num of channels + const int m = idxs_src.size(2); // m + const int n = idxs_dst.size(2); // n + + // AT_ASSERTM(m_ == m && n_ <= n && k_ == k, + // "input shape and predefined shape do not match: (%04d x %04d x %04d vs %04d x %04d x %04d).", k_, m_, n_, k, m, n); + + // printf("dimensions b=%06d, l=%06d, m=%06d, n=%06d, h=%06d, w= %06d, d= %06d \n", + // b, l, m, n, h, w, d); + auto vol_ht = at::zeros({b, l, c, h, w}, feats_src_dst.options()); + + AT_DISPATCH_FLOATING_TYPES(vol_ht.type(), "im2ht_cuda_forward", ([&] { + im2ht_cuda_forward(at::cuda::getCurrentCUDAStream(), + feats_src_dst.data() , + vol_ht.data(), + voxels_src.data() , + voxels_dst.data(), + idxs_src.data(), + idxs_dst.data(), + b, l, c, + m, n, + h, w, d + ); + + })); + // std::cout <<"output" < +at::Tensor +ht_cuda_backward( + const at::Tensor &grad_vol, + const at::Tensor &voxels_src, + const at::Tensor &voxels_dst, + const at::Tensor &idxs_src, + const at::Tensor &idxs_dst, + const int h, + const int w, + const int d + ) +{ + + AT_ASSERTM(grad_vol.is_contiguous(), "grad_output tensor has to be contiguous"); + AT_ASSERTM(voxels_src.is_contiguous(), "voxels tensor has to be contiguous"); + AT_ASSERTM(voxels_dst.is_contiguous(), "voxels tensor has to be contiguous"); + AT_ASSERTM(idxs_src.is_contiguous(), "idxs tensor has to be contiguous"); + AT_ASSERTM(idxs_dst.is_contiguous(), "idxs tensor has to be contiguous"); + + // AT_ASSERTM(grad_output.type().is_cuda(), "grad_output must be a CUDA tensor"); + + // grad_output: [b, l, c, h, w] + const int b = grad_vol.size(0); + const int l = grad_vol.size(1); + const int c = grad_vol.size(2); + + const int m = idxs_src.size(2); // m + const int n = idxs_dst.size(2); // n + + // printf("h_ = %04d, w_ = %04d, d_ = %04d vs h = %04d, w = %04d, d = %04d\n", h_, w_, d_, h, w, d); + // AT_ASSERTM(h_ == h && w_ == w && d_== d, + // "grad_out shape and predefined shape do not match: ", h_, ' ', w_, ' ', d_, 'vs', h, ' ', w, ' ', d); + + auto grad_feats_src_dst = at::zeros({b, l, n, c}, grad_vol.options()); + + AT_DISPATCH_FLOATING_TYPES(grad_vol.type(), "im2ht_cuda_backward", ([&] { + im2ht_cuda_backward(at::cuda::getCurrentCUDAStream(), + grad_feats_src_dst.data(), + grad_vol.data(), + voxels_src.data(), + voxels_dst.data(), + idxs_src.data(), + idxs_dst.data(), + b, l, c, + m, n, + h, w, d + ); + + })); + + return grad_feats_src_dst; +} diff --git a/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_cuda.h b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_cuda.h new file mode 100755 index 0000000..f8e0300 --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_cuda.h @@ -0,0 +1,29 @@ +#pragma once +#include + + +at::Tensor +ht_cuda_forward( + const at::Tensor &feats_src_dst, + const at::Tensor &voxels_src, + const at::Tensor &voxels_dst, + const at::Tensor &idxs_src, + const at::Tensor &idxs_dst, + const int h, + const int w, + const int d + ); + +// std::vector +at::Tensor +ht_cuda_backward( + const at::Tensor &grad_output, + const at::Tensor &voxels_src, + const at::Tensor &voxels_dst, + const at::Tensor &idxs_src, + const at::Tensor &idxs_dst, + const int h, + const int w, + const int d + ); + diff --git a/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_setup.sh b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_setup.sh new file mode 100755 index 0000000..73b833d --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_setup.sh @@ -0,0 +1,7 @@ +cd cpp_im2ht +rm -rf build +python setup.py clean +python setup.py build +python setup.py install --user +cd .. + diff --git a/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht.cpp b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht.cpp new file mode 100755 index 0000000..7ef830c --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht.cpp @@ -0,0 +1,64 @@ +#include "im2ht.h" +#include "ht_cuda.h" + +at::Tensor +// std::vector +im2ht_forward( + const at::Tensor &feats, + const at::Tensor &voxels_src, + const at::Tensor &voxels_dst, + const at::Tensor &idxs_src, + const at::Tensor &idxs_dst, + const int h, + const int w, + const int d + ) +{ + if ( + feats.type().is_cuda() && \ + voxels_src.type().is_cuda() && voxels_dst.type().is_cuda() && \ + idxs_src.type().is_cuda() && idxs_dst.type().is_cuda() + ) + { + return ht_cuda_forward(feats, + voxels_src, voxels_dst, + idxs_src, idxs_dst, + h, w, d + ); + } + AT_ERROR("Not implemented on the CPU"); +} + + +at::Tensor +// std::vector +im2ht_backward(const at::Tensor &grad_vol, + const at::Tensor &voxels_src, + const at::Tensor &voxels_dst, + const at::Tensor &idxs_src, + const at::Tensor &idxs_dst, + const int h, + const int w, + const int d + ) +{ + if ( + grad_vol.type().is_cuda() && \ + voxels_src.type().is_cuda() && voxels_dst.type().is_cuda() && \ + idxs_src.type().is_cuda() && idxs_dst.type().is_cuda() + ) + { + return ht_cuda_backward(grad_vol, + voxels_src, voxels_dst, + idxs_src, idxs_dst, + h, w, d + ); + } + AT_ERROR("Not implemented on the CPU"); +} + + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("im2ht_forward", &im2ht_forward, "Forward pass of ht"); + m.def("im2ht_backward", &im2ht_backward, "Backward pass of ht"); +} diff --git a/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht.h b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht.h new file mode 100755 index 0000000..ccb263d --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht.h @@ -0,0 +1,28 @@ +#pragma once +#include + +// std::vector +at::Tensor +im2ht_forward( + const at::Tensor &feats_src_dst, + const at::Tensor &voxels_src, + const at::Tensor &voxels_dst, + const at::Tensor &idxs_src, + const at::Tensor &idxs_dst, + const int h, + const int w, + const int d + ); + +at::Tensor +// std::vector +im2ht_backward( + const at::Tensor &grad_output, + const at::Tensor &voxels_src, + const at::Tensor &voxels_dst, + const at::Tensor &idxs_src, + const at::Tensor &idxs_dst, + const int h, + const int w, + const int d + ); diff --git a/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht_cuda.cuh b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht_cuda.cuh new file mode 100755 index 0000000..72d9b4e --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/im2ht_cuda.cuh @@ -0,0 +1,206 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +// #include +#include +// #include + +#include "cuda_utils.h" + +template +__global__ void im2ht_cuda_forward_kernel(const int64_t n_threads, + const scalar_t *feats, + scalar_t *vol_ht, + const scalar_t *voxels_src, + const scalar_t *voxels_dst, + const scalar_t *idxs_src, + const scalar_t *idxs_dst, + const int b, const int l, const int c, + const int m, const int n, + const int h, const int w, const int d + ) +{ + CUDA_KERNEL_LOOP(index, n_threads) + { + // [b, l, c, m, n] + int b_tmp = index / l / c / m / n % b; // b + int l_tmp = index / c / m / n % l; // l + int c_tmp = index / m / n % c; // c + int m_tmp = index / n % m; // m + int n_tmp = index % n; // n + + int src_tmp = idxs_src[b_tmp * (l * m) + l_tmp * m + m_tmp]; // point idx in [b, l, m] + if (src_tmp>=0 ) + { + int dst_tmp = idxs_dst[b_tmp * (l * n) + src_tmp * n + n_tmp]; // point idx in [b, l, n] + if (dst_tmp>=0) + { + // Z-Y-X after voxelization + int y_src_tmp = voxels_src[b_tmp * (l * 2) + src_tmp * 2 + 0]; // in [b, l, 2] + int x_src_tmp = voxels_src[b_tmp * (l * 2) + src_tmp * 2 + 1]; // in [b, l, 2] + + int y_dst_tmp = voxels_dst[b_tmp * (l * 2) + dst_tmp * 2 + 0]; // in [b, l, 2] + int x_dst_tmp = voxels_dst[b_tmp * (l * 2) + dst_tmp * 2 + 1]; // in [b, l, 2] + + scalar_t feat_tmp = feats[b_tmp * (l * n * c) + src_tmp * (n * c) + n_tmp * (c) + c_tmp]; // in [b, l, n, c] + + int bin_x = x_dst_tmp - x_src_tmp; + int bin_y = y_dst_tmp - y_src_tmp; + bin_x += w/2; + bin_y += h/2; + + if (y_src_tmp >=0 && x_src_tmp >=0 && y_dst_tmp >= 0 && x_dst_tmp >=0) + { + if (bin_x >=0 && bin_x < w && bin_y >= 0 && bin_y < h) + { + // (0, 0) is at position (h//2, w//2) + // vol_ht: [b, l, h, w]; + int64_t offset_vol_tmp = b_tmp * (l * c * h * w) + l_tmp * (c * h * w) + c_tmp * (h * w) + bin_y * w + bin_x; + atomicAdd(vol_ht+offset_vol_tmp, feat_tmp); + } + } + } + } + } +} + +template +__global__ void im2ht_cuda_backward_kernel(const int n_threads, + scalar_t* grad_feats, + const scalar_t* grad_vol_ht, + const scalar_t *voxels_src, + const scalar_t *voxels_dst, + const scalar_t *idxs_src, + const scalar_t *idxs_dst, + const int b, const int l, const int c, + const int m, const int n, + const int h, const int w, const int d + ) +{ + // todo: coalesce + CUDA_KERNEL_LOOP(index, n_threads) + { + // [b, l, c, m, n] + int b_tmp = index / l / c / m / n % b; // b + int l_tmp = index / c / m / n % l; // l + int c_tmp = index / m / n % c; // c + int m_tmp = index / n % m; // m + int n_tmp = index % n; // n + + int src_tmp = idxs_src[b_tmp * (l * m) + l_tmp * m + m_tmp]; // point idx in [b, l, m] + if (src_tmp>=0 ) + { + int dst_tmp = idxs_dst[b_tmp * (l * n) + src_tmp * n + n_tmp]; // point idx in [b, l, n] + if (dst_tmp>=0) + { + // Z-Y-X after voxelization + int y_src_tmp = voxels_src[b_tmp * (l * 2) + src_tmp * 2 + 0]; // in [b, l, 2] + int x_src_tmp = voxels_src[b_tmp * (l * 2) + src_tmp * 2 + 1]; // in [b, l, 2] + + int y_dst_tmp = voxels_dst[b_tmp * (l * 2) + dst_tmp * 2 + 0]; // in [b, l, 2] + int x_dst_tmp = voxels_dst[b_tmp * (l * 2) + dst_tmp * 2 + 1]; // in [b, l, 2] + + int bin_x = x_dst_tmp - x_src_tmp; + int bin_y = y_dst_tmp - y_src_tmp; + bin_x += w/2; + bin_y += h/2; + + if (y_src_tmp >=0 && x_src_tmp >=0 && y_dst_tmp >= 0 && x_dst_tmp >=0) + { + if (bin_x >=0 && bin_x < w && bin_y >= 0 && bin_y < h) + { + // (0, 0) is at position (h//2, w//2) + // vol_ht: [b, l, c, h, w]; + // output gradients + int offset_vol_tmp = b_tmp * (l * c * h * w) + l_tmp * (c * h * w) + c_tmp * (h * w) + bin_y * w + bin_x; // [b, l, c, h, w] + scalar_t grad_vol_tmp = grad_vol_ht[offset_vol_tmp]; // + // input grads + int offset_feat_tmp = b_tmp * (l * n * c) + src_tmp * (n * c) + n_tmp * (c) + c_tmp; // in [b, l, n, c] + atomicAdd(grad_feats+offset_feat_tmp, grad_vol_tmp); + } + } + } + } + } +} + +// %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% // +template +void im2ht_cuda_forward(cudaStream_t stream, + const scalar_t* feats, + scalar_t* vol_ht, + const scalar_t* voxels_src, + const scalar_t* voxels_dst, + const scalar_t* idxs_src, + const scalar_t* idxs_dst, + const int b, const int l, const int c, + const int m, const int n, + const int h, const int w, const int d + ) +{ + const int num_kernels = b * l * c * m * n; + // printf("dimensions num_kernels=%16d, b=%06d, l=%06d, m=%06d, n=%06d, h=%06d, w= %06d, d= %06d \n", + // num_kernels, b, l, m, n, h, w, d); + im2ht_cuda_forward_kernel + <<>>(num_kernels, + feats, + vol_ht, + voxels_src, + voxels_dst, + idxs_src, + idxs_dst, + b, l, c, + m, n, + h, w, d + ); + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + { + printf("error in im2ht_gpu_kernel: %s\n", cudaGetErrorString(err)); + } +} + +template +void im2ht_cuda_backward(cudaStream_t stream, + scalar_t* grad_feats, + const scalar_t* grad_vol, + const scalar_t* voxels_src, + const scalar_t* voxels_dst, + const scalar_t* idxs_src, + const scalar_t* idxs_dst, + const int b, const int l, const int c, + const int m, const int n, + const int h, const int w, const int d + ) +{ + + const int num_kernels = b * l * c * m * n; + // ***********************************************************************// + im2ht_cuda_backward_kernel + <<>>(num_kernels, + grad_feats, + grad_vol, + voxels_src, + voxels_dst, + idxs_src, + idxs_dst, + b, l, c, + m, n, + h, w, d + ); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + { + printf("error in ht2im_gpu_kernel: %s\n", cudaGetErrorString(err)); + } + +} + diff --git a/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/setup.py b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/setup.py new file mode 100755 index 0000000..6c2521f --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/setup.py @@ -0,0 +1,37 @@ +from setuptools import setup +from torch.utils.cpp_extension import BuildExtension, CUDAExtension + +setup( + name='im2ht', + # ext_modules=[ + # CUDAExtension(name='im2ht', sources=['im2ht.cpp', 'ht_cuda.cu'], + # extra_compile_args={'cxx': ['-g'], 'nvcc': ['-arch=sm_60']}), + # ], + ext_modules=[ + CUDAExtension(name='im2ht', sources=['im2ht.cpp', 'ht_cuda.cu']), + ], + cmdclass={ + 'build_ext': BuildExtension + }) + + + +# def load_cpp_ext(ext_name): +# root_dir = os.path.join(os.path.split(__file__)[0]) +# src_dir = os.path.join(root_dir, "cpp") +# tar_dir = os.path.join(src_dir, "build", ext_name) +# os.makedirs(tar_dir, exist_ok=True) +# srcs = glob(f"{src_dir}/*.cu") + glob(f"{src_dir}/*.cpp") + +# with warnings.catch_warnings(): +# warnings.simplefilter("ignore") +# from torch.utils.cpp_extension import load + +# ext = load( +# name=ext_name, +# sources=srcs, +# extra_cflags=["-O3"], +# extra_cuda_cflags=[], +# build_directory=tar_dir, +# ) +# return ext \ No newline at end of file diff --git a/src/models/basic/voteflow_plugin/hough_transformation/ht_cuda.py b/src/models/basic/voteflow_plugin/hough_transformation/ht_cuda.py new file mode 100755 index 0000000..c989bed --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/ht_cuda.py @@ -0,0 +1,63 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.autograd import gradcheck +from torch.nn.modules.utils import _pair +import scipy.io as sio +import random +import numpy as np +import matplotlib.pyplot as plt +import math +import time +import os +from .im2ht import IM2HT + +class HT_CUDA(nn.Module): + """ + for each voxel in src: + find its m knn voxels in src: + for each voxel_src: + find its n knn voxels in dst; + calculate a transaltion which determines the bin in voting (bin_y, bin_x); + vote incrementally by the amount of correponding src and dst features; + """ + """ mapping from points to histgrams + input: feats_src_dst [b, l, n, c] + input: voxels_src [b, l, 2] + input: voxels_dst [b, l, 2] + input: idxs_src [b, l, m] + input: idxs_dst [b, l, n] + input: bins_x int + input: bins_y int + input: bins_z int + output: HT feaures [b, l, h, w], c*2 stands for the concatenation of features from src and dst. + """ + def __init__(self, h, w, d): + super(HT_CUDA, self).__init__() + self.h, self.w, self.d = h, w, d + assert h%2==0 + assert w%2==0 + self.ht = IM2HT() + + self.__repr__() + + def __repr__(self): + return self.__class__.__name__ + f'( h: {self.h:04d}, w: {self.w:04d}, d: {self.d:04d} )' + + def forward(self, feats_src_dst, voxels_src, voxels_dst, idxs_src, idxs_dst): + # print('ht_cuda forward: ', feats.shape, idxs_fps.shape, idxs_src.shape, idxs_dst.shape, bins_x.shape, bins_y.shape, bins_z.shape) + # print(f'( k:{self.k:04d}, m: {self.m:04d}, n: {self.n:04d}, h: {self.h:04d}, w: {self.w:04d}, d: {self.d:04d} )') + assert feats_src_dst.dim()==4 + assert voxels_src.shape == voxels_dst.shape + + b, l, n = idxs_dst.shape + _, _, m = idxs_src.shape + + datatype = feats_src_dst.dtype + voxels_src = voxels_src.to(datatype) + voxels_dst = voxels_dst.to(datatype) + idxs_src = idxs_src.to(datatype) + idxs_dst = idxs_dst.to(datatype) + vol = self.ht(feats_src_dst, voxels_src, voxels_dst, idxs_src, idxs_dst, self.h, self.w, self.d) + # print('vol: ', vol.shape, vol.min(), vol.max()) + return vol diff --git a/src/models/basic/voteflow_plugin/hough_transformation/im2ht.py b/src/models/basic/voteflow_plugin/hough_transformation/im2ht.py new file mode 100755 index 0000000..802f059 --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/im2ht.py @@ -0,0 +1,118 @@ +import os +import math +import numpy as np +import warnings +from glob import glob + +import torch +from torch import nn +from torch.autograd import Function +from torch.nn.modules.utils import _pair +from torch.autograd.function import once_differentiable + +def load_cpp_ext(ext_name): + root_dir = os.path.join(os.path.split(__file__)[0]) + src_dir = os.path.join(root_dir, "cpp_im2ht") + tar_dir = os.path.join(src_dir, "build", ext_name) + os.makedirs(tar_dir, exist_ok=True) + srcs = glob(f"{src_dir}/*.cu") + glob(f"{src_dir}/*.cpp") + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + from torch.utils.cpp_extension import load + ext = load( + name=ext_name, + sources=srcs, + extra_cflags=["-O3"], + extra_cuda_cflags=[], + build_directory=tar_dir, + ) + return ext + +# defer calling load_cpp_ext to make CUDA_VISIBLE_DEVICES happy +im2ht = None + + +class IM2HTFunction(Function): + @staticmethod + def forward(ctx, + feats_src_dst, + voxels_src, + voxels_dst, + idxs_src, + idxs_dst, + h, w, d + ): + + ctx.h = h + ctx.w = w + ctx.d = d + ctx.save_for_backward(voxels_src, voxels_dst, idxs_src, idxs_dst) + vol = im2ht.im2ht_forward( + feats_src_dst, + voxels_src, + voxels_dst, + idxs_src, + idxs_dst, + h, + w, + d + ) + return vol + + @staticmethod + @once_differentiable + def backward(ctx, grad_vol): + voxels_src, voxels_dst, idxs_src, idxs_dst = ctx.saved_tensors # it is a list of saved tensors! + + grad_input_src_dst = im2ht.im2ht_backward( + grad_vol, + voxels_src, + voxels_dst, + idxs_src, + idxs_dst, + ctx.h, + ctx.w, + ctx.d + ) + + return ( + grad_input_src_dst, + None, # voxels_src, + None, # voxels_dst, + None, # idxs_src, + None, # idxs_dst, + None, # h + None, # w + None # d + ) + + +class IM2HT(nn.Module): + def __init__(self, ): + super(IM2HT, self).__init__() + + global im2ht + im2ht = load_cpp_ext("im2ht") + + # self.extra_repr() + # self.__repr__() + + # def extra_repr(self): + # s = ('size={ht_size}') + # # return s.format(**self.__dict__) + # return print(s.format(**self.__dict__)) + + # def __repr__(self): + # return self.__class__.__name__ + f'( m: {self.m:04d}, n: {self.n:04d}, h: {self.h:04d}, w: {self.w:04d}, d: {self.d:04d})' + + def forward(self, feats_src_dst, voxels_src, voxels_dst, idxs_src, idxs_dst, h, w, d): + # print('IM2HT forward', feats.shape, idxs_fps.shape, idxs_src.shape, idxs_dst.shape, bins_x.shape, bins_y.shape, bins_z.shape) + return IM2HTFunction.apply( + feats_src_dst.contiguous(), + voxels_src.contiguous(), + voxels_dst.contiguous(), + idxs_src.contiguous(), + idxs_dst.contiguous(), + h, w, d + ) diff --git a/src/models/basic/voteflow_plugin/hough_transformation/test_backward.py b/src/models/basic/voteflow_plugin/hough_transformation/test_backward.py new file mode 100755 index 0000000..a018c3c --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/test_backward.py @@ -0,0 +1,73 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.autograd import gradcheck + +import random +import numpy as np +import math +import scipy.io as sio +import glob + +from models.ht.ht_cuda import HT_CUDA +from models.ht.im2ht import IM2HTFunction + +random.seed(0) +np.random.seed(0) +torch.manual_seed(0) + + +if torch.cuda.is_available(): + # device_name = "cuda" + # torch.backends.cudnn.benchmark = True + # torch.backends.cudnn.deterministic = True + # torch.cuda.manual_seed(0) + + device_name = "cuda" + torch.cuda.manual_seed(0) + # torch.backends.cudnn.enabled = False #greatly slow down the speed + torch.backends.cudnn.deterministic = True + torch.cuda.manual_seed(0) + torch.cuda.empty_cache() + print("Let's use", torch.cuda.device_count(), "GPU(s)!") +else: + device_name = "cpu" + print("CUDA is not available") +device = torch.device(device_name) +print('device: ', device) + + +b = 2 +l = 64 +c=3 +m = 4 +n = 8 +max_x=0.5 +max_y=0.5 +max_z=0.2 +voxel=0.1 +nx = math.ceil(max_x / voxel) * 2 # assume 120km/h, along both directions +ny = math.ceil(max_y / voxel) * 2 # assume 120km/h, along both directions +nz = math.ceil(max_z / voxel) * 2 + +vote = HT_CUDA(nx, ny, nz) +vote = vote.double().to(device) + +print('grad check***********') + +feats = torch.rand(size=(b, l, n, 3), requires_grad=True).double().to(device) +voxels_src = torch.randint(size=(b, l, 2), low=0, high=nx).double().to(device) +voxels_dst = torch.randint(size=(b, l, 2), low=0, high=nx).double().to(device) +idxs_src = torch.randint(0, l, size=(b, l, m)).double().to(device) +idxs_dst = torch.randint(0, l, size=(b, l, n)).double().to(device) +print('input shapes: ', feats.shape, voxels_src.shape, voxels_dst.shape, idxs_src.shape, idxs_dst.shape) +# grad_output = torch.randn(vol.shape).double().to(device) +# vol.backward(gradient=grad_output) +# grad_input = fs_window.grad +# print(grad_input.shape) + +# only able to test inputs with tiny sizes +res = gradcheck(vote, (feats, voxels_src, voxels_dst, idxs_src, idxs_dst), raise_exception=True) +# res = gradcheck(IM2HTFunction.apply, (feats, bins_x, bins_y, bins_z, flags, idxs, m, n, h, w, d), raise_exception=True) +# # res=gradcheck(myconv, input, eps=1e-3, atol=1e-3, rtol=1e-2, raise_exception=True) +print('grad check', res) diff --git a/src/models/basic/voteflow_plugin/hough_transformation/test_forward.py b/src/models/basic/voteflow_plugin/hough_transformation/test_forward.py new file mode 100644 index 0000000..3d5af04 --- /dev/null +++ b/src/models/basic/voteflow_plugin/hough_transformation/test_forward.py @@ -0,0 +1,247 @@ +import numpy as np +import random +import glob +import warnings +import parmap +import os +import itertools +import shutil +import argparse +import copy +import sys +import yaml +import json +import gc +import datetime, time +import open3d as o3d +import seaborn as sns +import multiprocessing +import matplotlib.pyplot as plt +import plotly +from tqdm import tqdm +import torch +import torch.nn as nn +import pytorch3d.ops as pytorch3d_ops +import plotly.graph_objs as go +from util_model import Decoder +from util_loss import warped_pc_loss +from utils_visualization import visualize_pcd +warnings.filterwarnings('ignore') +np.set_printoptions(suppress=True) +# add cuda implementation +from ht.ht_cuda import HT_CUDA + +# this loads a pair of frames and gt flow. Note 1). the source has been compensated by ego motion; and 2) ground has been removed. +def dataloader_minimal(data_path): + data = np.load(data_path) + pcl_0 = data['pc1'] + pcl_1 = data['pc2'] + valid_0 = data['pc1_flows_valid_idx'] + valid_1 = data['pc2_flows_valid_idx'] + flow_0_1 = data['gt_flow_0_1'] + class_0 = data['pc1_classes'] + class_1 = data['pc2_classes'] + + pcl_0 = pcl_0[valid_0] + pcl_1 = pcl_1[valid_1] + flow_0_1 = flow_0_1[valid_0] + class_0 = class_0[valid_0] + class_1 = class_1[valid_1] + # print('pc range: ', + # pcl_0[:, 0].min(), pcl_0[:, 0].max(), + # pcl_0[:, 1].min(), pcl_0[:, 1].max(), + # pcl_1[:, 0].min(), pcl_1[:, 0].max(), + # pcl_1[:, 1].min(), pcl_1[:, 1].max(), + # ) + # print('class 0: ', np.unique(class_0)) + # visualize_pcd( + # np.concatenate([pcl_0, pcl_1, pcl_0+flow_0_1], axis=0), + # np.concatenate([np.zeros(len(pcl_0))+1, np.zeros(len(pcl_1))+2, np.zeros(len(pcl_0))+0], axis=0), + # num_colors=3, + # title=f'visualize input: src-g, dst-b, src+flow-r: {data_path}' + # ) + data_dict = { + 'point_src': pcl_0, + 'point_dst': pcl_1, + 'scene_flow': flow_0_1, + 'data_path': data_path + } + return data_dict + +if __name__ == "__main__": + + # Initialization + random.seed(0) + np.random.seed(0) + # fix open3d seed. question mark there: sometimes o3d outputs different results. + o3d.utility.random.seed(0) # only in open3d>=0.16 + multiprocessing.set_start_method('forkserver') # or 'spawn' + + torch.manual_seed(0) + if torch.cuda.is_available(): + assert torch.cuda.device_count()==1 + device_name = f'cuda:{torch.cuda.current_device()}' + torch.backends.cudnn.deterministic = True + torch.cuda.manual_seed(0) + torch.cuda.empty_cache() + gc.collect() + torch.set_default_tensor_type("torch.cuda.FloatTensor") + print(f"Let's use {torch.cuda.device_count()}, {torch.cuda.get_device_name()} GPU(s)!") + else: + print("CUDA is not available") + device_name = 'cpu' + device = torch.device(device_name) + print(f'device: {device}') + + model = Decoder().to(device) + b = 4 + l = 40000 + m = 9 # neighbors in src + n =16 # neighbors in dst + vote = HT_CUDA(m, n, 3.3, 3.3, 0.1, 0.1).to(device) + + files = glob.glob(os.path.join('demo.npz')) + print('total files: ', len(files)) + for file in files: + data = dataloader_minimal(file) + points_src = data['point_src'] + points_dst = data['point_dst'] + + # random downsampling + points_src = points_src[np.random.choice(np.arange(0, len(points_src)), l, replace=False)] + points_dst = points_dst[np.random.choice(np.arange(0, len(points_src)), l, replace=False)] + + # coordinates + points_src = torch.from_numpy(points_src).float().to(device) + points_dst = torch.from_numpy(points_dst).float().to(device) + print('points: ', points_src.shape, points_dst.shape) + + # feats. this is a minimal exmaple. in pratice, they are learned from some backbone model; + feats_src = torch.randn((len(points_src), 64)).float().to(device) + feats_dst = torch.randn((len(points_dst), 64)).float().to(device) + feats_src /= torch.linalg.norm(feats_src, dim=-1, keepdims=True) + feats_dst /= torch.linalg.norm(feats_dst, dim=-1, keepdims=True) + print('feats: ', feats_src.shape, feats_dst.shape) + + # batch, think about how to batch point clouds of various length + points_src, points_dst = points_src[None].expand(b, l, 3).contiguous(), points_dst[None].expand(b, l, 3).contiguous() + feats_src, feats_dst = feats_src[None].expand(b, l, 64).contiguous(), feats_dst[None].expand(b, l, 64).contiguous() + + # only need to calculate once per batch + with torch.no_grad(): + + # process coordinates + # knn of src in dst + knn_dst = pytorch3d_ops.knn_points(points_src, points_dst, lengths1=None, lengths2=None, K=16, return_nn=True, return_sorted=False) + dists_dst, idxs_dst, pts_dst = knn_dst[0], knn_dst[1], knn_dst[2] + print('dst: ', dists_dst.shape, idxs_dst.shape, pts_dst.shape) + + # translations, between src and its knn in dst + ts_dst_src = pts_dst - points_src[:, :, None, :] + print('translations: ', ts_dst_src.shape) + + # define a local window for each point in src, using knn for simplicity + knn_src = pytorch3d_ops.knn_points(points_src, points_src, lengths1=None, lengths2=None, K=9, return_nn=True, return_sorted=False) + dists_src, idxs_src, pts_src = knn_src[0], knn_src[1], knn_src[2] + print('src: ', dists_src.shape, idxs_src.shape, pts_src.shape) + + # all translations within in a window, between src and its knn in dst + ts_window = pytorch3d_ops.knn_gather(ts_dst_src.view(b, l, -1), idxs_src, lengths=None) + ts_window = ts_window.view(idxs_src.shape[0], idxs_src.shape[1], idxs_src.shape[2], 16, 3) + print('ts_window: ', ts_window.shape) + + # quantization + n_intervals = 3.33 // 0.1 # both directions + intervals = torch.arange(-0.1*n_intervals, 0.1*n_intervals + 1e-8, 0.1) # [-3.3, -3.2, ..., 3.2, 3.3] + intervals_z = torch.arange(-0.1, 0.1 + 1e-8, 0.1) # [-0.1, 0, 0.1] + print('intervals xy: ', n_intervals, 0.1*n_intervals, intervals.shape, intervals, intervals_z) + + # tricky there, check: https://pytorch.org/docs/stable/generated/torch.bucketize.html + x_idxs = torch.bucketize(ts_window[:, :, :, :, 0], intervals, right=True) + y_idxs = torch.bucketize(ts_window[:, :, :, :, 1], intervals, right=True) + z_idxs = torch.bucketize(ts_window[:, :, :, :, 2], intervals_z, right=True) + x_idxs -= 1 + y_idxs -= 1 + z_idxs -= 1 + print('ts_idxs: ', x_idxs.shape, y_idxs.shape, z_idxs.shape) + + # # # sanity check: find out shared translations + # for i in range(0, x_idxs.shape[0]): + # for j in range(0, x_idxs.shape[1]): + # x_idxs_temp = x_idxs[i, j] + # y_idxs_temp = y_idxs[i, j] + # z_idxs_temp = z_idxs[i, j] + # for p in range(0, ts_window.shape[2]): + # for q in range(0, ts_window.shape[3]): + # print('t, idx: ', ts_window[i, j, p, q], x_idxs[i, j, p, q], y_idxs[i, j, p, q], z_idxs[i, j, p, q]) + # idxs_temp = torch.stack([x_idxs_temp, y_idxs_temp, z_idxs_temp], dim=2) + # print('idxs_temp: ', idxs_temp.shape) + # bin_unq, count_unq = torch.unique(idxs_temp.view(-1, 3), dim=0, return_counts=True) + # print('unique bins: ', bin_unq, count_unq) + + # process feats: gradient backpropagation during training + # feats: knn of src in dst + feats_dst_nn = pytorch3d_ops.knn_gather(feats_dst, idxs_dst, lengths=None) + print('feat dst nn: ', feats_dst_nn.shape) + + # feats similarity, between src and its knn in dst + feats = (feats_src[:, :, None, :] * feats_dst_nn).sum(-1) + print('feats: ', feats.shape) + + # all feats within in a window, between src and its knn in dst + feats_window = pytorch3d_ops.knn_gather(feats, idxs_src, lengths=None) + print('feats_window: ', feats_window.shape) + + # # simple test: no feature learning + feats_window.fill_(1.0) + + # # # Hough voting: pytorch implemenation vs custom cuda implementation + # 1. test pytorch implementation + start_time_pytorch = time.time() + b_idxs = torch.arange(0, b)[:, None, None, None].expand(b, l, 9, 16).flatten() + p_idxs = torch.arange(0, l)[None, :, None, None].expand(b, l, 9, 16).flatten() + x_idxs = x_idxs.flatten() + y_idxs = y_idxs.flatten() + z_idxs = z_idxs.flatten() + + # delete out-of-range values later + valid = torch.all( + torch.stack([x_idxs=0, y_idxs=0, z_idxs=0], dim=1), + dim=1).flatten() + print('valid: ', valid.shape, valid.sum()) + # build a 3d volume (x/y: -3.3 +3.3, z:-0.1 +0.1) + vol = torch.zeros((b, l, len(intervals)-1, len(intervals)-1, len(intervals_z)-1)).float() + + # incremental voting: a 3d representation per point + # https://discuss.pytorch.org/t/indexing-with-repeating-indices-numpy-add-at/10223/13 + # option 1: + # idxs = b_idxs * vol.shape[1:].numel() + p_idxs * vol.shape[2:].numel() + x_idxs * vol.shape[3:].numel() + y_idxs * vol.shape[4:].numel() + z_idxs # flatten to 1d indices + # print('flatten to 1d indices: ', vol.shape, vol.shape[1:].numel(), vol.shape[2:].numel(), vol.shape[3:].numel(), vol.shape[4:].numel()) + # feats_window = feats_window.flatten() + # print('idxs, feats_window: ', idxs.shape, feats_window.shape) + # vol.flatten().put_(idxs[valid], feats_window[valid], accumulate=True) + # vol.flatten().put_(idxs[valid], torch.ones(sum(valid)).float(), accumulate=True) + # vol = vol.view(b, n_pts, len(intervals)-1, len(intervals)-1, len(intervals_z)-1) + # option 2: + vol.index_put_([b_idxs[valid], p_idxs[valid], x_idxs[valid], y_idxs[valid], z_idxs[valid]], feats_window.flatten()[valid], accumulate=True) + print('vote: ', vol.shape, vol.max(), vol.min(), time.time()-start_time_pytorch) + + # 2. test custom cuda implementation + start_time_cuda = time.time() + vol2, _, _, _ = vote(feats_window, ts_window) + vol2 = vol2.permute(0, 1, 3, 4, 2) + print('vote2: ', vol2.shape, vol2.max(), vol2.min(), time.time()-start_time_cuda) + + # check result + print('vol==vol2: ', torch.equal(vol, vol2), (vol-vol2).abs().max()) + + # add conv + fc layers for final prediction + flow = model.forward(vol.view(b*l, len(intervals)-1, len(intervals)-1, len(intervals_z)-1).permute(0, 3, 1, 2).contiguous()) + print('flow: ', flow.shape) + + # calculate loss during training + loss = warped_pc_loss(points_src.view(-1, 3)+flow, points_dst.view(-1, 3)) + print('loss: ', loss.shape, loss) + + + diff --git a/src/models/basic/voteflow_plugin/utils.py b/src/models/basic/voteflow_plugin/utils.py new file mode 100644 index 0000000..ec1f843 --- /dev/null +++ b/src/models/basic/voteflow_plugin/utils.py @@ -0,0 +1,48 @@ +import math +import numpy as np +import torch +from torch_scatter import scatter_max +import pytorch3d.ops as pytorch3d_ops + +def calculate_unq_voxels(coords, image_dims): + unqs, idxs = torch.unique(coords[:, 1]*image_dims[0]+coords[:, 2], return_inverse=True, sorted=True) + unqs_voxel = torch.stack([unqs//image_dims[0], unqs%image_dims[0]], dim=1) + return unqs_voxel, idxs + +# https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/ops/utils.html +def batched_masked_gather(x: torch.Tensor, idxs: torch.Tensor, mask: torch.Tensor,fill_value=-1.0) -> torch.Tensor: + assert x.dim() == 3 # [b, m, c] + assert idxs.dim() == 3 # [b, n, k] + assert idxs.shape == mask.shape + b, m, c = x.shape + + idxs_masked = idxs.clone() + idxs_masked[~mask] = 0 + l, n, k = idxs.shape + y = pytorch3d_ops.knn_gather(x, idxs_masked) # [b, n, k, c] + y[~mask, :] = fill_value + + return y + +def pad_to_batch(idxs, l): + if idxs is None: + return idxs + if idxs.dim()==2: + assert idxs.shape[0]<=l + if idxs.shape[0] torch.Tensor: + x = self.conv(x) + x = F.relu(x) + return x + +class ConvBNBlock(nn.Module): + def __init__(self, in_num_channels: int, out_num_channels: int, + kernel_size=3, stride=1, padding=1): + super().__init__() + self.conv = nn.Conv2d(in_num_channels, out_num_channels, kernel_size, stride, padding, bias=True) + self.bn = nn.BatchNorm2d(out_num_channels) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.conv(x) + x = self.bn(x) + x = F.relu(x) + return x + +class VolConvBN(nn.Module): + def __init__(self, h, w, hidden_dim=16, dim_output=64): + super().__init__() + assert h%2==0 + assert w%2==0 + self.conv1 = ConvBNBlock(in_num_channels=1, out_num_channels=hidden_dim, stride=2) + self.conv2 = ConvBNBlock(in_num_channels=hidden_dim, out_num_channels=hidden_dim, stride=2) + self.linear = nn.Linear(math.ceil(h/4) * math.ceil(w/4) * hidden_dim, dim_output) + self.relu = nn.ReLU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + b, l, _, h, w = x.shape + x = self.conv1(x.view(b*l, 1, h, w)) + x = self.conv2(x) + # print(x.shape) + x = self.linear(x.view(b*l, -1)) + # print('after linear:', x.shape) + x = self.relu(x) + return x.view(b, l, -1) + + +class VoteFlowLinearDecoder(nn.Module): + def __init__(self, dim_input=16, layer_size=1, filter_size=128): + super().__init__() + # self.linear = nn.Linear(m*dim_input, dim_output) + offset_encoder_channels = 128 + self.offset_encoder = nn.Linear(3, offset_encoder_channels) + filter_size=filter_size + + decoder = nn.ModuleList() + decoder.append(torch.nn.Linear(dim_input + offset_encoder_channels, filter_size)) + decoder.append(torch.nn.GELU()) + for _ in range(layer_size-1): + decoder.append(torch.nn.Linear(filter_size, filter_size)) + decoder.append(torch.nn.ReLU()) + decoder.append(torch.nn.Linear(filter_size, 3)) + + self.decoder = nn.Sequential(*decoder) + print(self.decoder) + + def forward(self, x: torch.Tensor, pts_offsets: torch.Tensor) -> torch.Tensor: + b, l, _ = x.shape + pts_offsets_feats = self.offset_encoder(pts_offsets) + x = torch.cat([x, pts_offsets_feats], dim=-1) + # print('decoder conv: ', x.shape) + x = self.decoder(x) + + return x + + +class ConvWithNorms(nn.Module): + + def __init__(self, in_num_channels: int, out_num_channels: int, + kernel_size: int, stride: int, padding: int): + super().__init__() + self.conv = nn.Conv2d(in_num_channels, out_num_channels, kernel_size, + stride, padding) + self.batchnorm = nn.BatchNorm2d(out_num_channels) + self.nonlinearity = nn.GELU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + conv_res = self.conv(x) + if conv_res.shape[2] == 1 and conv_res.shape[3] == 1: + # This is a hack to get around the fact that batchnorm doesn't support + # 1x1 convolutions + batchnorm_res = conv_res + else: + batchnorm_res = self.batchnorm(conv_res) + return self.nonlinearity(batchnorm_res) + +class BilinearDecoder(nn.Module): + + def __init__(self, scale_factor: int): + super().__init__() + self.scale_factor = scale_factor + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return nn.functional.interpolate(x, + scale_factor=self.scale_factor, + mode="bilinear", + align_corners=False) + +class UpsampleSkip(nn.Module): + + def __init__(self, skip_channels: int, latent_channels: int, + out_channels: int): + super().__init__() + self.u1_u2 = nn.Sequential( + nn.Conv2d(skip_channels, latent_channels, 1, 1, 0), + BilinearDecoder(2)) + self.u3 = nn.Conv2d(latent_channels, latent_channels, 1, 1, 0) + self.u4_u5 = nn.Sequential( + nn.Conv2d(2 * latent_channels, out_channels, 3, 1, 1), + nn.Conv2d(out_channels, out_channels, 3, 1, 1)) + + def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + u2_res = self.u1_u2(a) + u3_res = self.u3(b) + u5_res = self.u4_u5(torch.cat([u2_res, u3_res], dim=1)) + return u5_res + +class FastFlow3DUNet(nn.Module): + """ + Standard UNet with a few modifications: + - Uses Bilinear interpolation instead of transposed convolutions + """ + + def __init__(self, in_num_channels, out_num_channels) -> None: + super().__init__() + + self.encoder_step_1 = nn.Sequential(ConvWithNorms(in_num_channels, 64, 3, 2, 1), + ConvWithNorms(64, 64, 3, 1, 1), + ConvWithNorms(64, 64, 3, 1, 1), + ConvWithNorms(64, 64, 3, 1, 1)) + self.encoder_step_2 = nn.Sequential(ConvWithNorms(64, 128, 3, 2, 1), + ConvWithNorms(128, 128, 3, 1, 1), + ConvWithNorms(128, 128, 3, 1, 1), + ConvWithNorms(128, 128, 3, 1, 1), + ConvWithNorms(128, 128, 3, 1, 1), + ConvWithNorms(128, 128, 3, 1, 1)) + self.encoder_step_3 = nn.Sequential(ConvWithNorms(128, 256, 3, 2, 1), + ConvWithNorms(256, 256, 3, 1, 1), + ConvWithNorms(256, 256, 3, 1, 1), + ConvWithNorms(256, 256, 3, 1, 1), + ConvWithNorms(256, 256, 3, 1, 1), + ConvWithNorms(256, 256, 3, 1, 1)) + self.decoder_step1 = UpsampleSkip(512, 256, 256) + self.decoder_step2 = UpsampleSkip(256, 128, 128) + self.decoder_step3 = UpsampleSkip(128, 64, 64) + self.decoder_step4 = nn.Conv2d(64, out_num_channels, 3, 1, 1) + + def forward(self, pc0_B: torch.Tensor, + pc1_B: torch.Tensor) -> torch.Tensor: + + expected_channels = 32 + assert pc0_B.shape[ + 1] == expected_channels, f"Expected {expected_channels} channels, got {pc0_B.shape[1]}" + assert pc1_B.shape[ + 1] == expected_channels, f"Expected {expected_channels} channels, got {pc1_B.shape[1]}" + + pc0_F = self.encoder_step_1(pc0_B) + pc0_L = self.encoder_step_2(pc0_F) + pc0_R = self.encoder_step_3(pc0_L) + + pc1_F = self.encoder_step_1(pc1_B) + pc1_L = self.encoder_step_2(pc1_F) + pc1_R = self.encoder_step_3(pc1_L) + + Rstar = torch.cat([pc0_R, pc1_R], + dim=1) # torch.Size([1, 512, 64, 64]) + Lstar = torch.cat([pc0_L, pc1_L], + dim=1) # torch.Size([1, 256, 128, 128]) + Fstar = torch.cat([pc0_F, pc1_F], + dim=1) # torch.Size([1, 128, 256, 256]) + Bstar = torch.cat([pc0_B, pc1_B], + dim=1) # torch.Size([1, 64, 512, 512]) + + S = self.decoder_step1(Rstar, Lstar) + T = self.decoder_step2(S, Fstar) + U = self.decoder_step3(T, Bstar) + V = self.decoder_step4(U) + + return V \ No newline at end of file diff --git a/src/models/voteflow.py b/src/models/voteflow.py new file mode 100644 index 0000000..670a3c6 --- /dev/null +++ b/src/models/voteflow.py @@ -0,0 +1,376 @@ +import math +import numpy as np + +import dztimer +import torch +import torch.nn as nn +import pytorch3d.ops as pytorch3d_ops + +from .basic.encoder import DynamicEmbedder +from .basic import cal_pose0to1 + +from .basic.voteflow_plugin.hough_transformation import HT_CUDA +from .basic.voteflow_plugin.voteflow_module import VolConvBN, Decoder, FastFlow3DUNet +from .basic.voteflow_plugin.utils import calculate_unq_voxels, batched_masked_gather, pad_to_batch + +import warnings +warnings.filterwarnings('ignore') +np.set_printoptions(suppress=True) + +class VoteFlow(nn.Module): + def __init__(self, + nframes=1, + m=8, + n=64, + using_voting=True, + input_channels=32, + output_channels=64, + point_cloud_range = [-51.2, -51.2, -3, 51.2, 51.2, 3], + voxel_size=(0.2, 0.2, 6), + grid_feature_size = [512, 512], + decoder_layers=1, + use_ball_query=False, + use_separate_feats_voting=False, + vol_conv_hidden_dim=16, + **kwargs): + super().__init__() + + assert len(point_cloud_range)==6 + assert len(voxel_size)==3 + + # a hack there, may and may not have any impact depends on the setting. + assert voxel_size[0]==voxel_size[1] + self.using_voting = using_voting + print('Using voting:', self.using_voting) + + # General architecture: + pseudo_image_dims = grid_feature_size[:2] #int((point_cloud_range[3]-point_cloud_range[0])/voxel_size[0]), int((point_cloud_range[4]-point_cloud_range[1])/voxel_size[1])) # ignore z dimension + + self.point_cloud_range = point_cloud_range + self.voxel_size = voxel_size + self.pseudo_image_dims = pseudo_image_dims + #self.backbone = Backbone(input_channels, output_channels) + self.backbone = FastFlow3DUNet(input_channels, output_channels) ## output_channel 64 + + if self.using_voting: + self.decoder = Decoder(dim_input= output_channels * 2 + input_channels * 2, layer_size=decoder_layers) + print('decoder:', self.decoder) + else: + self.decoder = Decoder(dim_input= output_channels + input_channels * 2, layer_size=decoder_layers) + print('decoder:', self.decoder) + + + self.embedder = DynamicEmbedder(voxel_size=voxel_size, + pseudo_image_dims=pseudo_image_dims, + point_cloud_range=point_cloud_range, + feat_channels=input_channels) + + # Voting Module settings + if self.using_voting: + ## how many bins inside a voxel after quantization + ## assume 72km/h (20m/s), along x/y; 0.1m along z + e = 1e-8 + radius = 2 + nx = math.ceil((radius+e)*nframes / voxel_size[0]) + ny = math.ceil((radius+e)*nframes / voxel_size[1]) + nz = math.ceil((0.1+e) / voxel_size[2]) # +/-0.1 + + self.nx = nx*2 # +/-x + self.ny = ny*2 # +/-y + self.nz = nz*2 # +/-z + print('n x/y/z: ', self.nx, self.ny, self.nz) + + self.nframes = nframes + self.m = m # m knn within src, for each src point + self.using_ball_query = use_ball_query + self.n = n # n knn between src and dst, for each src voxel + self.radius_dst = max(nx, ny) # define a search window for a src voxel in dst voxels for calculating translations + + if self.using_ball_query: + self.radius_src = math.ceil(max((radius+e)/voxel_size[0], (radius+e)/voxel_size[1])) # define a search window (in meters) within src voxels, aka the rigid motion window + print(f'using ball query to search window radius in source pc: {self.radius_src}, m={self.m};') + else: + print(f'using knn to search window radius in source pc, m={self.m};') + print(f'using ball query to search in target pc, n={self.n}, search window radius: {self.radius_dst}.') + self.use_separate_feats = use_separate_feats_voting + print(f'using separate features for voting: {self.use_separate_feats}') + + self.vote = HT_CUDA(self.ny, self.nx, self.nz) + + self.volconv = VolConvBN(self.ny, self.nx, hidden_dim=vol_conv_hidden_dim, dim_output=output_channels) + + self.timer = dztimer.Timing() + self.timer.start("Total") + + def process_points_per_pair(self, voxel_info_src, voxel_info_dst): + valid_point_idxs_src = voxel_info_src['point_idxes'] # [N_valid_pts] + valid_point_idxs_dst = voxel_info_dst['point_idxes'] + valid_voxel_coords_src = voxel_info_src['voxel_coords'] #[N_valid_pts, 3] + valid_voxel_coords_dst = voxel_info_dst['voxel_coords'] + + unq_voxel_coords_src, point_voxel_idxs_src = calculate_unq_voxels(valid_voxel_coords_src, self.pseudo_image_dims) + unq_voxel_coords_dst, point_voxel_idxs_dst = calculate_unq_voxels(valid_voxel_coords_dst, self.pseudo_image_dims) + + # unq_voxel_coords_src # [N_valid_voxels, 2] + # point_voxel_idxs_src # [N_valid_pts], the index of the unique voxel for each point + if self.using_voting: + dists_dst, knn_idxs_dst, _ = pytorch3d_ops.ball_query(unq_voxel_coords_src[None].float(), unq_voxel_coords_dst[None].float(), lengths1=None, lengths2=None, K=self.n, return_nn=False, radius=self.radius_dst) + if self.using_ball_query: + dists_src, knn_idxs_src, _ = pytorch3d_ops.ball_query(unq_voxel_coords_src[None].float(), unq_voxel_coords_src[None].float(), lengths1=None, lengths2=None, K=self.m, return_nn=False, radius=self.radius_src) + else: + dists_src, knn_idxs_src, _ = pytorch3d_ops.knn_points(unq_voxel_coords_src[None].float(), unq_voxel_coords_src[None].float(), lengths1=None, lengths2=None, K=self.m, return_nn=False, return_sorted=False) + else: + knn_idxs_dst=[None] + knn_idxs_src=[None] + + return valid_point_idxs_src, valid_point_idxs_dst, point_voxel_idxs_src, point_voxel_idxs_dst, unq_voxel_coords_src, unq_voxel_coords_dst, knn_idxs_src[0], knn_idxs_dst[0] + + def preprocessing(self, points_src, points_dst, voxel_info_list_src, voxel_info_list_dst): + point_masks_src = [] + point_masks_dst = [] + point_voxel_idxs_src = [] + point_voxel_idxs_dst = [] + point_offsets_src = [] + + unq_voxels_src = [] + unq_voxels_dst = [] + knn_idxs_src = [] + knn_idxs_dst = [] + + l_points = 0 # for padding + l_voxels = 0 # for padding + for points_src_, points_dst_, voxel_info_src, voxel_info_dst in zip(points_src, points_dst, voxel_info_list_src, voxel_info_list_dst): + point_idxs_src_, point_idxs_dst_, \ + point_voxel_idxs_src_, point_voxel_idxs_dst_, \ + unq_voxels_src_, unq_voxels_dst_, \ + knn_idxs_src_, knn_idxs_dst_ = self.process_points_per_pair(voxel_info_src, voxel_info_dst) + # print('process_points_per_pair: ', point_idxs_src_.shape, point_voxel_idxs_src_.shape, unq_voxels_src_.shape, knn_idxs_dst_.shape) + + # generate a mask for easy batching and sequential processing + point_masks_src_ = torch.zeros((len(points_src_)), device=points_src_.device, dtype=points_src_.dtype) + point_masks_src_[point_idxs_src_] = 1 + point_masks_dst_ = torch.zeros((len(points_dst_)), device=points_dst_.device, dtype=points_dst_.dtype) + point_masks_dst_[point_idxs_dst_] = 1 + point_masks_src.append(point_masks_src_) + point_masks_dst.append(point_masks_dst_) + + point_voxel_idxs_src.append(point_voxel_idxs_src_) + point_voxel_idxs_dst.append(point_voxel_idxs_dst_) + point_offsets_src.append(voxel_info_src['point_offsets']) + unq_voxels_src.append(unq_voxels_src_) + unq_voxels_dst.append(unq_voxels_dst_) + knn_idxs_src.append(knn_idxs_src_) + knn_idxs_dst.append(knn_idxs_dst_) + + l_voxels = max(l_voxels, max(unq_voxels_src_.shape[0], unq_voxels_dst_.shape[0])) + l_points = max(l_points, max(point_voxel_idxs_src_.shape[0], point_voxel_idxs_dst_.shape[0])) + # print('l', l_points, l_voxels) + + # padding + for i, (point_voxel_idxs_src_, point_voxel_idxs_dst_, point_offsets_src_, unq_voxels_src_, unq_voxels_dst_, knn_idxs_src_, knn_idxs_dst_) \ + in enumerate( zip(point_voxel_idxs_src, point_voxel_idxs_dst, point_offsets_src, unq_voxels_src, unq_voxels_dst, knn_idxs_src, knn_idxs_dst) ): + unq_voxels_src_= pad_to_batch(unq_voxels_src_, l_voxels) + unq_voxels_dst_ = pad_to_batch(unq_voxels_dst_, l_voxels) + point_voxel_idxs_src_= pad_to_batch(point_voxel_idxs_src_, l_points) + point_voxel_idxs_dst_= pad_to_batch(point_voxel_idxs_dst_, l_points) + point_offsets_src_ = pad_to_batch(point_offsets_src_, l_points) + + knn_idxs_src_ = pad_to_batch(knn_idxs_src_, l_voxels) + knn_idxs_dst_= pad_to_batch(knn_idxs_dst_, l_voxels) + + unq_voxels_src[i] = unq_voxels_src_ + unq_voxels_dst[i] = unq_voxels_dst_ + point_voxel_idxs_src[i] = point_voxel_idxs_src_ + point_voxel_idxs_dst[i] = point_voxel_idxs_dst_ + point_offsets_src[i] = point_offsets_src_ + + knn_idxs_src[i] = knn_idxs_src_ + knn_idxs_dst[i] = knn_idxs_dst_ + + unq_voxels_src = torch.stack(unq_voxels_src, dim=0) + unq_voxels_dst = torch.stack(unq_voxels_dst, dim=0) + point_voxel_idxs_src = torch.stack(point_voxel_idxs_src, dim=0) + point_voxel_idxs_dst = torch.stack(point_voxel_idxs_dst, dim=0) + point_offsets_src = torch.stack(point_offsets_src, dim=0) + point_masks_src = torch.stack(point_masks_src, dim=0) + point_masks_dst = torch.stack(point_masks_dst, dim=0) + + if self.using_voting: + knn_idxs_src = torch.stack(knn_idxs_src, dim=0) + knn_idxs_dst = torch.stack(knn_idxs_dst, dim=0) + else: + knn_idxs_src = None + knn_idxs_dst = None + + return point_masks_src, point_masks_dst, point_voxel_idxs_src, point_voxel_idxs_dst, point_offsets_src, unq_voxels_src, unq_voxels_dst, knn_idxs_src, knn_idxs_dst + + def extract_voxel_from_image(self, image, voxels): + # image: [b, c, h, w]; voxels: [b, num, 2] + idxs = voxels[:, :, 0] * self.pseudo_image_dims[0] + voxels[:, :, 1] + + mask = voxels.min(-1)[0]>=0 + + b, c, h, w = image.shape + + feats_per_voxel = batched_masked_gather(image.view(b, c, h*w).permute(0, 2, 1), idxs[:, :, None].long(), mask[:, :, None], fill_value=0) + + # print('point per voxel: ', feats_per_voxel.shape) + return feats_per_voxel[:, :, 0, :] # [b, num , c] + + def extract_point_from_voxel(self, voxels, idxs): + # voxels: [b, l, c]; idxs: [b, num] + feats_per_point = batched_masked_gather(voxels, idxs[:,:,None].long(), idxs[:, :, None]>=0, fill_value=0) + # print('point per point: ', feats_per_point.shape) + return feats_per_point[:, :, 0, :]# [b, num, c] + + def extract_point_from_image(self, image, voxels, point_idxs): + feats_per_voxel = self.extract_voxel_from_image(image, voxels) + feats_per_point = self.extract_point_from_voxel(feats_per_voxel, point_idxs) + # print('point per point merged: ', feats_per_point.shape) + return feats_per_point + + # adapted from zeroflow + def concat_feats(self, feats_vote, feats_before, feats_after): + feats = torch.cat([feats_vote, feats_before, feats_after], dim=-1) + return feats + + def _model_forward(self, points_src, points_dst): + # assert points_src.shape==points_dst.shape + + self.timer[1][0].start("Voxelization") + pseudoimages_src, voxel_infos_lst_src = self.embedder(points_src) + pseudoimages_dst, voxel_infos_lst_dst = self.embedder(points_dst) + self.timer[1][0].stop() + + self.timer[1][1].start("Preprocessing") + with torch.no_grad(): + point_masks_src, point_masks_dst, \ + point_voxel_idxs_src, point_voxel_idxs_dst, \ + point_offsets_src, \ + voxels_src, voxels_dst, \ + knn_idxs_src, knn_idxs_dst = self.preprocessing(points_src, points_dst, voxel_infos_lst_src, voxel_infos_lst_dst) + self.timer[1][1].stop() + + self.timer[1][2].start("Feature extraction") + pseudoimages_grid = self.backbone(pseudoimages_src, pseudoimages_dst) + + feats_point_src_init = self.extract_point_from_image(torch.cat([pseudoimages_src, pseudoimages_dst], dim=1), voxels_src, point_voxel_idxs_src) + feats_point_src_grid = self.extract_point_from_image(pseudoimages_grid, voxels_src, point_voxel_idxs_src) + + self.timer[1][2].stop() + + if self.using_voting: + self.timer[1][3].start("Gathering") + if self.use_separate_feats: + # print('using separate features for voting') + feats_voxel_src = self.extract_voxel_from_image(pseudoimages_src, voxels_src) # [B, N_valid_voxels, C] + feats_voxel_dst = self.extract_voxel_from_image(pseudoimages_dst, voxels_dst) # [B, N_valid_voxels, C] + else: + feats_voxel_src = self.extract_voxel_from_image(pseudoimages_grid, voxels_src) # [B, N_valid_voxels, C] + feats_voxel_dst = self.extract_voxel_from_image(pseudoimages_grid, voxels_dst) # [B, N_valid_voxels, C] + feats_voxel_dst_inflate = batched_masked_gather(feats_voxel_dst, knn_idxs_dst.long(), knn_idxs_dst>=0, fill_value=0) + + corr_src_dst = torch.nn.functional.cosine_similarity(feats_voxel_src[:, :, None, :], feats_voxel_dst_inflate, dim=-1) + self.timer[1][3].stop() + + self.timer[1][4].start("Voting") + + voting_vols= self.vote(corr_src_dst[:, :, :, None], voxels_src, voxels_dst, knn_idxs_src, knn_idxs_dst) # [b, l, c, ny, nx] + + vols = self.volconv(voting_vols) + feats_point_vol = self.extract_point_from_voxel(vols, point_voxel_idxs_src) + + self.timer[1][4].stop() + + self.timer[1][5].start("Decoding") + if self.decoder=='simple_decoder': + flows = self.decoder(feats_point_vol) + else: + feats_cat = torch.cat([feats_point_vol, feats_point_src_init, feats_point_src_grid], -1) + + flows = self.decoder(feats_cat, point_offsets_src) + self.timer[1][5].stop() + else: + feats_voxel_src = None + feats_voxel_dst = None + corr_src_dst = None + voting_vols = None + + self.timer[1][5].start("Decoding") + feats_cat = torch.cat([feats_point_src_init, feats_point_src_grid], -1) + flows = self.decoder(feats_cat, point_offsets_src) + self.timer[1][5].stop() + + pc0_points_lst = [e["points"] for e in voxel_infos_lst_src] + pc1_points_lst = [e["points"] for e in voxel_infos_lst_dst] + + pc0_valid_point_idxes = [e["point_idxes"] for e in voxel_infos_lst_src] + pc1_valid_point_idxes = [e["point_idxes"] for e in voxel_infos_lst_dst] + + flows_reshape = [] + for (flow, valid_pts) in zip(flows, pc0_valid_point_idxes): + flow = flow[:valid_pts.shape[0], :] + flows_reshape.append(flow) + + model_res = { + "pseudoimages_src": pseudoimages_src, + "pseudoimages_dst": pseudoimages_dst, + "pseudoimages_grid": pseudoimages_grid, + "feats_voxel_src": feats_voxel_src, + "feats_voxel_dst": feats_voxel_dst, + "corr": corr_src_dst, + "voxels_src": voxels_src, + "voting_vol": voting_vols, + "points_src_offset": point_offsets_src, + "points_src_voxel_idx": point_voxel_idxs_src, + "flow": flows_reshape, + "pc0_points_lst": pc0_points_lst, + "pc1_points_lst": pc1_points_lst, + "pc0_valid_point_idxes": pc0_valid_point_idxes, + "pc1_valid_point_idxes": pc1_valid_point_idxes, + } + + return model_res + + + def forward(self, batch): + """ + input: using the batch from dataloader, which is a dict + Detail: [pc0, pc1, pose0, pose1] + output: the predicted flow, pose_flow, and the valid point index of pc0 + """ + self.timer[0].start("Data Preprocess") + batch_sizes = len(batch["pose0"]) + + pose_flows = [] + transform_pc0s = [] + for batch_id in range(batch_sizes): + selected_pc0 = batch["pc0"][batch_id] + self.timer[0][0].start("pose") + with torch.no_grad(): + if 'ego_motion' in batch: + pose_0to1 = batch['ego_motion'][batch_id] + else: + pose_0to1 = cal_pose0to1(batch["pose0"][batch_id], batch["pose1"][batch_id]) + self.timer[0][0].stop() + + self.timer[0][1].start("transform") + # transform selected_pc0 to pc1 + transform_pc0 = selected_pc0 @ pose_0to1[:3, :3].T + pose_0to1[:3, 3] + self.timer[0][1].stop() + pose_flows.append(transform_pc0 - selected_pc0) + transform_pc0s.append(transform_pc0) + + pc0s = torch.stack(transform_pc0s, dim=0) + pc1s = batch["pc1"] + self.timer[0].stop() + + self.timer[1].start("Model Forward") + model_res = self._model_forward(pc0s, pc1s) + self.timer[1].stop() + + ret_dict = model_res + ret_dict["pose_flow"] = pose_flows + + return ret_dict + \ No newline at end of file From 2696dff13b59f2ae90ad241111a7a911603c17d9 Mon Sep 17 00:00:00 2001 From: Shiming Date: Wed, 28 May 2025 08:54:59 +0000 Subject: [PATCH 03/10] a workable version and lack the learning rate scheduler --- README.md | 16 +++++++-- conf/model/voteflow.yaml | 18 ++++++++++ src/models/__init__.py | 4 ++- src/models/basic/voteflow_plugin/__init__.py | 2 +- src/models/voteflow.py | 37 ++++++++++---------- 5 files changed, 54 insertions(+), 23 deletions(-) create mode 100644 conf/model/voteflow.yaml diff --git a/README.md b/README.md index 08d6bee..45a55f2 100644 --- a/README.md +++ b/README.md @@ -132,11 +132,23 @@ mamba activate opensf ### VoteFLow Extra pakcges needed for VoteFlow, [pytorch3d](https://pytorch3d.org/) (prefer 0.7.7) and [torch-scatter](https://github.com/rusty1s/pytorch_scatter?tab=readme-ov-file) (prefer 2.1.2): +Install Pytorch3d +```bash +conda install -c iopath iopath +conda install -c bottler nvidiacub +conda install pytorch3d -c pytorch3d +``` +Install torch-scatter ```bash -pip install pytorch3d pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.0+cu117.html ``` - +```bash +python train.py model=voteflow lr=2e-4 epochs=12 batch_size=16 model.target.m=8 model.target.n=128 loss_fn=seflowLoss "add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}" +``` +Pretrained weight can be downloaded through +``` +PLACEHOLDER +``` ### Flow4D Train Flow4D with the leaderboard submit config. [Runtime: Around 18 hours in 4x RTX 3090 GPUs.] diff --git a/conf/model/voteflow.yaml b/conf/model/voteflow.yaml new file mode 100644 index 0000000..9ce5d3a --- /dev/null +++ b/conf/model/voteflow.yaml @@ -0,0 +1,18 @@ +name: voteflow + +target: + _target_: src.models.VoteFlow + using_voting: True + nframes: 1 + m: 8 + n: 128 + input_channels: 32 + output_channels: 64 + point_cloud_range: ${point_cloud_range} + voxel_size: ${voxel_size} + grid_feature_size: [512, 512] + decoder_layers: 1 + use_ball_query: False + vol_conv_hidden_dim: 16 + +val_monitor: val/Dynamic/Mean \ No newline at end of file diff --git a/src/models/__init__.py b/src/models/__init__.py index 503a046..3fdd7ec 100644 --- a/src/models/__init__.py +++ b/src/models/__init__.py @@ -12,6 +12,7 @@ from .deflow import DeFlow from .fastflow3d import FastFlow3D +from .voteflow import VoteFlow # following need install extra package: # * pip install spconv-cu117 @@ -27,4 +28,5 @@ from .ssf import SSF except ImportError as e: print("\033[93m--- WARNING [model]: Model with torch scatter is not imported, as it requires some lib which is not installed.") - print(f"Detail error message\033[0m: {e}. Just ignore this warning if code runs without these models.") \ No newline at end of file + print(f"Detail error message\033[0m: {e}. Just ignore this warning if code runs without these models.") + diff --git a/src/models/basic/voteflow_plugin/__init__.py b/src/models/basic/voteflow_plugin/__init__.py index 02e9652..a04b2bc 100644 --- a/src/models/basic/voteflow_plugin/__init__.py +++ b/src/models/basic/voteflow_plugin/__init__.py @@ -1,3 +1,3 @@ from .hough_transformation import HT_CUDA from .utils import calculate_unq_voxels, batched_masked_gather, pad_to_batch -from .voteflow_module import VolConvBN, Decoder, FastFlow3DUNet \ No newline at end of file +from .voteflow_module import VolConvBN, VoteFlowLinearDecoder, FastFlow3DUNet \ No newline at end of file diff --git a/src/models/voteflow.py b/src/models/voteflow.py index 670a3c6..05c7b29 100644 --- a/src/models/voteflow.py +++ b/src/models/voteflow.py @@ -1,3 +1,11 @@ +""" +This file is directly copied from: +https://github.com/tudelft-iv/VoteFlow + +with slightly modification to have unified format with all benchmark. +""" + + import math import numpy as np @@ -10,7 +18,7 @@ from .basic import cal_pose0to1 from .basic.voteflow_plugin.hough_transformation import HT_CUDA -from .basic.voteflow_plugin.voteflow_module import VolConvBN, Decoder, FastFlow3DUNet +from .basic.voteflow_plugin.voteflow_module import VolConvBN, VoteFlowLinearDecoder, FastFlow3DUNet from .basic.voteflow_plugin.utils import calculate_unq_voxels, batched_masked_gather, pad_to_batch import warnings @@ -30,7 +38,6 @@ def __init__(self, grid_feature_size = [512, 512], decoder_layers=1, use_ball_query=False, - use_separate_feats_voting=False, vol_conv_hidden_dim=16, **kwargs): super().__init__() @@ -53,10 +60,10 @@ def __init__(self, self.backbone = FastFlow3DUNet(input_channels, output_channels) ## output_channel 64 if self.using_voting: - self.decoder = Decoder(dim_input= output_channels * 2 + input_channels * 2, layer_size=decoder_layers) + self.decoder = VoteFlowLinearDecoder(dim_input= output_channels * 2 + input_channels * 2, layer_size=decoder_layers) print('decoder:', self.decoder) else: - self.decoder = Decoder(dim_input= output_channels + input_channels * 2, layer_size=decoder_layers) + self.decoder = VoteFlowLinearDecoder(dim_input= output_channels + input_channels * 2, layer_size=decoder_layers) print('decoder:', self.decoder) @@ -92,8 +99,6 @@ def __init__(self, else: print(f'using knn to search window radius in source pc, m={self.m};') print(f'using ball query to search in target pc, n={self.n}, search window radius: {self.radius_dst}.') - self.use_separate_feats = use_separate_feats_voting - print(f'using separate features for voting: {self.use_separate_feats}') self.vote = HT_CUDA(self.ny, self.nx, self.nz) @@ -261,13 +266,10 @@ def _model_forward(self, points_src, points_dst): if self.using_voting: self.timer[1][3].start("Gathering") - if self.use_separate_feats: - # print('using separate features for voting') - feats_voxel_src = self.extract_voxel_from_image(pseudoimages_src, voxels_src) # [B, N_valid_voxels, C] - feats_voxel_dst = self.extract_voxel_from_image(pseudoimages_dst, voxels_dst) # [B, N_valid_voxels, C] - else: - feats_voxel_src = self.extract_voxel_from_image(pseudoimages_grid, voxels_src) # [B, N_valid_voxels, C] - feats_voxel_dst = self.extract_voxel_from_image(pseudoimages_grid, voxels_dst) # [B, N_valid_voxels, C] + + feats_voxel_src = self.extract_voxel_from_image(pseudoimages_src, voxels_src) # [B, N_valid_voxels, C] + feats_voxel_dst = self.extract_voxel_from_image(pseudoimages_dst, voxels_dst) # [B, N_valid_voxels, C] + feats_voxel_dst_inflate = batched_masked_gather(feats_voxel_dst, knn_idxs_dst.long(), knn_idxs_dst>=0, fill_value=0) corr_src_dst = torch.nn.functional.cosine_similarity(feats_voxel_src[:, :, None, :], feats_voxel_dst_inflate, dim=-1) @@ -283,12 +285,9 @@ def _model_forward(self, points_src, points_dst): self.timer[1][4].stop() self.timer[1][5].start("Decoding") - if self.decoder=='simple_decoder': - flows = self.decoder(feats_point_vol) - else: - feats_cat = torch.cat([feats_point_vol, feats_point_src_init, feats_point_src_grid], -1) - - flows = self.decoder(feats_cat, point_offsets_src) + + feats_cat = torch.cat([feats_point_vol, feats_point_src_init, feats_point_src_grid], -1) + flows = self.decoder(feats_cat, point_offsets_src) self.timer[1][5].stop() else: feats_voxel_src = None From 5e1d0da4ee88d2bd218b3d233803897f809d68bd Mon Sep 17 00:00:00 2001 From: Shiming Date: Mon, 2 Jun 2025 14:14:08 +0000 Subject: [PATCH 04/10] added lr scheduler and uses the default unet --- src/models/voteflow.py | 5 +++-- src/trainer.py | 11 +++++++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/src/models/voteflow.py b/src/models/voteflow.py index 05c7b29..8f32e6a 100644 --- a/src/models/voteflow.py +++ b/src/models/voteflow.py @@ -15,10 +15,11 @@ import pytorch3d.ops as pytorch3d_ops from .basic.encoder import DynamicEmbedder +from .basic.unet import FastFlow3DUNet from .basic import cal_pose0to1 from .basic.voteflow_plugin.hough_transformation import HT_CUDA -from .basic.voteflow_plugin.voteflow_module import VolConvBN, VoteFlowLinearDecoder, FastFlow3DUNet +from .basic.voteflow_plugin.voteflow_module import VolConvBN, VoteFlowLinearDecoder from .basic.voteflow_plugin.utils import calculate_unq_voxels, batched_masked_gather, pad_to_batch import warnings @@ -57,7 +58,7 @@ def __init__(self, self.voxel_size = voxel_size self.pseudo_image_dims = pseudo_image_dims #self.backbone = Backbone(input_channels, output_channels) - self.backbone = FastFlow3DUNet(input_channels, output_channels) ## output_channel 64 + self.backbone = FastFlow3DUNet() ## output_channel 64 if self.using_voting: self.decoder = VoteFlowLinearDecoder(dim_input= output_channels * 2 + input_channels * 2, layer_size=decoder_layers) diff --git a/src/trainer.py b/src/trainer.py index 8f048f3..e4df550 100644 --- a/src/trainer.py +++ b/src/trainer.py @@ -62,6 +62,7 @@ def __init__(self, cfg, eval=False): self.batch_size = int(cfg.batch_size) if 'batch_size' in cfg else 1 self.lr = cfg.lr if 'lr' in cfg else None + self.lr_scheduler = cfg.lr_scheduler if 'lr_scheduler' in cfg else None self.epochs = cfg.epochs if 'epochs' in cfg else None self.metrics = OfficialMetrics() @@ -171,8 +172,14 @@ def train_validation_step_(self, batch, res_dict): pass def configure_optimizers(self): - optimizer = optim.Adam(self.model.parameters(), lr=self.lr) - return optimizer + if self.lr is None: + optimizer = optim.Adam(self.model.parameters(), lr=self.lr) + return optimizer + elif self.lr_scheduler == 'step': + optimizer = optim.Adam(self.model.parameters(), lr=self.lr) + lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=self.trainer.max_epochs//2, gamma=0.1) + return {'optimizer': optimizer, + 'lr_scheduler': lr_scheduler} def on_train_epoch_start(self): self.time_start_train_epoch = time.time() From b96af0fc04417c0617358491ff96eb174ec74a10 Mon Sep 17 00:00:00 2001 From: Shiming Date: Tue, 26 Aug 2025 13:48:02 +0000 Subject: [PATCH 05/10] finalized the configuration --- conf/config.yaml | 1 + conf/model/voteflow.yaml | 2 +- src/models/voteflow.py | 4 ++-- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/conf/config.yaml b/conf/config.yaml index 1e34f76..5386c50 100644 --- a/conf/config.yaml +++ b/conf/config.yaml @@ -27,6 +27,7 @@ gradient_clip_val: 5.0 # optimizer ==> Adam lr: 2e-6 +lr_scheduler: step # choices: [cosine, step, linear] loss_fn: deflowLoss # choices: [ff3dLoss, zeroflowLoss, deflowLoss, seflowLoss] add_seloss: # {chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0} diff --git a/conf/model/voteflow.yaml b/conf/model/voteflow.yaml index 9ce5d3a..3362d8e 100644 --- a/conf/model/voteflow.yaml +++ b/conf/model/voteflow.yaml @@ -11,7 +11,7 @@ target: point_cloud_range: ${point_cloud_range} voxel_size: ${voxel_size} grid_feature_size: [512, 512] - decoder_layers: 1 + decoder_layers: 4 use_ball_query: False vol_conv_hidden_dim: 16 diff --git a/src/models/voteflow.py b/src/models/voteflow.py index 8f32e6a..5d034b9 100644 --- a/src/models/voteflow.py +++ b/src/models/voteflow.py @@ -30,14 +30,14 @@ class VoteFlow(nn.Module): def __init__(self, nframes=1, m=8, - n=64, + n=128, using_voting=True, input_channels=32, output_channels=64, point_cloud_range = [-51.2, -51.2, -3, 51.2, 51.2, 3], voxel_size=(0.2, 0.2, 6), grid_feature_size = [512, 512], - decoder_layers=1, + decoder_layers=4, use_ball_query=False, vol_conv_hidden_dim=16, **kwargs): From 96b2bd24c23e64302d31af48394804e50d21f2c3 Mon Sep 17 00:00:00 2001 From: Shiming Date: Tue, 26 Aug 2025 16:48:38 +0000 Subject: [PATCH 06/10] finalized the voteflow and updated the checkpoint link --- README.md | 18 ++++++++---------- conf/config.yaml | 2 +- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index b2a7211..8d3020b 100644 --- a/README.md +++ b/README.md @@ -11,16 +11,16 @@ OpenSceneFlow is a codebase for point cloud scene flow estimation. It is also an official implementation of the following papers (sorted by the time of publication): -- **VoteFlow:VoteFlow: Enforcing Local Rigidity in Self-Supervised Scene Flow** -*Yancong Lin\*, Shiming Wang\*, Liangliang Nan, Julian Kooij, Holger Caesar* -Conference on Computer Vision and Pattern Recognition (**CVPR**) 2025 -[ Backbone ] [ Self-Supervised ] - [ [arXiv](https://arxiv.org/abs/2503.22328) ] [ [Project](https://github.com/tudelft-iv/VoteFlow/)] - - **HiMo: High-Speed Objects Motion Compensation in Point Clouds** (SeFlow++) *Qingwen Zhang, Ajinkya Khoche, Yi Yang, Li Ling, Sina Sharif Mansouri, Olov Andersson, Patric Jensfelt* Preprint; Under review; 2025 [ Strategy ] [ Self-Supervised ] - [ [arXiv](https://arxiv.org/abs/2503.00803) ] [ [Project](https://kin-zhang.github.io/HiMo/) ] +- **VoteFlow:VoteFlow: Enforcing Local Rigidity in Self-Supervised Scene Flow** +*Yancong Lin\*, Shiming Wang\*, Liangliang Nan, Julian Kooij, Holger Caesar* +Conference on Computer Vision and Pattern Recognition (**CVPR**) 2025 +[ Backbone ] [ Self-Supervised ] - [ [arXiv](https://arxiv.org/abs/2503.22328) ] [ [Project](https://github.com/tudelft-iv/VoteFlow/)] → [here](#VoteFLow) + - **Flow4D: Leveraging 4D Voxel Network for LiDAR Scene Flow Estimation** *Jaeyeul Kim, Jungwan Woo, Ukcheol Shin, Jean Oh, Sunghoon Im* IEEE Robotics and Automation Letters (**RA-L**) 2025 @@ -146,12 +146,10 @@ Install torch-scatter pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.0+cu117.html ``` ```bash -python train.py model=voteflow lr=2e-4 epochs=12 batch_size=16 model.target.m=8 model.target.n=128 loss_fn=seflowLoss "add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}" -``` -Pretrained weight can be downloaded through -``` -PLACEHOLDER +python train.py model=voteflow lr=2e-4 lr_scheduler=step epochs=12 batch_size=16 model.target.m=8 model.target.n=128 loss_fn=seflowLoss "add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}" ``` +Pretrained weight can be downloaded through the [link](https://surfdrive.surf.nl/files/index.php/s/wxk1GxwvDKe9TWW) + ### Flow4D Train Flow4D with the leaderboard submit config. [Runtime: Around 18 hours in 4x RTX 3090 GPUs.] diff --git a/conf/config.yaml b/conf/config.yaml index 5386c50..4a80895 100644 --- a/conf/config.yaml +++ b/conf/config.yaml @@ -27,7 +27,7 @@ gradient_clip_val: 5.0 # optimizer ==> Adam lr: 2e-6 -lr_scheduler: step # choices: [cosine, step, linear] +lr_scheduler: None # choices: [cosine, step, linear] loss_fn: deflowLoss # choices: [ff3dLoss, zeroflowLoss, deflowLoss, seflowLoss] add_seloss: # {chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0} From bd685ee949a7ff9194cf5985303dddb554d7585b Mon Sep 17 00:00:00 2001 From: Kin Date: Wed, 27 Aug 2025 00:20:11 +0200 Subject: [PATCH 07/10] style: remove useless import or repeated fn in common py already. * please add training setup and hour for user as a reference. * todo: check later on env again. --- README.md | 10 +- src/models/__init__.py | 8 +- src/models/basic/voteflow_plugin/__init__.py | 2 +- .../hough_transformation/ht_cuda.py | 11 -- .../hough_transformation/im2ht.py | 15 --- .../basic/voteflow_plugin/voteflow_module.py | 115 ------------------ 6 files changed, 14 insertions(+), 147 deletions(-) diff --git a/README.md b/README.md index 8d3020b..0036be9 100644 --- a/README.md +++ b/README.md @@ -135,16 +135,18 @@ mamba activate opensf ### VoteFLow Extra pakcges needed for VoteFlow, [pytorch3d](https://pytorch3d.org/) (prefer 0.7.7) and [torch-scatter](https://github.com/rusty1s/pytorch_scatter?tab=readme-ov-file) (prefer 2.1.2): -Install Pytorch3d + ```bash +# Install Pytorch3d conda install -c iopath iopath conda install -c bottler nvidiacub conda install pytorch3d -c pytorch3d -``` -Install torch-scatter -```bash + +# Install torch-scatter pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.0+cu117.html ``` + +Train VoteFlow with the leaderboard submit config. [Runtime: Around ? hours in ?x ? GPUs.] ```bash python train.py model=voteflow lr=2e-4 lr_scheduler=step epochs=12 batch_size=16 model.target.m=8 model.target.n=128 loss_fn=seflowLoss "add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}" ``` diff --git a/src/models/__init__.py b/src/models/__init__.py index 8e63dfb..410b773 100644 --- a/src/models/__init__.py +++ b/src/models/__init__.py @@ -12,9 +12,15 @@ from .deflow import DeFlow from .fastflow3d import FastFlow3D -from .voteflow import VoteFlow from .nsfp import NSFP +# check README for package: +try: + from .voteflow import VoteFlow +except ImportError as e: + print("\033[93m--- WARNING [model]: VoteFlow is not imported, as it requires some lib which is not installed.") + print(f"Detail error message\033[0m: {e}. Just ignore this warning if code runs without these models.") + # following need install extra package: # * pip install spconv-cu117 try: diff --git a/src/models/basic/voteflow_plugin/__init__.py b/src/models/basic/voteflow_plugin/__init__.py index a04b2bc..55a35ae 100644 --- a/src/models/basic/voteflow_plugin/__init__.py +++ b/src/models/basic/voteflow_plugin/__init__.py @@ -1,3 +1,3 @@ from .hough_transformation import HT_CUDA from .utils import calculate_unq_voxels, batched_masked_gather, pad_to_batch -from .voteflow_module import VolConvBN, VoteFlowLinearDecoder, FastFlow3DUNet \ No newline at end of file +from .voteflow_module import VolConvBN, VoteFlowLinearDecoder \ No newline at end of file diff --git a/src/models/basic/voteflow_plugin/hough_transformation/ht_cuda.py b/src/models/basic/voteflow_plugin/hough_transformation/ht_cuda.py index c989bed..0acc055 100755 --- a/src/models/basic/voteflow_plugin/hough_transformation/ht_cuda.py +++ b/src/models/basic/voteflow_plugin/hough_transformation/ht_cuda.py @@ -1,15 +1,4 @@ -import torch import torch.nn as nn -import torch.nn.functional as F -from torch.autograd import gradcheck -from torch.nn.modules.utils import _pair -import scipy.io as sio -import random -import numpy as np -import matplotlib.pyplot as plt -import math -import time -import os from .im2ht import IM2HT class HT_CUDA(nn.Module): diff --git a/src/models/basic/voteflow_plugin/hough_transformation/im2ht.py b/src/models/basic/voteflow_plugin/hough_transformation/im2ht.py index 802f059..e4d12e7 100755 --- a/src/models/basic/voteflow_plugin/hough_transformation/im2ht.py +++ b/src/models/basic/voteflow_plugin/hough_transformation/im2ht.py @@ -1,13 +1,9 @@ import os -import math -import numpy as np import warnings from glob import glob -import torch from torch import nn from torch.autograd import Function -from torch.nn.modules.utils import _pair from torch.autograd.function import once_differentiable def load_cpp_ext(ext_name): @@ -95,17 +91,6 @@ def __init__(self, ): global im2ht im2ht = load_cpp_ext("im2ht") - # self.extra_repr() - # self.__repr__() - - # def extra_repr(self): - # s = ('size={ht_size}') - # # return s.format(**self.__dict__) - # return print(s.format(**self.__dict__)) - - # def __repr__(self): - # return self.__class__.__name__ + f'( m: {self.m:04d}, n: {self.n:04d}, h: {self.h:04d}, w: {self.w:04d}, d: {self.d:04d})' - def forward(self, feats_src_dst, voxels_src, voxels_dst, idxs_src, idxs_dst, h, w, d): # print('IM2HT forward', feats.shape, idxs_fps.shape, idxs_src.shape, idxs_dst.shape, bins_x.shape, bins_y.shape, bins_z.shape) return IM2HTFunction.apply( diff --git a/src/models/basic/voteflow_plugin/voteflow_module.py b/src/models/basic/voteflow_plugin/voteflow_module.py index 518ae17..bef7287 100644 --- a/src/models/basic/voteflow_plugin/voteflow_module.py +++ b/src/models/basic/voteflow_plugin/voteflow_module.py @@ -81,118 +81,3 @@ def forward(self, x: torch.Tensor, pts_offsets: torch.Tensor) -> torch.Tensor: x = self.decoder(x) return x - - -class ConvWithNorms(nn.Module): - - def __init__(self, in_num_channels: int, out_num_channels: int, - kernel_size: int, stride: int, padding: int): - super().__init__() - self.conv = nn.Conv2d(in_num_channels, out_num_channels, kernel_size, - stride, padding) - self.batchnorm = nn.BatchNorm2d(out_num_channels) - self.nonlinearity = nn.GELU() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - conv_res = self.conv(x) - if conv_res.shape[2] == 1 and conv_res.shape[3] == 1: - # This is a hack to get around the fact that batchnorm doesn't support - # 1x1 convolutions - batchnorm_res = conv_res - else: - batchnorm_res = self.batchnorm(conv_res) - return self.nonlinearity(batchnorm_res) - -class BilinearDecoder(nn.Module): - - def __init__(self, scale_factor: int): - super().__init__() - self.scale_factor = scale_factor - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return nn.functional.interpolate(x, - scale_factor=self.scale_factor, - mode="bilinear", - align_corners=False) - -class UpsampleSkip(nn.Module): - - def __init__(self, skip_channels: int, latent_channels: int, - out_channels: int): - super().__init__() - self.u1_u2 = nn.Sequential( - nn.Conv2d(skip_channels, latent_channels, 1, 1, 0), - BilinearDecoder(2)) - self.u3 = nn.Conv2d(latent_channels, latent_channels, 1, 1, 0) - self.u4_u5 = nn.Sequential( - nn.Conv2d(2 * latent_channels, out_channels, 3, 1, 1), - nn.Conv2d(out_channels, out_channels, 3, 1, 1)) - - def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - u2_res = self.u1_u2(a) - u3_res = self.u3(b) - u5_res = self.u4_u5(torch.cat([u2_res, u3_res], dim=1)) - return u5_res - -class FastFlow3DUNet(nn.Module): - """ - Standard UNet with a few modifications: - - Uses Bilinear interpolation instead of transposed convolutions - """ - - def __init__(self, in_num_channels, out_num_channels) -> None: - super().__init__() - - self.encoder_step_1 = nn.Sequential(ConvWithNorms(in_num_channels, 64, 3, 2, 1), - ConvWithNorms(64, 64, 3, 1, 1), - ConvWithNorms(64, 64, 3, 1, 1), - ConvWithNorms(64, 64, 3, 1, 1)) - self.encoder_step_2 = nn.Sequential(ConvWithNorms(64, 128, 3, 2, 1), - ConvWithNorms(128, 128, 3, 1, 1), - ConvWithNorms(128, 128, 3, 1, 1), - ConvWithNorms(128, 128, 3, 1, 1), - ConvWithNorms(128, 128, 3, 1, 1), - ConvWithNorms(128, 128, 3, 1, 1)) - self.encoder_step_3 = nn.Sequential(ConvWithNorms(128, 256, 3, 2, 1), - ConvWithNorms(256, 256, 3, 1, 1), - ConvWithNorms(256, 256, 3, 1, 1), - ConvWithNorms(256, 256, 3, 1, 1), - ConvWithNorms(256, 256, 3, 1, 1), - ConvWithNorms(256, 256, 3, 1, 1)) - self.decoder_step1 = UpsampleSkip(512, 256, 256) - self.decoder_step2 = UpsampleSkip(256, 128, 128) - self.decoder_step3 = UpsampleSkip(128, 64, 64) - self.decoder_step4 = nn.Conv2d(64, out_num_channels, 3, 1, 1) - - def forward(self, pc0_B: torch.Tensor, - pc1_B: torch.Tensor) -> torch.Tensor: - - expected_channels = 32 - assert pc0_B.shape[ - 1] == expected_channels, f"Expected {expected_channels} channels, got {pc0_B.shape[1]}" - assert pc1_B.shape[ - 1] == expected_channels, f"Expected {expected_channels} channels, got {pc1_B.shape[1]}" - - pc0_F = self.encoder_step_1(pc0_B) - pc0_L = self.encoder_step_2(pc0_F) - pc0_R = self.encoder_step_3(pc0_L) - - pc1_F = self.encoder_step_1(pc1_B) - pc1_L = self.encoder_step_2(pc1_F) - pc1_R = self.encoder_step_3(pc1_L) - - Rstar = torch.cat([pc0_R, pc1_R], - dim=1) # torch.Size([1, 512, 64, 64]) - Lstar = torch.cat([pc0_L, pc1_L], - dim=1) # torch.Size([1, 256, 128, 128]) - Fstar = torch.cat([pc0_F, pc1_F], - dim=1) # torch.Size([1, 128, 256, 256]) - Bstar = torch.cat([pc0_B, pc1_B], - dim=1) # torch.Size([1, 64, 512, 512]) - - S = self.decoder_step1(Rstar, Lstar) - T = self.decoder_step2(S, Fstar) - U = self.decoder_step3(T, Bstar) - V = self.decoder_step4(U) - - return V \ No newline at end of file From ff6bd101c38540fdcb24c340bd0f8d010bc6dddb Mon Sep 17 00:00:00 2001 From: Kin Date: Wed, 27 Aug 2025 00:29:31 +0200 Subject: [PATCH 08/10] docs: fix some typo here. --- README.md | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/README.md b/README.md index 0036be9..273e160 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ It is also an official implementation of the following papers (sorted by the tim Preprint; Under review; 2025 [ Strategy ] [ Self-Supervised ] - [ [arXiv](https://arxiv.org/abs/2503.00803) ] [ [Project](https://kin-zhang.github.io/HiMo/) ] -- **VoteFlow:VoteFlow: Enforcing Local Rigidity in Self-Supervised Scene Flow** +- **VoteFlow: Enforcing Local Rigidity in Self-Supervised Scene Flow** *Yancong Lin\*, Shiming Wang\*, Liangliang Nan, Julian Kooij, Holger Caesar* Conference on Computer Vision and Pattern Recognition (**CVPR**) 2025 [ Backbone ] [ Self-Supervised ] - [ [arXiv](https://arxiv.org/abs/2503.22328) ] [ [Project](https://github.com/tudelft-iv/VoteFlow/)] → [here](#VoteFLow) @@ -138,8 +138,6 @@ Extra pakcges needed for VoteFlow, [pytorch3d](https://pytorch3d.org/) (prefer 0 ```bash # Install Pytorch3d -conda install -c iopath iopath -conda install -c bottler nvidiacub conda install pytorch3d -c pytorch3d # Install torch-scatter From 546c4e2cb371532cd5a45866e8a96d23258cde8e Mon Sep 17 00:00:00 2001 From: Kin Date: Wed, 27 Aug 2025 00:57:48 +0200 Subject: [PATCH 09/10] docs: clean up some fn and readme upload ckpt to huggingface for easy wget. --- README.md | 7 ++++++- src/models/basic/voteflow_plugin/utils.py | 3 --- src/models/basic/voteflow_plugin/voteflow_module.py | 1 - src/models/voteflow.py | 1 - 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 273e160..b7134c5 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,7 @@ Additionally, *OpenSceneFlow* integrates following excellent works: [ICLR'24 Zer - [x] [NSFP](https://arxiv.org/abs/2111.01253): NeurIPS 2021, faster 3x than original version because of [our CUDA speed up](assets/cuda/README.md), same (slightly better) performance. - [x] [FastNSF](https://arxiv.org/abs/2304.09121): ICCV 2023. SSL optimization-based. - [ ] [ICP-Flow](https://arxiv.org/abs/2402.17351): CVPR 2024. SSL optimization-based. Done coding, public after review. +- [ ] [EulerFlow](https://arxiv.org/abs/2410.02031): ICLR 2025. SSL optimization-based. In my plan, haven't coding yet. @@ -148,7 +149,11 @@ Train VoteFlow with the leaderboard submit config. [Runtime: Around ? hours in ? ```bash python train.py model=voteflow lr=2e-4 lr_scheduler=step epochs=12 batch_size=16 model.target.m=8 model.target.n=128 loss_fn=seflowLoss "add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}" ``` -Pretrained weight can be downloaded through the [link](https://surfdrive.surf.nl/files/index.php/s/wxk1GxwvDKe9TWW) + +Pretrained weight can be downloaded through: +```bash +wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/voteflow_best.ckpt +``` ### Flow4D diff --git a/src/models/basic/voteflow_plugin/utils.py b/src/models/basic/voteflow_plugin/utils.py index ec1f843..e823fd1 100644 --- a/src/models/basic/voteflow_plugin/utils.py +++ b/src/models/basic/voteflow_plugin/utils.py @@ -1,7 +1,4 @@ -import math -import numpy as np import torch -from torch_scatter import scatter_max import pytorch3d.ops as pytorch3d_ops def calculate_unq_voxels(coords, image_dims): diff --git a/src/models/basic/voteflow_plugin/voteflow_module.py b/src/models/basic/voteflow_plugin/voteflow_module.py index bef7287..25e448c 100644 --- a/src/models/basic/voteflow_plugin/voteflow_module.py +++ b/src/models/basic/voteflow_plugin/voteflow_module.py @@ -1,7 +1,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -from typing import List, Tuple, Dict import math def set_grad(var): diff --git a/src/models/voteflow.py b/src/models/voteflow.py index 5d034b9..0b708d3 100644 --- a/src/models/voteflow.py +++ b/src/models/voteflow.py @@ -57,7 +57,6 @@ def __init__(self, self.point_cloud_range = point_cloud_range self.voxel_size = voxel_size self.pseudo_image_dims = pseudo_image_dims - #self.backbone = Backbone(input_channels, output_channels) self.backbone = FastFlow3DUNet() ## output_channel 64 if self.using_voting: From 0a1f546c537fa0a40c71951777ef46a783883dfa Mon Sep 17 00:00:00 2001 From: Shiming Date: Wed, 27 Aug 2025 08:55:30 +0000 Subject: [PATCH 10/10] updated train config, the training hours and deleted the cuda test file --- README.md | 4 ++-- .../hough_transformation/cpp_im2ht/ht_setup.sh | 7 ------- 2 files changed, 2 insertions(+), 9 deletions(-) delete mode 100755 src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_setup.sh diff --git a/README.md b/README.md index b7134c5..a1aa2ae 100644 --- a/README.md +++ b/README.md @@ -145,9 +145,9 @@ conda install pytorch3d -c pytorch3d pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.0+cu117.html ``` -Train VoteFlow with the leaderboard submit config. [Runtime: Around ? hours in ?x ? GPUs.] +Train VoteFlow with the leaderboard submit config. [Runtime: Around 32 hours in 4 x V100 GPUs.] ```bash -python train.py model=voteflow lr=2e-4 lr_scheduler=step epochs=12 batch_size=16 model.target.m=8 model.target.n=128 loss_fn=seflowLoss "add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}" +python train.py model=voteflow lr=2e-4 lr_scheduler=step epochs=12 batch_size=4 model.target.m=8 model.target.n=128 loss_fn=seflowLoss "add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}" ``` Pretrained weight can be downloaded through: diff --git a/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_setup.sh b/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_setup.sh deleted file mode 100755 index 73b833d..0000000 --- a/src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_setup.sh +++ /dev/null @@ -1,7 +0,0 @@ -cd cpp_im2ht -rm -rf build -python setup.py clean -python setup.py build -python setup.py install --user -cd .. -