Merge ICP-Flow into codebase

README.md

@@ -41,21 +41,22 @@ International Conference on Robotics and Automation (**ICRA**) 2025
 European Conference on Computer Vision (**ECCV**) 2024  
 [ Strategy ] [ Self-Supervised ] - [ [arXiv](https://arxiv.org/abs/2407.01702) ] [ [Project](https://github.com/KTH-RPL/SeFlow) ] → [here](#seflow)
 
+
 - **DeFlow: Decoder of Scene Flow Network in Autonomous Driving**  
 *Qingwen Zhang, Yi Yang, Heng Fang, Ruoyu Geng, Patric Jensfelt*  
 International Conference on Robotics and Automation (**ICRA**) 2024  
 [ Backbone ] [ Supervised ] - [ [arXiv](https://arxiv.org/abs/2401.16122) ] [ [Project](https://github.com/KTH-RPL/DeFlow) ] → [here](#deflow)
 
 🎁 <b>One repository, All methods!</b> 
-Additionally, *OpenSceneFlow* integrates following excellent works: [ICLR'24 ZeroFlow](https://arxiv.org/abs/2305.10424), [ICCV'23 FastNSF](https://arxiv.org/abs/2304.09121), [RA-L'21 FastFlow3D](https://arxiv.org/abs/2103.01306), [NeurIPS'21 NSFP](https://arxiv.org/abs/2111.01253). (More on the way...)
+Additionally, *OpenSceneFlow* integrates following excellent works: [ICLR'24 ZeroFlow](https://arxiv.org/abs/2305.10424), [CVPR'24 ICP-Flow](https://arxiv.org/abs/2402.17351), [ICCV'23 FastNSF](https://arxiv.org/abs/2304.09121), [RA-L'21 FastFlow3D](https://arxiv.org/abs/2103.01306), [NeurIPS'21 NSFP](https://arxiv.org/abs/2111.01253). (More on the way...)
 
 <details> <summary> Summary of them:</summary>
 
 - [x] [FastFlow3D](https://arxiv.org/abs/2103.01306): RA-L 2021, a basic backbone model.
 - [x] [ZeroFlow](https://arxiv.org/abs/2305.10424): ICLR 2024, their pre-trained weight can covert into our format easily through [the script](tools/zerof2ours.py).
 - [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.
+- [x] [FastNSF](https://arxiv.org/abs/2304.09121): ICCV 2023. SSL Optimization-based.
+- [x] [ICP-Flow](https://arxiv.org/abs/2402.17351): CVPR 2024. SSL Optimization-based.
 - [ ] [EulerFlow](https://arxiv.org/abs/2410.02031): ICLR 2025. SSL optimization-based. In my plan, haven't coding yet.
 
 </details>
@@ -145,7 +146,7 @@ wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/flow4d_best.ckp
 
 #### SSF
 
-Extra pakcges needed for SSF model:
+Extra packages needed for SSF model:
 ```bash
 pip install mmengine-lite
 pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.0+cu117.html
@@ -223,6 +224,20 @@ python train.py model=deflowpp save_top_model=3 val_every=3 voxel_size="[0.2, 0.
 wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/seflowpp_best.ckpt
 ```
 
+
+### Optimization-based Unsupervised Methods
+
+For all optimization-based methods, you can directly run `eval.py`/`save.py` to get the result without training, while the running might take really long time, maybe tmux for run it. For multi-program running, the master port can be set through `+master_port=12346`.
+
+```bash
+# you can change another model by passing model name.
+python eval.py model=fastnsf
+
+# or save the result directly
+python save.py model=fastnsf
+```
+
+
 ## 3. Evaluation
 
 You can view Wandb dashboard for the training and evaluation results or upload result to online leaderboard.

assets/cuda/README.md

@@ -5,6 +5,8 @@ Faster our code in CUDA.
 
 - chamfer3D: 3D chamfer distance within two point cloud, by Qingwen Zhang involved when she was working on SeFlow.
 - mmcv: directly from mmcv, not our code.
+- mmdet: only python file, no need to compile
+- histlib: from Yancong's [ICP-Flow](https://github.com/yanconglin/ICP-Flow) project.
 
 ---
 

assets/cuda/histlib/__init__.py

@@ -0,0 +1,71 @@
+from torch import nn
+from torch.autograd import Function
+import torch
+import importlib
+
+import os, time
+import hist
+
+def histf(X, Y, min_x, min_y, min_z, max_x, max_y, max_z, len_x, len_y, len_z, mini_batch=8):
+    # print('hist cuda params: ', X.shape, Y.shape,
+    #       min_x, min_y, min_z,
+    #       max_x, max_y, max_z,
+    #       len_x, len_y, len_z,
+    #       )
+    histogram = hist.hist(X.contiguous(), Y.contiguous(), 
+                          min_x, min_y, min_z,
+                          max_x, max_y, max_z, 
+                          len_x, len_y, len_z, 
+                          mini_batch
+                          )
+    return histogram
+
+
+torch.manual_seed(2022)
+
+########################
+def run_test():
+    pts = torch.randn(3, 1000, 3)
+    indicators = torch.randint(0, 2, size=(3, 1000, 1))
+    pts1 = torch.cat([pts, indicators], dim=-1)
+    pts2  = pts1.clone()
+    pts2[:, :,0] += 5.
+    pts2[:, :,1] += -3.
+    pts2[:, :,2] += -0.2
+
+    range_x = 10.
+    range_y = 10.
+    range_z = 0.5
+    thres =0.1
+    # bins_x = torch.linspace(-range_x, range_x, int(2*range_x/thres)+1)
+    # bins_y = torch.linspace(-range_y, range_y, int(2*range_y/thres)+1)
+    # bins_z = torch.linspace(-range_z, range_z, int(2*range_z/thres)+1)
+    bins_x = torch.arange(-range_x, range_x+thres, thres)
+    bins_y = torch.arange(-range_y, range_y+thres, thres)
+    bins_z = torch.arange(-range_z, range_z+thres, thres)
+    print('bins_x: ', bins_x)
+    print('bins_z: ', bins_z)
+    pts1 = pts1.cuda()
+    pts2 = pts2.cuda()
+    bins_x = bins_x.cuda()
+    bins_y = bins_y.cuda()
+    bins_z = bins_z.cuda()
+
+    t_hists = histf(pts1, pts2, 
+               -range_x, -range_y, -range_z,
+               range_x, range_y, range_z,
+               len(bins_x), len(bins_y), len(bins_z),
+               )
+    print('output shape: ', t_hists.shape)
+    b, h, w, d = t_hists.shape
+    for t_hist in t_hists:
+        t_argmax = torch.argmax(t_hist)
+        print(f't_argmax: {t_argmax}, {t_hist.max()} {h}, {w}, {d}, {t_argmax//d//w%h}, {t_argmax//d%w}, {t_argmax%d}')
+        print('t_argmax', t_argmax//d//w%h, t_argmax//d%w, t_argmax%d, bins_x[t_argmax//d//w%h], bins_y[t_argmax//d%w], bins_z[t_argmax%d])
+
+if __name__ == '__main__':
+
+    print("Pytorch version: ", torch.__version__)
+    print("GPU version: ", torch.cuda.get_device_name())
+
+    run_test()

assets/cuda/histlib/hist.cu

@@ -0,0 +1,90 @@
+#include <vector>
+#include "hist_cuda_core.cuh"
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+// #include <THC/THC.h>
+// #include <THC/THCAtomics.cuh>
+// #include <THC/THCDeviceUtils.cuh>
+
+// extern THCState *state;
+
+// author: Charles Shang
+// https://github.com/torch/cunn/blob/master/lib/THCUNN/generic/SpatialConvolutionMM.cu
+
+
+at::Tensor
+hist_cuda(const at::Tensor &X, const at::Tensor &Y,
+        const float min_x, const float min_y, const float min_z,
+        const float max_x, const float max_y, const float max_z,
+        const int len_x, const int len_y, const int len_z,
+        const int mini_batch
+            )
+{
+    // THCAssertSameGPU(THCudaTensor_checkGPU(state, 5, input, weight, bias, offset, mask));
+
+    AT_ASSERTM(X.is_contiguous(), "input tensor has to be contiguous");
+    AT_ASSERTM(Y.is_contiguous(), "input tensor has to be contiguous");
+
+    AT_ASSERTM(X.type().is_cuda(), "input must be a CUDA tensor");
+    AT_ASSERTM(Y.type().is_cuda(), "input must be a CUDA tensor");
+
+    const int batch = X.size(0); 
+    const int num_X = X.size(1);
+    const int dim = X.size(2);
+    const int num_Y = Y.size(1);
+
+    AT_ASSERTM((X.size(0) == Y.size(0)), "batch_X (%d) != batch_Y (%d).", X.size(0), Y.size(0));
+    AT_ASSERTM((X.size(2) == Y.size(2)), "dim_X (%d) != dim_Y (%d).", X.size(2), Y.size(2));
+
+    AT_ASSERTM((dim == 4), "dim (%d) != 4; 3 for (x, y, z); 1 for indicator,padded or not.", dim);
+
+    // printf("len: %d %d %f \n", len_x, len_y, len_z);
+    // printf("hist cuda coord: %f, %f, %f; %f, %f, %f; %f, %f, %f. \n", val_x, val_y, val_z, p_x, p_y, p_z, len_x, len_y, len_z);
+
+    // auto bins = at::zeros({batch, len_x, len_y, len_z}, X.options());
+    // AT_DISPATCH_FLOATING_TYPES(X.type(), "hist_cuda_core", ([&] {
+    //     hist_cuda_core(at::cuda::getCurrentCUDAStream(),
+    //                                     X.data<scalar_t>(), Y.data<scalar_t>(),
+    //                                     batch, dim, num_X, num_Y,
+    //                                     min_x, min_y, min_z, 
+    //                                     max_x, max_y, max_z, 
+    //                                     len_x, len_y, len_z, 
+    //                                     bins.data<scalar_t>());
+    //     }));
+
+    auto bins = at::zeros({batch, len_x, len_y, len_z}, X.options());
+
+    int iters = batch / mini_batch; 
+    if (batch % mini_batch != 0) 
+    { 
+        iters += 1; 
+    }
+
+    for (int i=0; i<iters; ++i)
+    {
+        int mini_batch_ = mini_batch;
+        if ((i+1) * mini_batch > batch) 
+        {
+            mini_batch_ = batch - i * mini_batch; 
+        }
+        // printf("iter: %d %d %d %d %d \n", i, iters, mini_batch_, mini_batch, batch);
+        AT_DISPATCH_FLOATING_TYPES(X.type(), "hist_cuda_core", ([&] {
+            hist_cuda_core(at::cuda::getCurrentCUDAStream(),
+                                            X.data<scalar_t>() + i*mini_batch*num_X*dim, 
+                                            Y.data<scalar_t>() + i*mini_batch*num_Y*dim, 
+                                            mini_batch_, dim, num_X, num_Y,
+                                            min_x, min_y, min_z, 
+                                            max_x, max_y, max_z, 
+                                            len_x, len_y, len_z, 
+                                            bins.data<scalar_t>()+i*mini_batch*len_x*len_y*len_z);
+            }));
+    }
+
+
+
+    return bins;
+}

assets/cuda/histlib/hist.h

@@ -0,0 +1,13 @@
+#pragma once
+#include <torch/extension.h>
+
+at::Tensor
+hist(const at::Tensor &X, const at::Tensor &Y,
+    const float min_x, const float min_y, const float min_z,
+    const float max_x, const float max_y, const float max_z,
+    const int len_x, const int len_y, const int len_z,
+    const int mini_batch
+    );
+
+
+

assets/cuda/histlib/hist_cuda.cpp

@@ -0,0 +1,27 @@
+#include "hist.h"
+#include "hist_cuda.h"
+
+at::Tensor
+hist(const at::Tensor &X, const at::Tensor &Y,
+        const float min_x, const float min_y, const float min_z,
+        const float max_x, const float max_y, const float max_z,
+        const int len_x, const int len_y, const int len_z,
+        const int mini_batch
+        )
+{
+
+    if (X.type().is_cuda() && Y.type().is_cuda())
+    {
+        return hist_cuda(X, Y,
+                        min_x, min_y, min_z, 
+                        max_x, max_y, max_z, 
+                        len_x, len_y, len_z,
+                        mini_batch
+                        );
+    }
+    AT_ERROR("Not implemented on the CPU");
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("hist", &hist, "hist");
+}

assets/cuda/histlib/hist_cuda.h

@@ -0,0 +1,10 @@
+#pragma once
+#include <torch/extension.h>
+
+at::Tensor
+hist_cuda(const at::Tensor &X, const at::Tensor &Y,
+        const float min_x, const float min_y, const float min_z,
+        const float max_x, const float max_y, const float max_z,
+        const int len_x, const int len_y, const int len_z,
+        const int mini_batch
+        );

assets/cuda/histlib/hist_cuda_core.cuh

@@ -0,0 +1,100 @@
+#include <cstdio>
+#include <algorithm>
+#include <cstring>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+// #include <THC/THC.h>
+#include <THC/THCAtomics.cuh>
+// #include <THC/THCDeviceUtils.cuh>
+
+#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;
+}
+
+template <typename scalar_t>
+__global__ void hist_cuda_kernel(const int n,
+                                              const scalar_t* X,
+                                              const scalar_t* Y,
+                                              const int batch, const int dim, 
+                                              const int num_X, const int num_Y, 
+                                              const float min_x, const float min_y, const float min_z, 
+                                              const float max_x, const float max_y, const float max_z, 
+                                              const int len_x, const int len_y, const int len_z,
+                                              scalar_t* bins
+                                              )
+{
+  CUDA_KERNEL_LOOP(index, n)
+  {
+    // index index of output matrix
+    // launch in parallel:  batch * numX * numY;
+    // printf("hist cuda bin size: %d, %d, %d, %d. \n", batch, len_x, len_y, len_z);
+    const int b = index / num_X / num_Y % batch;
+    const int i = index / num_Y % num_X;
+    const int j = index  % num_Y;
+
+    scalar_t flag_x = X[b*num_X*dim+i*dim+3];
+    scalar_t flag_y = Y[b*num_Y*dim+j*dim+3];
+    if (flag_x>0.0 && flag_y>0.0)
+    {
+      scalar_t val_x = X[b*num_X*dim+i*dim+0] - Y[b*num_Y*dim+j*dim+0];
+      scalar_t val_y = X[b*num_X*dim+i*dim+1] - Y[b*num_Y*dim+j*dim+1];
+      scalar_t val_z = X[b*num_X*dim+i*dim+2] - Y[b*num_Y*dim+j*dim+2];
+      if (val_x >= min_x && val_x < max_x && val_y >= min_y && val_y < max_y && val_z >= min_z && val_z < max_z)
+      {
+        // [): left included; right excluded.
+        int p_x = __float2int_rd( (val_x-min_x) / (max_x-min_x) * __int2float_rd(len_x));
+        int p_y = __float2int_rd( (val_y-min_y) / (max_y-min_y) * __int2float_rd(len_y));
+        int p_z = __float2int_rd( (val_z-min_z) / (max_z-min_z) * __int2float_rd(len_z));
+
+        // printf("hist cuda coord: %d, %d, %d, %d; %d, %d, %d, %d. \n", batch, len_x, len_y, len_z, b, p_x, p_y, p_z);
+        int bin_id = b*len_x*len_y*len_z + p_x*len_y*len_z  +  p_y*len_z + p_z;
+        atomicAdd(bins + bin_id, 1);
+      }
+  }
+  }
+}
+
+template <typename scalar_t>
+void hist_cuda_core(cudaStream_t stream,
+                              const scalar_t* X, const scalar_t* Y,
+                              const int batch, const int dim, 
+                              const int num_X, const int num_Y, 
+                              const float min_x, const float min_y, const float min_z, 
+                              const float max_x, const float max_y, const float max_z, 
+                              const int len_x, const int len_y, const int len_z,
+                              scalar_t* bins
+                              ) 
+{
+  const int num_kernels = batch * num_X * num_Y;
+  // printf("num kernels: %d\n", num_kernels);
+
+  // printf("hist cuda core: %f, %f, %f; %f, %f, %f; %f, %f, %f. \n", min_x, min_y, min_z, max_x, max_y, max_z, len_x, len_y, len_z);
+  // printf("hist cuda core: ", min_x, min_y, min_z, max_x, max_y, max_z, len_x, len_y, len_z, " \n");
+  hist_cuda_kernel<scalar_t>
+      <<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS, 0, stream>>>(
+      num_kernels, 
+      X, Y,
+      batch, dim,
+      num_X, num_Y,
+      min_x, min_y, min_z, 
+      max_x, max_y, max_z, 
+      len_x, len_y, len_z, 
+      bins
+      );
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess)
+  {
+    printf("error in hist_cuda_core: %s\n", cudaGetErrorString(err));
+  }
+}
+

assets/cuda/histlib/setup.py

@@ -0,0 +1,15 @@
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+setup(
+    name='hist',
+    ext_modules=[
+        CUDAExtension('hist', [
+            "/".join(__file__.split('/')[:-1] + ['hist_cuda.cpp']), # must named as xxx_cuda.cpp
+            "/".join(__file__.split('/')[:-1] + ['hist.cu']),
+        ]),
+    ],
+    cmdclass={
+        'build_ext': BuildExtension
+    },
+    version='1.0.1')