Merge VoteFlow into codebase

.gitignore

@@ -25,4 +25,7 @@ cancel.sh
 # cuda build files
 *.egg-info
 build*
-dist*
+dist*
+
+data
+*__pycache__

README.md

@@ -16,6 +16,11 @@ 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: 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  
@@ -46,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.
 
 </details>
 
@@ -128,6 +134,27 @@ 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
+# Install Pytorch3d
+conda install pytorch3d -c pytorch3d
+
+# 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 32 hours in 4 x V100 GPUs.]
+```bash
+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:
+```bash
+wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/voteflow_best.ckpt
+```
+
 ### Flow4D
 
 Train Flow4D with the leaderboard submit config. [Runtime: Around 18 hours in 4x RTX 3090 GPUs.]
@@ -327,6 +354,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! ❤️

conf/config.yaml

@@ -27,6 +27,7 @@ gradient_clip_val: 5.0
 
 # optimizer ==> Adam
 lr: 2e-6
+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}
 

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: 4
+  use_ball_query: False
+  vol_conv_hidden_dim: 16
+
+val_monitor: val/Dynamic/Mean

src/models/__init__.py

@@ -14,6 +14,13 @@
 from .fastflow3d import FastFlow3D
 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:

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, VoteFlowLinearDecoder

src/models/basic/voteflow_plugin/hough_transformation/__init__.py

@@ -0,0 +1,4 @@
+from .ht_cuda import HT_CUDA
+__all__ = [
+    'HT_CUDA',
+]

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 <ATen/ATen.h>
+// #include <ATen/cuda/CUDAContext.h>
+// #include <cmath>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+
+#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)

src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_cuda.cu

@@ -0,0 +1,116 @@
+#include <vector>
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cmath>
+#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<scalar_t>() ,
+                    vol_ht.data<scalar_t>(),
+                    voxels_src.data<scalar_t>() ,
+                    voxels_dst.data<scalar_t>(),
+                    idxs_src.data<scalar_t>(),
+                    idxs_dst.data<scalar_t>(),
+                    b, l, c,
+                    m, n, 
+                    h, w, d
+                    );
+
+    }));
+    // std::cout <<"output" <<output.sum() << std::endl;
+    // printf("output", output.sum());
+    return vol_ht;
+}
+
+
+
+// std::vector<at::Tensor>
+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<scalar_t>(),
+                    grad_vol.data<scalar_t>(),
+                    voxels_src.data<scalar_t>(),
+                    voxels_dst.data<scalar_t>(),
+                    idxs_src.data<scalar_t>(),
+                    idxs_dst.data<scalar_t>(),
+                    b, l, c,
+                    m, n, 
+                    h, w, d
+                );
+
+    }));
+
+    return grad_feats_src_dst; 
+}

src/models/basic/voteflow_plugin/hough_transformation/cpp_im2ht/ht_cuda.h

@@ -0,0 +1,29 @@
+#pragma once
+#include <torch/extension.h>
+
+
+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>
+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
+                );
+