From b29f50bea654fe19c29a54c039cf0013a00bd1c0 Mon Sep 17 00:00:00 2001 From: Kin Date: Mon, 11 Aug 2025 22:11:14 +0200 Subject: [PATCH 1/9] cuda(histlib): cuda library from icp-flow project. --- assets/cuda/README.md | 2 + assets/cuda/histlib/hist.cu | 90 ++++++++++++++++++++++ assets/cuda/histlib/hist.h | 13 ++++ assets/cuda/histlib/hist_cuda.cpp | 27 +++++++ assets/cuda/histlib/hist_cuda.h | 10 +++ assets/cuda/histlib/hist_cuda_core.cuh | 100 +++++++++++++++++++++++++ assets/cuda/histlib/setup.py | 15 ++++ 7 files changed, 257 insertions(+) create mode 100644 assets/cuda/histlib/hist.cu create mode 100644 assets/cuda/histlib/hist.h create mode 100644 assets/cuda/histlib/hist_cuda.cpp create mode 100644 assets/cuda/histlib/hist_cuda.h create mode 100644 assets/cuda/histlib/hist_cuda_core.cuh create mode 100755 assets/cuda/histlib/setup.py diff --git a/assets/cuda/README.md b/assets/cuda/README.md index 621bcd3..5c1b05f 100644 --- a/assets/cuda/README.md +++ b/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. --- diff --git a/assets/cuda/histlib/hist.cu b/assets/cuda/histlib/hist.cu new file mode 100644 index 0000000..153ea43 --- /dev/null +++ b/assets/cuda/histlib/hist.cu @@ -0,0 +1,90 @@ +#include +#include "hist_cuda_core.cuh" + +#include +#include +#include +#include + +// #include +// #include +// #include + +// 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(), Y.data(), + // 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()); + // })); + + 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 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() + i*mini_batch*num_X*dim, + Y.data() + 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()+i*mini_batch*len_x*len_y*len_z); + })); + } + + + + return bins; +} diff --git a/assets/cuda/histlib/hist.h b/assets/cuda/histlib/hist.h new file mode 100644 index 0000000..31c15d6 --- /dev/null +++ b/assets/cuda/histlib/hist.h @@ -0,0 +1,13 @@ +#pragma once +#include + +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 + ); + + + diff --git a/assets/cuda/histlib/hist_cuda.cpp b/assets/cuda/histlib/hist_cuda.cpp new file mode 100644 index 0000000..23a2137 --- /dev/null +++ b/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"); +} diff --git a/assets/cuda/histlib/hist_cuda.h b/assets/cuda/histlib/hist_cuda.h new file mode 100644 index 0000000..80e5ff3 --- /dev/null +++ b/assets/cuda/histlib/hist_cuda.h @@ -0,0 +1,10 @@ +#pragma once +#include + +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 + ); diff --git a/assets/cuda/histlib/hist_cuda_core.cuh b/assets/cuda/histlib/hist_cuda_core.cuh new file mode 100644 index 0000000..7101fb3 --- /dev/null +++ b/assets/cuda/histlib/hist_cuda_core.cuh @@ -0,0 +1,100 @@ +#include +#include +#include + +#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) + +const int CUDA_NUM_THREADS = 1024; +inline int GET_BLOCKS(const int N) +{ + return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS; +} + +template +__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 +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 + <<>>( + 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)); + } +} + diff --git a/assets/cuda/histlib/setup.py b/assets/cuda/histlib/setup.py new file mode 100755 index 0000000..17c4fbb --- /dev/null +++ b/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') From 4cf1168339322a6e7b2e72fdfcddc93274dcb19c Mon Sep 17 00:00:00 2001 From: Kin Date: Mon, 11 Aug 2025 22:36:17 +0200 Subject: [PATCH 2/9] docs: update README with icp-flow in the official implementation. * conf(optimization-based): update all config files. * todo: update model file, double check with yancong and Qingwen confirm that: icp-flow results can be reproduced and tested. --- README.md | 29 +++++++++++++- assets/cuda/histlib/__init__.py | 71 +++++++++++++++++++++++++++++++++ conf/model/fastnsf.yaml | 1 - conf/model/icpflow.yaml | 17 ++++++++ conf/model/nsfp.yaml | 1 - 5 files changed, 115 insertions(+), 4 deletions(-) create mode 100644 assets/cuda/histlib/__init__.py create mode 100644 conf/model/icpflow.yaml diff --git a/README.md b/README.md index 51d5307..5d00ea6 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,11 @@ 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) +- **ICP-Flow: LiDAR Scene Flow Estimation with ICP** +*Yancong Lin, Holger Caesar* +Conference on Computer Vision and Pattern Recognition (**CVPR**) 2024 +[ Optimization-based ] [ Self-Supervised ] - [ [arXiv](https://arxiv.org/abs/2402.17351) ] [ [Project](https://github.com/yanconglin/ICP-Flow) ] → [here](#icp-flow) + - **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 @@ -45,7 +50,6 @@ Additionally, *OpenSceneFlow* integrates following excellent works: [ICLR'24 Zer - [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. @@ -143,7 +147,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 torch-scatter # torch-scatter might not working, then reinstall by: @@ -179,6 +183,21 @@ Pretrained weight can be downloaded through: wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/seflow_best.ckpt ``` +### ICP-Flow + +ICP-Flow is a optimization-based method, you can directly run `eval.py`/`save.py` to get the result. + +Extra packages needed for ICP-Flow model: +```bash +pip install pytorch3d assets/cuda/histlib +``` + +Then run as: +```bash +python eval.py model=icpflow +python save.py model=icpflow +``` + ### DeFlow Train DeFlow with the leaderboard submit config. [Runtime: Around 6-8 hours in 4x A100 GPUs.] Please change `batch_size&lr` accoordingly if you don't have enough GPU memory. (e.g. `batch_size=6` for 24GB GPU) @@ -327,6 +346,12 @@ And our excellent collaborators works contributed to this codebase also: journal={arXiv preprint arXiv:2501.17821}, year={2025} } +@article{lin2024icp, + title={ICP-Flow: LiDAR Scene Flow Estimation with ICP}, + author={Lin, Yancong and Caesar, Holger}, + booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)}, + year={2024} +} ``` Thank you for your support! โค๏ธ diff --git a/assets/cuda/histlib/__init__.py b/assets/cuda/histlib/__init__.py new file mode 100644 index 0000000..e10f897 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/conf/model/fastnsf.yaml b/conf/model/fastnsf.yaml index fd837bc..9805769 100644 --- a/conf/model/fastnsf.yaml +++ b/conf/model/fastnsf.yaml @@ -13,5 +13,4 @@ target: early_patience: 10 # default is 10, set 30 to get better results point_cloud_range: [-51.2, -51.2, -3, 51.2, 51.2, 3] -val_monitor: val/Dynamic/Mean iter_only: True \ No newline at end of file diff --git a/conf/model/icpflow.yaml b/conf/model/icpflow.yaml new file mode 100644 index 0000000..91d5548 --- /dev/null +++ b/conf/model/icpflow.yaml @@ -0,0 +1,17 @@ +name: icpflow + +target: + _target_: src.models.ICPFlow + thres_dist: 0.1 + translation_frame: 1.67 + chunk_size: 50 + thres_iou: 0.2 + max_points: 10000 + min_cluster_size: 20 + thres_box: 0.1 + thres_rot: 0.1 + thres_error: 0.2 + speed: 1.67 + point_cloud_range: [-51.2, -51.2, -3, 51.2, 51.2, 3] + +iter_only: True \ No newline at end of file diff --git a/conf/model/nsfp.yaml b/conf/model/nsfp.yaml index 6163b6b..bd6c173 100644 --- a/conf/model/nsfp.yaml +++ b/conf/model/nsfp.yaml @@ -11,5 +11,4 @@ target: early_patience: 30 # default is 30, set 100 to get better results (zeroflow setting.) point_cloud_range: [-51.2, -51.2, -3, 51.2, 51.2, 3] -val_monitor: val/Dynamic/Mean iter_only: True \ No newline at end of file From 21ca00c061264951ae3a32b1fd4e7317b38e7cad Mon Sep 17 00:00:00 2001 From: Kin Date: Mon, 11 Aug 2025 22:38:06 +0200 Subject: [PATCH 3/9] docs: fix small typo, and starting updating model file. --- README.md | 2 +- src/models/__init__.py | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 5d00ea6..b644b6a 100644 --- a/README.md +++ b/README.md @@ -33,7 +33,7 @@ European Conference on Computer Vision (**ECCV**) 2024 - **ICP-Flow: LiDAR Scene Flow Estimation with ICP** *Yancong Lin, Holger Caesar* -Conference on Computer Vision and Pattern Recognition (**CVPR**) 2024 +Conference on Computer Vision and Pattern Recognition (**CVPR**) 2024 [ Optimization-based ] [ Self-Supervised ] - [ [arXiv](https://arxiv.org/abs/2402.17351) ] [ [Project](https://github.com/yanconglin/ICP-Flow) ] → [here](#icp-flow) - **DeFlow: Decoder of Scene Flow Network in Autonomous Driving** diff --git a/src/models/__init__.py b/src/models/__init__.py index 4735f1c..6e3eaff 100644 --- a/src/models/__init__.py +++ b/src/models/__init__.py @@ -36,3 +36,11 @@ from .fastnsf import FastNSF except ImportError: print("--- WARNING [model]: FastNSF is not imported, as it requires FastGeodis lib which is not installed.") + +# following need install extra package: +# * pip install pytorch3d assets/cuda/histlib +try: + from .icpflow import ICPFlow +except ImportError: + print("--- WARNING [model]: ICPFlow is not imported, as it requires pytorch3d lib which is not installed.") + print(f"Detail error message\033[0m: {e}. Just ignore this warning if code runs without these models.") From 345db96551f21343095533c1099abc48153f47bc Mon Sep 17 00:00:00 2001 From: Kin Date: Sat, 16 Aug 2025 11:23:48 +0200 Subject: [PATCH 4/9] feat(icp): core icp files, tested successfully. --- README.md | 2 +- src/models/basic/icpflow_lib.py | 819 ++++++++++++++++++++++++++++++++ src/models/icpflow.py | 187 ++++++++ src/runner.py | 1 + 4 files changed, 1008 insertions(+), 1 deletion(-) create mode 100644 src/models/basic/icpflow_lib.py create mode 100644 src/models/icpflow.py diff --git a/README.md b/README.md index b644b6a..d4e04d4 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ ๐Ÿ’ž If you find [*OpenSceneFlow*](https://github.com/KTH-RPL/OpenSceneFlow) useful to your research, please cite [**our works** ๐Ÿ“–](#cite-us) and [give a star ๐ŸŒŸ](https://github.com/KTH-RPL/OpenSceneFlow) as encouragement. (เฉญหŠ๊’ณโ€‹ห‹)เฉญโœง 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): +It is also an official implementation of the following papers (sorted by the time of publication): - **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* diff --git a/src/models/basic/icpflow_lib.py b/src/models/basic/icpflow_lib.py new file mode 100644 index 0000000..86f2e0b --- /dev/null +++ b/src/models/basic/icpflow_lib.py @@ -0,0 +1,819 @@ + +""" +This file is from: https://github.com/yanconglin/ICP-Flow +with slightly modification to have unified format with all benchmark. +""" + +import torch +import os, sys, time, fire +BASE_DIR = os.path.abspath(os.path.join( os.path.dirname( __file__ ), '../..' )) +sys.path.append(BASE_DIR) + +import pytorch3d.transforms as pytorch3d_t +import numpy as np +import pytorch3d.ops as p3d + +def nearest_neighbor_batch(src, dst): + assert src.dim()==3 + assert dst.dim()==3 + assert len(src)==len(dst) + b, num, dim = src.shape + assert src.shape[2]>=3 + assert dst.shape[2]>=3 + result = p3d.knn_points(src[:, :, 0:3], dst[:, :, 0:3], K=1) + dst_idxs = result.idx + distances = result.dists + return dst_idxs.view(b, num), distances.view(b, num).sqrt() + +# modified the original pytorch3d icp +# check the dif between open3d icp and pytorch3d icp +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from typing import Optional, Tuple, TYPE_CHECKING, Union +import warnings +from typing import List, NamedTuple, Optional, TYPE_CHECKING, Union + +import torch +from pytorch3d.ops import knn_points +from pytorch3d.ops import utils as oputil +from pytorch3d.structures import utils as strutil +if TYPE_CHECKING: + from pytorch3d.structures.pointclouds import Pointclouds + +# named tuples for inputs/outputs +class SimilarityTransform(NamedTuple): + R: torch.Tensor + T: torch.Tensor + s: torch.Tensor + + +class ICPSolution(NamedTuple): + converged: bool + rmse: Union[torch.Tensor, None] + Xt: torch.Tensor + RTs: SimilarityTransform + t_history: List[SimilarityTransform] + + +def iterative_closest_point( + X: Union[torch.Tensor, "Pointclouds"], + Y: Union[torch.Tensor, "Pointclouds"], + init_transform: Optional[SimilarityTransform] = None, + thres: float = 0.1, + max_iterations: int = 100, + relative_rmse_thr: float = 1e-6, + estimate_scale: bool = False, + allow_reflection: bool = False, + verbose: bool = False, +) -> ICPSolution: + """ + Executes the iterative closest point (ICP) algorithm [1, 2] in order to find + a similarity transformation (rotation `R`, translation `T`, and + optionally scale `s`) between two given differently-sized sets of + `d`-dimensional points `X` and `Y`, such that: + + `s[i] X[i] R[i] + T[i] = Y[NN[i]]`, + + for all batch indices `i` in the least squares sense. Here, Y[NN[i]] stands + for the indices of nearest neighbors from `Y` to each point in `X`. + Note, however, that the solution is only a local optimum. + + Args: + **X**: Batch of `d`-dimensional points + of shape `(minibatch, num_points_X, d)` or a `Pointclouds` object. + **Y**: Batch of `d`-dimensional points + of shape `(minibatch, num_points_Y, d)` or a `Pointclouds` object. + **init_transform**: A named-tuple `SimilarityTransform` of tensors + `R`, `T, `s`, where `R` is a batch of orthonormal matrices of + shape `(minibatch, d, d)`, `T` is a batch of translations + of shape `(minibatch, d)` and `s` is a batch of scaling factors + of shape `(minibatch,)`. + **max_iterations**: The maximum number of ICP iterations. + **relative_rmse_thr**: A threshold on the relative root mean squared error + used to terminate the algorithm. + **estimate_scale**: If `True`, also estimates a scaling component `s` + of the transformation. Otherwise assumes the identity + scale and returns a tensor of ones. + **allow_reflection**: If `True`, allows the algorithm to return `R` + which is orthonormal but has determinant==-1. + **verbose**: If `True`, prints status messages during each ICP iteration. + + Returns: + A named tuple `ICPSolution` with the following fields: + **converged**: A boolean flag denoting whether the algorithm converged + successfully (=`True`) or not (=`False`). + **rmse**: Attained root mean squared error after termination of ICP. + **Xt**: The point cloud `X` transformed with the final transformation + (`R`, `T`, `s`). If `X` is a `Pointclouds` object, returns an + instance of `Pointclouds`, otherwise returns `torch.Tensor`. + **RTs**: A named tuple `SimilarityTransform` containing + a batch of similarity transforms with fields: + **R**: Batch of orthonormal matrices of shape `(minibatch, d, d)`. + **T**: Batch of translations of shape `(minibatch, d)`. + **s**: batch of scaling factors of shape `(minibatch, )`. + **t_history**: A list of named tuples `SimilarityTransform` + the transformation parameters after each ICP iteration. + + References: + [1] Besl & McKay: A Method for Registration of 3-D Shapes. TPAMI, 1992. + [2] https://en.wikipedia.org/wiki/Iterative_closest_point + """ + Xt = X[:, :, 0:3] + Yt = Y[:, :, 0:3] + b, size_X, dim = Xt.shape + if (Xt.shape[2] != Yt.shape[2]) or (Xt.shape[0] != Yt.shape[0]): + raise ValueError( + "Point sets X and Y have to have the same " + + "number of batches and data dimensions." + ) + + mask_X = X[:, :, -1]>0.0 + mask_Y = Y[:, :, -1]>0.0 + num_points_X = mask_X.sum(dim=-1) + num_points_Y = mask_Y.sum(dim=-1) + + # clone the initial point cloud + Xt_init = Xt.clone() + mask_X_init = mask_X.clone() + + if init_transform is not None: + # parse the initial transform from the input and apply to Xt + try: + R, T, s = init_transform + assert ( + R.shape == torch.Size((b, dim, dim)) + and T.shape == torch.Size((b, dim)) + and s.shape == torch.Size((b,)) + ) + except Exception: + raise ValueError( + "The initial transformation init_transform has to be " + "a named tuple SimilarityTransform with elements (R, T, s). " + "R are dim x dim orthonormal matrices of shape " + "(minibatch, dim, dim), T is a batch of dim-dimensional " + "translations of shape (minibatch, dim) and s is a batch " + "of scalars of shape (minibatch,)." + ) from None + # apply the init transform to the input point cloud + Xt = _apply_similarity_transform(Xt, R, T, s) + else: + # initialize the transformation with identity + R = oputil.eyes(dim, b, device=Xt.device, dtype=Xt.dtype) + T = Xt.new_zeros((b, dim)) + s = Xt.new_ones(b) + + prev_rmse = None + rmse = None + iteration = -1 + converged = False + + # initialize the transformation history + t_history = [] + + # the main loop over ICP iterations + for iteration in range(max_iterations): + knn_result = knn_points( + Xt, Yt, lengths1=num_points_X, lengths2=num_points_Y, K=1, return_nn=True, norm=2 + ) + Xt_nn_points = knn_result.knn[:, :, 0, :] + # print('knn points: ', valid.shape, Xt_init.shape, Xt_nn_points.shape) + + valid = knn_result.dists[:, :, 0]<=thres**2 + mask_X = torch.logical_and(mask_X_init, valid) + + X1 = Xt_init * mask_X[:, :, None] + X2 = Xt_nn_points * mask_X[:, :, None] + # print('knn points: ', valid.shape, mask_X.shape, Xt_nn_points.shape) + + # get the alignment of the nearest neighbors from Yt with Xt + R, T, s = corresponding_points_alignment( + X1, + X2, + weights=mask_X, + estimate_scale=estimate_scale, + allow_reflection=allow_reflection, + ) + + # apply the estimated similarity transform to Xt_init + Xt = _apply_similarity_transform(Xt_init, R, T, s) + # if iteration%2==0: + # src = Xt[0] + # dst = Yt[0] + # visualize_pcd(np.concatenate([src.cpu().numpy(), dst.cpu().numpy()], axis=0), + # np.concatenate([np.zeros((len(src)))+1, np.zeros((len(dst)))+2], axis=0), + # num_colors=3, + # title=f'registration, iter: {iteration} {len(src)} vs {len(dst)}') + + # add the current transformation to the history + t_history.append(SimilarityTransform(R, T, s)) + + # compute the root mean squared error + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`. + Xt_sq_diff = ((Xt - Xt_nn_points) ** 2).sum(2) + rmse = oputil.wmean(Xt_sq_diff[:, :, None], mask_X).sqrt()[:, 0, 0] + + # compute the relative rmse + if prev_rmse is None: + relative_rmse = rmse.new_ones(b) + else: + relative_rmse = (prev_rmse - rmse) / prev_rmse + + if verbose: + rmse_msg = ( + f"ICP iteration {iteration}: mean/max rmse = " + + f"{rmse.mean():1.2e}/{rmse.max():1.2e} " + + f"; mean relative rmse = {relative_rmse.mean():1.2e}" + ) + print(rmse_msg) + + # check for convergence + if (relative_rmse <= relative_rmse_thr).all(): + converged = True + break + + # update the previous rmse + prev_rmse = rmse + + if verbose: + if converged: + print(f"ICP has converged in {iteration + 1} iterations.") + else: + print(f"ICP has not converged in {max_iterations} iterations.") + + if oputil.is_pointclouds(X): + Xt = X.update_padded(Xt) # type: ignore + + return ICPSolution(converged, rmse, Xt, SimilarityTransform(R, T, s), t_history) + + +# threshold for checking that point crosscorelation +# is full rank in corresponding_points_alignment +AMBIGUOUS_ROT_SINGULAR_THR = 1e-15 + + +def corresponding_points_alignment( + X: Union[torch.Tensor, "Pointclouds"], + Y: Union[torch.Tensor, "Pointclouds"], + weights: Union[torch.Tensor, List[torch.Tensor], None] = None, + estimate_scale: bool = False, + allow_reflection: bool = False, + eps: float = 1e-9, +) -> SimilarityTransform: + """ + Finds a similarity transformation (rotation `R`, translation `T` + and optionally scale `s`) between two given sets of corresponding + `d`-dimensional points `X` and `Y` such that: + + `s[i] X[i] R[i] + T[i] = Y[i]`, + + for all batch indexes `i` in the least squares sense. + + The algorithm is also known as Umeyama [1]. + + Args: + **X**: Batch of `d`-dimensional points of shape `(minibatch, num_point, d)` + or a `Pointclouds` object. + **Y**: Batch of `d`-dimensional points of shape `(minibatch, num_point, d)` + or a `Pointclouds` object. + **weights**: Batch of non-negative weights of + shape `(minibatch, num_point)` or list of `minibatch` 1-dimensional + tensors that may have different shapes; in that case, the length of + i-th tensor should be equal to the number of points in X_i and Y_i. + Passing `None` means uniform weights. + **estimate_scale**: If `True`, also estimates a scaling component `s` + of the transformation. Otherwise assumes an identity + scale and returns a tensor of ones. + **allow_reflection**: If `True`, allows the algorithm to return `R` + which is orthonormal but has determinant==-1. + **eps**: A scalar for clamping to avoid dividing by zero. Active for the + code that estimates the output scale `s`. + + Returns: + 3-element named tuple `SimilarityTransform` containing + - **R**: Batch of orthonormal matrices of shape `(minibatch, d, d)`. + - **T**: Batch of translations of shape `(minibatch, d)`. + - **s**: batch of scaling factors of shape `(minibatch, )`. + + References: + [1] Shinji Umeyama: Least-Suqares Estimation of + Transformation Parameters Between Two Point Patterns + """ + + # make sure we convert input Pointclouds structures to tensors + Xt, num_points = oputil.convert_pointclouds_to_tensor(X) + Yt, num_points_Y = oputil.convert_pointclouds_to_tensor(Y) + + if (Xt.shape != Yt.shape) or (num_points != num_points_Y).any(): + raise ValueError( + "Point sets X and Y have to have the same \ + number of batches, points and dimensions." + ) + if weights is not None: + if isinstance(weights, list): + if any(np != w.shape[0] for np, w in zip(num_points, weights)): + raise ValueError( + "number of weights should equal to the " + + "number of points in the point cloud." + ) + weights = [w[..., None] for w in weights] + weights = strutil.list_to_padded(weights)[..., 0] + + if Xt.shape[:2] != weights.shape: + raise ValueError("weights should have the same first two dimensions as X.") + + b, n, dim = Xt.shape + + if (num_points < Xt.shape[1]).any() or (num_points < Yt.shape[1]).any(): + # in case we got Pointclouds as input, mask the unused entries in Xc, Yc + mask = ( + torch.arange(n, dtype=torch.int64, device=Xt.device)[None] + < num_points[:, None] + ).type_as(Xt) + weights = mask if weights is None else mask * weights.type_as(Xt) + + # compute the centroids of the point sets + Xmu = oputil.wmean(Xt, weight=weights, eps=eps) + Ymu = oputil.wmean(Yt, weight=weights, eps=eps) + + # mean-center the point sets + Xc = Xt - Xmu + Yc = Yt - Ymu + + total_weight = torch.clamp(num_points, 1) + # special handling for heterogeneous point clouds and/or input weights + if weights is not None: + Xc *= weights[:, :, None] + Yc *= weights[:, :, None] + total_weight = torch.clamp(weights.sum(1), eps) + + if (num_points < (dim + 1)).any(): + warnings.warn( + "The size of one of the point clouds is <= dim+1. " + + "corresponding_points_alignment cannot return a unique rotation." + ) + + # compute the covariance XYcov between the point sets Xc, Yc + XYcov = torch.bmm(Xc.transpose(2, 1), Yc) + XYcov = XYcov / total_weight[:, None, None] + + # decompose the covariance matrix XYcov + U, S, V = torch.svd(XYcov) + + # catch ambiguous rotation by checking the magnitude of singular values + # if (S.abs() <= AMBIGUOUS_ROT_SINGULAR_THR).any() and not ( + # num_points < (dim + 1) + # ).any(): + # warnings.warn( + # "Excessively low rank of " + # + "cross-correlation between aligned point clouds. " + # + "corresponding_points_alignment cannot return a unique rotation." + # ) + + # identity matrix used for fixing reflections + E = torch.eye(dim, dtype=XYcov.dtype, device=XYcov.device)[None].repeat(b, 1, 1) + + if not allow_reflection: + # reflection test: + # checks whether the estimated rotation has det==1, + # if not, finds the nearest rotation s.t. det==1 by + # flipping the sign of the last singular vector U + R_test = torch.bmm(U, V.transpose(2, 1)) + E[:, -1, -1] = torch.det(R_test) + + # find the rotation matrix by composing U and V again + R = torch.bmm(torch.bmm(U, E), V.transpose(2, 1)) + + if estimate_scale: + # estimate the scaling component of the transformation + trace_ES = (torch.diagonal(E, dim1=1, dim2=2) * S).sum(1) + Xcov = (Xc * Xc).sum((1, 2)) / total_weight + + # the scaling component + s = trace_ES / torch.clamp(Xcov, eps) + + # translation component + T = Ymu[:, 0, :] - s[:, None] * torch.bmm(Xmu, R)[:, 0, :] + else: + # translation component + T = Ymu[:, 0, :] - torch.bmm(Xmu, R)[:, 0, :] + + # unit scaling since we do not estimate scale + s = T.new_ones(b) + + return SimilarityTransform(R, T, s) + + +def _apply_similarity_transform( + X: torch.Tensor, R: torch.Tensor, T: torch.Tensor, s: torch.Tensor +) -> torch.Tensor: + """ + Applies a similarity transformation parametrized with a batch of orthonormal + matrices `R` of shape `(minibatch, d, d)`, a batch of translations `T` + of shape `(minibatch, d)` and a batch of scaling factors `s` + of shape `(minibatch,)` to a given `d`-dimensional cloud `X` + of shape `(minibatch, num_points, d)` + """ + X = s[:, None, None] * torch.bmm(X, R) + T[:, None, :] + return X + +# ----------------------------------------------------------------------------------------------------------- +def setdiff1d(t1, t2): + # indices = torch.new_ones(t1.shape, dtype = torch.uint8) + # for elem in t2: + # indices = indices & (t1 != elem) + # intersection = t1[indices] + # assuming t2 is a subset of t1 + t1_unique = torch.unique(t1) + t2_unique = torch.unique(t2) + assert len(t1_unique)>=len(t2_unique) + t12, counts = torch.cat([t1_unique, t2_unique]).unique(return_counts=True) + diff = t12[torch.where(counts.eq(1))] + return diff + +def get_bbox_tensor(points): + x = torch.abs(points[:, 0].max() - points[:, 0].min()) + y = torch.abs(points[:, 1].max() - points[:, 1].min()) + z = torch.abs(points[:, 2].max() - points[:, 2].min()) + return sorted([x, y, z]) + +def sanity_check(args, src_points, dst_points, src_labels, dst_labels, pairs): + pairs_true = [] + for pair in pairs: + src = src_points[src_labels==pair[0]] + dst = dst_points[dst_labels==pair[1]] + # print('sanity check :', pair, len(src), len(dst), src.mean(0), dst.mean(0), torch.linalg.norm(dst.mean(0) - src.mean(0)), args.translation_frame) + + # scenario 1: either src or dst does not exist, return None + # scenario 2: both src or dst exist, but they are not matchable because of ground points/too few points/size mismatch, return False + # scenario 3: both src or dst exist, and they are are matchable, return True + if min(len(src), len(dst))args.translation_frame: continue # x/y translation + + src_bbox = get_bbox_tensor(src) + dst_bbox = get_bbox_tensor(dst) + # print('sanity check bbox:', src_bbox, dst_bbox) + if min(src_bbox[0], dst_bbox[0]) < args.thres_box * max(src_bbox[0], dst_bbox[0]): continue + if min(src_bbox[1], dst_bbox[1]) < args.thres_box * max(src_bbox[1], dst_bbox[1]): continue + if min(src_bbox[2], dst_bbox[2]) < args.thres_box * max(src_bbox[2], dst_bbox[2]): continue + + pairs_true.append(pair) + if len(pairs_true)>0: + return torch.vstack(pairs_true) + else: + return torch.zeros((0,2)) + +def match_pcds(args, src_points, dst_points, src_labels, dst_labels): + + src_labels_unq = torch.unique(src_labels, return_counts=False).long() + dst_labels_unq = torch.unique(dst_labels, return_counts=False).long() + labels_unq = torch.unique(torch.cat([src_labels_unq, dst_labels_unq], axis=0), return_counts=False) + + # # # stage 1: match static: overlapped clusters + pairs = torch.stack([labels_unq, labels_unq], dim=1) + mask = pairs.min(dim=1)[0]>=0 # remove ground + pairs = pairs[mask] + pairs_true = sanity_check(args, src_points, dst_points, src_labels, dst_labels, pairs) + # print('sanity check: sta: ', len(pairs), len(pairs_true), pairs_true) + + if len(pairs_true)>0: + pairs_sta, transformations_sta = match_pairs(args, src_points, dst_points, src_labels, dst_labels, pairs_true) + else: + pairs_sta, transformations_sta = torch.tensor([]).reshape(0, 10), torch.tensor([]).reshape(0, 4, 4) + # print('pairs_sta: ', len(pairs), len(pairs_sta), pairs_sta[:, 0:2]) + + # stage 2: match dynamic + # # # remove matched near-static pairs: + if len(pairs_sta)0: + src_labels_unq = setdiff1d(src_labels_unq, pairs_sta[:, 0]) + dst_labels_unq = setdiff1d(dst_labels_unq, pairs_sta[:, 1]) + + pairs = torch.stack([src_labels_unq.repeat_interleave(len(dst_labels_unq)), dst_labels_unq.repeat(len(src_labels_unq))], dim=1) + pairs_true = sanity_check(args, src_points, dst_points, src_labels, dst_labels, pairs) + else: + pairs_true = torch.zeros(0, 2) + # print('dynamic src_labels, dst_labels: ', src_labels_unq.long(), dst_labels_unq.long(), pairs_true) + + if len(pairs_true)>0: + pairs_dyn, transformations_dyn = match_pairs(args, src_points, dst_points, src_labels, dst_labels, pairs_true) + # print('dynamic paired_idxs: ', len(pairs), len(pairs_true), len(pairs_dyn), pairs_true, pairs_dyn) + else: + pairs_dyn, transformations_dyn = torch.tensor([]).reshape(0, 10), torch.tensor([]).reshape(0, 4, 4) + + pairs_matched = torch.cat([pairs_sta, pairs_dyn], dim=0) + transformations_matched = torch.cat([transformations_sta, transformations_dyn], dim=0) + # assert len(pairs_matched)>0 # likely to be bugs or outliers + + return pairs_matched, transformations_matched + +def track(args, point_src, point_dst, label_src, label_dst): + pairs, transformations = match_pcds(args, point_src, point_dst, label_src, label_dst) + # print(f'match_pcds pairs: {pairs}, {transformations}') + # print(f'match_pcds pairs: {torch.round(pairs[:, 0:2], decimals=2)}') + return pairs, transformations + + +def random_choice(m, n): + assert m>=n + perm = torch.randperm(m) + return perm[0:n] + +def pad_segment(seg, max_points): + padding = seg.new_zeros((max_points, 1)) + 1.0 + if len(seg) > max_points: + sample_idxs = random_choice(len(seg), max_points) + seg = seg[sample_idxs, :] + elif len(seg) < max_points: + padding[len(seg):] = 0.0 + seg = torch.cat([seg, seg.new_zeros((max_points-len(seg), 3))+1e8], dim=0) + else: + pass + assert len(seg)==max_points + return torch.cat([seg, padding], axis=1) + +# tok=1 already works decently. topk=5 is for ablation study and works slightly better +def topk_nms(x, k=5, kernel_size=11): + b, h, w, d = x.shape + x = x.unsqueeze(1) + xp = torch.nn.functional.max_pool3d(x, kernel_size=kernel_size, stride=1, padding=(kernel_size-1)//2) + mask = (x == xp).float().clamp(min=0.0) + xp = x * mask + votes, idxs = torch.topk(xp.view(b, -1), dim=1, k=k) + del xp, mask + return votes, idxs.long() + +def transform_points_batch(xyz, pose): + assert xyz.dim()==3 + assert pose.dim()==3 + assert xyz.shape[2]==4 + assert pose.shape[1]==4 + assert pose.shape[2]==4 + assert len(xyz)==len(pose) + b, n, _ = xyz.shape + # print('transform points batch: ', xyz.shape, pose.shape) + xyzh = torch.cat([xyz[:, :, 0:3], xyz.new_ones((b, n, 1))], dim=-1) + xyzh_tmp = torch.bmm(xyzh, pose.permute(0, 2, 1)) + return torch.cat([xyzh_tmp[:, :, 0:3], xyz[:, :, -1:]], dim=-1) + +def pytorch3d_icp(args, src, dst): + icp_result = iterative_closest_point(src, dst, + init_transform=None, + thres=args.thres_dist, + max_iterations=100, + relative_rmse_thr=1e-6, + estimate_scale=False, + allow_reflection=False, + verbose=False) + + Rs = icp_result.RTs.R + ts = icp_result.RTs.T + + Rts = torch.cat([Rs, ts[:, None, :]], dim=1) + Rts = torch.cat([Rts.permute(0, 2, 1), Rts.new_zeros(len(ts), 1, 4)], dim=1) + Rts[:, 3, 3]=1.0 + + # print('pytorch3d icp Rt: ', Rt) + # src_tmp = transform_points_tensor(src, Rt) + # visualize_pcd(np.concatenate([src_tmp.cpu().numpy(), dst.cpu().numpy()], axis=0), + # np.concatenate([np.zeros((len(src)))+1, np.zeros((len(dst)))+2], axis=0), + # num_colors=3, + # title=f'registration debug, size: {len(src)} vs {len(dst)}') + return Rts + +def apply_icp(args, src, dst, init_poses): + src_tmp = transform_points_batch(src, init_poses) + + Rts = pytorch3d_icp(args, src_tmp, dst) + Rts = torch.bmm(Rts, init_poses) + + # # # pytorch 3d icp might go wrong ! to fix! + mask_src = src[:, : , -1] > 0.0 + _, error_init = nearest_neighbor_batch(src_tmp, dst) + error_init = (error_init * mask_src).sum(dim=1) / mask_src.sum(dim=1) + + src_tmp = transform_points_batch(src, Rts) + _, error_icp = nearest_neighbor_batch(src_tmp, dst) + error_icp = (error_icp * mask_src).sum(dim=1) / mask_src.sum(dim=1) + invalid = error_icp>=error_init + Rts[invalid] = init_poses[invalid] + + return Rts + +from assets.cuda.histlib import histf +def estimate_init_pose_batch(args, src, dst): + pcd1 = src[:, :, 0:3] + pcd2 = dst[:, :, 0:3] + mask1 = src[:, : , -1] > 0.0 + mask2 = dst[:, : , -1] > 0.0 + + ########################################################################################### + eps = 1e-8 + # https://pytorch.org/docs/stable/generated/torch.arange.html#torch-arange + bins_x = torch.arange(-args.translation_frame, args.translation_frame+args.thres_dist-eps, args.thres_dist) + bins_y = torch.arange(-args.translation_frame, args.translation_frame+args.thres_dist-eps, args.thres_dist) + bins_z = torch.arange(-args.thres_dist, args.thres_dist+args.thres_dist-eps, args.thres_dist) + # print(f'bins: {bins_x.min()} {bins_x.max()} {bins_x.shape}, {bins_z.min()} {bins_z.max()} {bins_z}') + + # bug there: when batch size is large! + t_hist = histf(dst, src, + bins_x.min(), bins_y.min(), bins_z.min(), + bins_x.max(), bins_y.max(), bins_z.max(), + len(bins_x), len(bins_y), len(bins_z)) + b, h, w, d = t_hist.shape + # print(f't_hist.shape: {t_hist.shape} {bins_x.shape} {bins_y.shape} {bins_z.shape} {t_hist.max()}') + ########################################################################################### + + t_maxs, t_argmaxs = topk_nms(t_hist) + t_dynamic = torch.stack([ bins_x[t_argmaxs//d//w%h], bins_y[t_argmaxs//d%w], bins_z[t_argmaxs%d] ], dim=-1) + args.thres_dist//2 + # print(f't_dynamic: {t_dynamic.shape}, {t_maxs} {h}, {w}, {d}, {t_dynamic}') + del t_hist, bins_x, bins_y, bins_z + + n = pcd1.shape[1] + t_both = torch.cat([t_dynamic, t_dynamic.new_zeros(b, 1, 3)], dim=1) + k = t_both.shape[1] + + pcd1_ = pcd1[:, None, :, :] + t_both[:, :, None, :] + pcd2_ = pcd2[:, None, :, :].expand(-1, k, -1, -1) + + _, errors = nearest_neighbor_batch(pcd1_.reshape(b*k, n, 3), pcd2_.reshape(b*k, n, 3)) + _, errors_inv = nearest_neighbor_batch(pcd2_.reshape(b*k, n, 3), pcd1_.reshape(b*k, n, 3)) + + # using errors + errors = (errors.view(b, k, n) * mask1[:, None, :]).sum(dim=-1) / mask1[:, None, :].sum(dim=-1) + errors_inv = (errors_inv.view(b, k, n) * mask2[:, None, :]).sum(dim=-1) / mask2[:, None, :].sum(dim=-1) + errors = torch.minimum(errors, errors_inv) + error, idx = errors.min(dim=-1) + error_best = error + t_best = t_both[torch.arange(0, b, device=idx.device), idx, :] + del pcd1_, pcd2_, errors + + transformation = torch.eye(4)[None].repeat(b, 1, 1) + transformation[:, 0:3, -1] = t_best + return transformation + +def estimate_init_pose(args, src, dst): + transformations = [] + assert len(src) == len(dst) + n = len(src)//args.chunk_size if len(src) % args.chunk_size ==0 else len(src)//args.chunk_size +1 + for k in range(0, n): + transformation = estimate_init_pose_batch(args, + src[k*args.chunk_size: (k+1)*args.chunk_size], + dst[k*args.chunk_size: (k+1)*args.chunk_size]) + transformations.append(transformation) + transformations = torch.vstack(transformations) + return transformations + +def match_eval(args, pcd1, pcd2, transformations): + pcd1_tmp = transform_points_batch(pcd1, transformations) + src = pcd1_tmp + dst = pcd2 + src_mask = pcd1[:, :, -1]>0.0 + dst_mask = pcd2[:, :, -1]>0.0 + src_dst_idxs, src_error = nearest_neighbor_batch(src, dst) + dst_src_idxs, dst_error = nearest_neighbor_batch(dst, src) + + src_inlier = torch.logical_and(src_error < args.thres_dist, src_mask).float() + dst_inlier = torch.logical_and(dst_error < args.thres_dist, dst_mask).float() + + src_ratio = torch.sum(src_inlier, dim=1) / torch.sum(src_mask, dim=1) + dst_ratio = torch.sum(dst_inlier, dim=1) / torch.sum(dst_mask, dim=1) + + src_iou = torch.sum(src_inlier, dim=1) / (torch.sum(src_mask, dim=1) + torch.sum(dst_mask, dim=1) - torch.sum(dst_inlier, dim=1)) + dst_iou = torch.sum(dst_inlier, dim=1) / (torch.sum(src_mask, dim=1) + torch.sum(dst_mask, dim=1) - torch.sum(src_inlier, dim=1)) + + src_error = (src_error * src_mask).sum(1) / src_mask.sum(1) + dst_error = (dst_error * dst_mask).sum(1) / dst_mask.sum(1) + + src_mean = (src[:, :, 0:3] * src_mask[:, :, None]).sum(dim=1) / src_mask.sum(dim=1, keepdim=True) + src_ori_mean = (pcd1[:, :, 0:3] * src_mask[:, :, None]).sum(dim=1) / src_mask.sum(dim=1, keepdim=True) + # dst_mean = (dst * dst_mask).sum(dim=1) / dst_mask.sum(dim=1, keepdim=True) + translations = src_mean - src_ori_mean + rotations = pytorch3d_t.matrix_to_euler_angles(transformations[:, 0:3, 0:3], convention='ZYX') * 180./np.pi + + return torch.stack([src_error, dst_error], dim=1), \ + torch.stack([src_inlier.sum(1), dst_inlier.sum(1)], dim=1), \ + torch.stack([src_ratio, dst_ratio], dim=1), \ + torch.stack([src_iou, dst_iou], dim=1), \ + translations, rotations + +def hist_icp(args, src, dst): + mask1 = src[:, :, -1]>0.0 + mask2 = dst[:, :, -1]>0.0 + # alwyas match the smaller one to the larger one + idxs = mask1.sum(dim=1)>mask2.sum(dim=1) + src_ = src.clone() + dst_ = dst.clone() + src_[idxs] = dst[idxs] + dst_[idxs] = src[idxs] + + with torch.no_grad(): + init_poses_ = estimate_init_pose(args, src_, dst_) + transformations_ = apply_icp(args, src_, dst_, init_poses_) + + if sum(idxs)>0: + transformations = transformations_.clone() + transformations[idxs] = torch.linalg.inv(transformations_[idxs]) + else: + transformations = transformations_ + return transformations + +def check_transformation(args, translation, rotation, iou): + # print('check transformation: ', translation, rotation, iou, args.translation_frame) + # # # check translation + if torch.linalg.norm(translation) > args.translation_frame: + return False + + # # # check iou + if ioumax_rot: # roll and pitch + return False + + return True + + +def match_segments_descend(matrix_metric): + src_idxs = torch.arange(0, len(matrix_metric)) + dst_idxs = torch.argmin(matrix_metric, dim=1) + return src_idxs, dst_idxs + +def match_pairs(args, src_points, dst_points, src_labels, dst_labels, pairs): + src_labels_unq = torch.unique(src_labels) + dst_labels_unq = torch.unique(dst_labels) + matrix_errors = torch.zeros((len(src_labels_unq), len(dst_labels_unq), 2)) + 1e8 + matrix_inliers = torch.zeros((len(src_labels_unq), len(dst_labels_unq), 2)) + 0.0 + matrix_ratios = torch.zeros((len(src_labels_unq), len(dst_labels_unq), 2)) + 0.0 + matrix_ious = torch.zeros((len(src_labels_unq), len(dst_labels_unq), 2)) + 0.0 + matrix_transformations = torch.zeros((len(src_labels_unq), len(dst_labels_unq), 4, 4)) + + assert len(pairs)>0 + segs_src = [] + segs_dst = [] + for pair in pairs: + src = src_points[src_labels==pair[0], 0:3] + dst = dst_points[dst_labels==pair[1], 0:3] + src = pad_segment(src, args.max_points) + dst = pad_segment(dst, args.max_points) + # always match the smaller one to the larger one + segs_src.append(src) + segs_dst.append(dst) + + segs_src = torch.stack(segs_src, dim=0) + segs_dst = torch.stack(segs_dst, dim=0) + transformations = hist_icp(args, segs_src, segs_dst) + errors, inliers, ratios, ious, translations, rotations = match_eval(args, segs_src, segs_dst, transformations) + # reject unreliable matches + num_matches = 0 + for k, (pair, error, inlier, ratio, iou, translation, rotation, transformation) in \ + enumerate(zip(pairs, errors, inliers, ratios, ious, translations, rotations, transformations)): + # print('check per pair: ', pair, error, inlier, ratio, iou, translation, rotation, args.translation_max ) + if not check_transformation(args, translation, rotation, min(iou)): + continue + src_idx = torch.nonzero(src_labels_unq == pair[0]) + dst_idx = torch.nonzero(dst_labels_unq == pair[1]) + matrix_errors[src_idx, dst_idx, :] = error + matrix_inliers[src_idx, dst_idx, :] = inlier + matrix_ratios[src_idx, dst_idx, :] = ratio + matrix_ious[src_idx, dst_idx, :] = iou + matrix_transformations[src_idx, dst_idx] = transformation + num_matches += 1 + + if num_matches>0: + matrix_errors_min, _ = matrix_errors.min(-1) + src_idxs, dst_idxs = match_segments_descend(matrix_errors_min) + valid = matrix_errors_min[src_idxs, dst_idxs] < args.thres_error + src_idxs = src_idxs[valid] + dst_idxs = dst_idxs[valid] + + # matrix_ious_min, _ = matrix_ious.min(-1) + # src_idxs, dst_idxs = match_segments_ascend(matrix_ious_min) + # valid = matrix_ious_min[src_idxs, dst_idxs] >= args.thres_iou + # src_idxs = src_idxs[valid] + # dst_idxs = dst_idxs[valid] + + pairs = torch.cat([src_labels_unq[src_idxs][:, None], dst_labels_unq[dst_idxs][:, None], + matrix_errors[src_idxs, dst_idxs], + matrix_inliers[src_idxs, dst_idxs], + matrix_ratios[src_idxs, dst_idxs], + matrix_ious[src_idxs, dst_idxs]], + axis=1) + + transformations = matrix_transformations[src_idxs, dst_idxs] + + else: + pairs = torch.tensor([]).reshape(0, 2+matrix_errors.shape[-1]+matrix_inliers.shape[-1]+matrix_ratios.shape[-1]+matrix_ious.shape[-1]) + transformations = torch.tensor([]).reshape(0, 4, 4) + + return pairs, transformations + + diff --git a/src/models/icpflow.py b/src/models/icpflow.py new file mode 100644 index 0000000..95c110f --- /dev/null +++ b/src/models/icpflow.py @@ -0,0 +1,187 @@ + +""" +This file is from: https://github.com/yanconglin/ICP-Flow/blob/main/utils_flow.py +with slightly modification to have unified format with all benchmark. +""" + +import dztimer, torch +import torch.nn as nn +import numpy as np +import hdbscan +import torch + +from .basic import cal_pose0to1 +from .basic.icpflow_lib import track + +def flow_estimation(src_points, dst_points, + src_labels, dst_labels, + pairs, transformations, pose): + unqs = np.unique(src_labels.astype(int)) + # sanity check: src labels include -1e8 and -1; segment labels start from 0. + # -1e8: ground, + # -1: unmatched/unclustered segments in non-ground points + # assert len(unqs)-2 == max(src_label)+1 + # print('unqs: ', len(unqs), unqs) + + assert len(src_points) == len(src_labels) + flow = np.zeros((len(src_points), 3)) + for unq in unqs: + idxs = src_labels==unq + xyz = src_points[src_labels==unq, 0:3] + if unq in pairs[:, 0]: + idx = np.flatnonzero(unq==pairs[:,0]) + assert len(idx)==1 + idx = idx.item() + transformation = transformations[idx] + else: + transformation = np.eye(4) + + flow_i = calculate_flow_rigid(xyz, transformation @ pose) + flow[idxs] = flow_i + + return flow + +def calculate_flow_rigid(points, transformation): + points_tmp = transform_points(points, transformation) + flow = points_tmp - points + return flow + +def transform_points(xyz, pose): + assert xyz.shape[1]==3 + xyzh = np.concatenate([xyz, np.ones((len(xyz), 1))], axis=1) + xyzh_tmp = xyzh @ pose.T + return xyzh_tmp[:, 0:3] + +def cluster_pcd(points, min_cluster_size=20, num_clusters=200): + + clusterer = hdbscan.HDBSCAN(algorithm='best', alpha=1., approx_min_span_tree=True, + gen_min_span_tree=True, leaf_size=100, + metric='euclidean', min_cluster_size=min_cluster_size, min_samples=None + ) + clusterer.fit(points[:, 0:3]) + labels = clusterer.labels_.copy() + lbls, counts = np.unique(labels, return_counts=True) + cluster_info = np.array(list(zip(lbls[1:], counts[1:]))) + cluster_info = cluster_info[cluster_info[:,1].argsort()] + + clusters_labels = cluster_info[::-1][:num_clusters, 0] + labels[np.in1d(labels, clusters_labels, invert=True)] = -1 # unclustered points + + return labels + +def cluster_labels(pose0, pose1, pc0_wo_ground, pc1_wo_ground): + + point_src = pc0_wo_ground + point_dst = pc1_wo_ground + + pose = cal_pose0to1(pose0, pose1) + point_src_ego = point_src @ pose[:3, :3].T + pose[:3, 3] + points_tmp = np.concatenate([point_dst.cpu().numpy(), point_src_ego.cpu().numpy()], axis=0) + + label_tmp = cluster_pcd(points_tmp) + label_src = label_tmp[len(point_dst):] + label_dst = label_tmp[0:len(point_dst)] + + return point_src_ego, point_dst, label_src, label_dst, pose.cpu().numpy() + +from dataclasses import dataclass +@dataclass +class dataargs: + thres_dist: float = 0.1 + translation_frame: float = 1.67 + chunk_size: int = 50 + thres_iou: float = 0.2 + max_points: int = 10000 + min_cluster_size: int = 20 + thres_box: float = 0.1 + thres_rot: float = 0.1 + thres_error: float = 0.2 + speed: float = 1.67 + +class ICPFlow(nn.Module): + def __init__(self, thres_dist: float = 0.1, + translation_frame: float = 1.67, + chunk_size: int = 50, + thres_iou: float = 0.2, + max_points: int = 10000, + min_cluster_size: int = 20, + thres_box: float = 0.1, + thres_rot: float = 0.1, + thres_error: float = 0.2, + speed: float = 1.67, + point_cloud_range = [-51.2, -51.2, -3, 51.2, 51.2, 3] + ): + super().__init__() + self.args = dataargs(thres_dist, translation_frame, chunk_size, + thres_iou, max_points, min_cluster_size, thres_box, + thres_rot, thres_error, speed) + self.timer = dztimer.Timing() + self.timer.start("NSFP Model Inference") + self.point_cloud_range = point_cloud_range + print(f"\n---LOG[model]: ICPFlow setup successfully, hyperparameters Settings: thres_dist={thres_dist}, min_cluster_size={min_cluster_size}, speed={speed}.") + + def range_limit_(self, pc): + """ + Limit the point cloud to the given range. + """ + mask = (pc[:, 0] >= self.point_cloud_range[0]) & (pc[:, 0] <= self.point_cloud_range[3]) & \ + (pc[:, 1] >= self.point_cloud_range[1]) & (pc[:, 1] <= self.point_cloud_range[4]) & \ + (pc[:, 2] >= self.point_cloud_range[2]) & (pc[:, 2] <= self.point_cloud_range[5]) + return pc[mask], mask + + def forward(self, batch): + batch_sizes = len(batch["pose0"]) + device = batch['pose0'][0].device + pose_flows = [] + batch_final_flow = [] + for batch_id in range(batch_sizes): + self.timer[0].start("Data Processing") + pc0 = batch["pc0"][batch_id] + pc1 = batch["pc1"][batch_id] + selected_pc0, rm0 = self.range_limit_(pc0) + selected_pc1, rm1 = self.range_limit_(pc1) + self.timer[0][0].start("pose") + 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) + + self.timer[0].stop() + + point_src, point_dst, label_src, label_dst, pose = cluster_labels(batch["pose0"][batch_id], batch["pose1"][batch_id], selected_pc0, selected_pc1) + translation_frame = max(self.args.speed, np.linalg.norm(pose[0:3, -1])) * 2 + + with torch.inference_mode(True): + with torch.device(device): + point_src = point_src.float().to(device) + point_dst = point_dst.float().to(device) + pairs, transformations = track(self.args, point_src, point_dst, torch.from_numpy(label_src).float().to(device), \ + torch.from_numpy(label_dst).float().to(device)) + pairs = pairs.cpu().numpy() + transformations = transformations.cpu().numpy() + + # assigning the same labels to corresponding instances; useful for visualization + # tracker_src, tracker_dst = trackers2labels(label_src, labels_dst[0], pairs) + point_src = selected_pc0.cpu().numpy() + point_dst = selected_pc1.cpu().numpy() + + flow = flow_estimation(src_points=point_src, dst_points=point_dst, + src_labels=label_src, dst_labels=label_dst, + pairs=pairs, transformations=transformations, pose=pose) + + final_flow = torch.zeros_like(pc0) + final_flow[rm0] = torch.from_numpy(flow).float().to(device) - pose_flows[batch_id] + batch_final_flow.append(final_flow) + + res_dict = {"flow": batch_final_flow, + "pose_flow": pose_flows + } + + return res_dict \ No newline at end of file diff --git a/src/runner.py b/src/runner.py index 98d5644..04077fd 100644 --- a/src/runner.py +++ b/src/runner.py @@ -300,6 +300,7 @@ def _run_process(cfg, mode): def _spawn_wrapper(rank, world_size, cfg, mode): torch.cuda.set_device(rank) + # FIXME(Qingwen): better to set these through command, since we might have more nodes to connected. os.environ['RANK'] = str(rank) os.environ['WORLD_SIZE'] = str(world_size) os.environ['MASTER_ADDR'] = 'localhost' From 8d99f5a7d972bad0d2efdb0e8dcc96afa581ad28 Mon Sep 17 00:00:00 2001 From: Qingwen Zhang <35365764+Kin-Zhang@users.noreply.github.com> Date: Fri, 12 Sep 2025 23:04:10 +0200 Subject: [PATCH 5/9] hotfix(runner): Refactor metrics accumulation to use speed buckets --- src/runner.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/runner.py b/src/runner.py index 8a6414c..e0a5321 100644 --- a/src/runner.py +++ b/src/runner.py @@ -267,13 +267,13 @@ def _run_process(cfg, mode): final_metrics.epe_3way[key].extend(val_list) for class_idx, class_name in enumerate(metrics_obj.bucketedMatrix.class_names): - for range_idx, range_bucket in enumerate(metrics_obj.bucketedMatrix.range_buckets): - count = metrics_obj.bucketedMatrix.count_storage_matrix[class_idx, range_idx] + for speed_idx, speed_bucket in enumerate(metrics_obj.bucketedMatrix.speed_buckets): + count = metrics_obj.bucketedMatrix.count_storage_matrix[class_idx, speed_idx] if count > 0: - avg_epe = metrics_obj.bucketedMatrix.epe_storage_matrix[class_idx, range_idx] - avg_range = metrics_obj.bucketedMatrix.range_storage_matrix[class_idx, range_idx] + avg_epe = metrics_obj.bucketedMatrix.epe_storage_matrix[class_idx, speed_idx] + avg_speed = metrics_obj.bucketedMatrix.speed_storage_matrix[class_idx, speed_idx] final_metrics.bucketedMatrix.accumulate_value( - class_name, range_bucket, avg_epe, avg_range, count + class_name, speed_bucket, avg_epe, avg_speed, count ) for class_idx, class_name in enumerate(metrics_obj.distanceMatrix.class_names): for range_idx, range_bucket in enumerate(metrics_obj.distanceMatrix.range_buckets): From dda21fd8b6ff45c395373af5c829f13541e1634d Mon Sep 17 00:00:00 2001 From: Qingwen Zhang <35365764+Kin-Zhang@users.noreply.github.com> Date: Sun, 14 Sep 2025 19:16:37 +0100 Subject: [PATCH 6/9] Improve ICPFlow import error message --- src/models/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/models/__init__.py b/src/models/__init__.py index 3eff0ca..5288e5b 100644 --- a/src/models/__init__.py +++ b/src/models/__init__.py @@ -48,6 +48,6 @@ # * pip install pytorch3d assets/cuda/histlib try: from .icpflow import ICPFlow -except ImportError: +except ImportError as e: print("--- WARNING [model]: ICPFlow is not imported, as it requires pytorch3d lib which is not installed.") print(f"Detail error message\033[0m: {e}. Just ignore this warning if code runs without these models.") From 658876121239bc0d63e65847c04a46d779921619 Mon Sep 17 00:00:00 2001 From: Kin Date: Fri, 21 Nov 2025 10:25:52 +0100 Subject: [PATCH 7/9] docs: move to reproduce As the official review didn't finish, will open another request once Yancong have time to review/improve. --- README.md | 26 ++++++-------------------- 1 file changed, 6 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index 8f7b3fb..1d424df 100644 --- a/README.md +++ b/README.md @@ -41,10 +41,6 @@ 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) -- **ICP-Flow: LiDAR Scene Flow Estimation with ICP** -*Yancong Lin, Holger Caesar* -Conference on Computer Vision and Pattern Recognition (**CVPR**) 2024 -[ Optimization-based ] [ Self-Supervised ] - [ [arXiv](https://arxiv.org/abs/2402.17351) ] [ [Project](https://github.com/yanconglin/ICP-Flow) ] → [here](#icp-flow) - **DeFlow: Decoder of Scene Flow Network in Autonomous Driving** *Qingwen Zhang, Yi Yang, Heng Fang, Ruoyu Geng, Patric Jensfelt* @@ -59,7 +55,8 @@ Additionally, *OpenSceneFlow* integrates following excellent works: [ICLR'24 Zer - [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. +- [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. @@ -232,17 +229,12 @@ wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/seflowpp_best.c 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. -#### ICP-Flow - -Extra packages needed for ICP-Flow model: ```bash -pip install pytorch3d assets/cuda/histlib -``` +# you can change another model by passing model name. +python eval.py model=fastnsf -Then run as: -```bash -python eval.py model=icpflow -python save.py model=icpflow +# or save the result directly +python save.py model=fastnsf ``` @@ -394,12 +386,6 @@ And our excellent collaborators works contributed to this codebase also: booktitle={CVPR}, year={2025}, } -@article{lin2024icp, - title={ICP-Flow: LiDAR Scene Flow Estimation with ICP}, - author={Lin, Yancong and Caesar, Holger}, - booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)}, - year={2024} -} ``` Thank you for your support! โค๏ธ From 0ff427d68de559e3ddacf7ec35c84b7764053611 Mon Sep 17 00:00:00 2001 From: Kin Date: Fri, 21 Nov 2025 10:33:45 +0100 Subject: [PATCH 8/9] fix(runner): val mode add valid to align waymo validation run. * update master port also for multi-program run. --- README.md | 4 ++-- src/runner.py | 13 ++++++------- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 1d424df..8168458 100644 --- a/README.md +++ b/README.md @@ -48,7 +48,7 @@ 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) ๐ŸŽ One repository, All methods! -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...)
Summary of them: @@ -227,7 +227,7 @@ wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/seflowpp_best.c ### 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 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. diff --git a/src/runner.py b/src/runner.py index e0a5321..518f3c6 100644 --- a/src/runner.py +++ b/src/runner.py @@ -101,11 +101,12 @@ def __init__(self, cfg, rank, world_size, mode): self.mode = mode self.model.to(self.device) - self.metrics = OfficialMetrics() if self.mode in ['val', 'eval'] else None + self.metrics = OfficialMetrics() if self.mode in ['val', 'eval', 'valid'] else None + self.res_name = cfg.get('res_name', cfg.model.name) self.save_res_path = cfg.get('save_res_path', None) def _setup_dataloader(self): - if self.mode in ['val', 'test', 'eval']: + if self.mode in ['val', 'test', 'eval', 'valid']: dataset_path = self.cfg.dataset_path + f"/{self.cfg.data_mode}" is_eval_mode = True else: # 'save' @@ -153,7 +154,7 @@ def _process_step(self, batch): final_flow = pose_flow.clone() final_flow[~batch['gm0']] = res_dict['flow'] + pose_flow[~batch['gm0']] - if self.mode in ['val', 'eval']: + if self.mode in ['val', 'eval', 'valid']: eval_mask = batch['eval_mask'].squeeze() gt_flow = batch["flow"] v1_dict = evaluate_leaderboard(final_flow[eval_mask], pose_flow[eval_mask], pc0[eval_mask], \ @@ -257,7 +258,7 @@ def _run_process(cfg, mode): gathered_metrics_objects = [runner.metrics] if rank == 0: - if mode in ['val', 'eval']: + if mode in ['val', 'eval', 'valid']: final_metrics = OfficialMetrics() print(f"\n--- [LOG] Finished processing. Aggregating results from {world_size} GPUs with {len(gathered_metrics_objects)} metrics objects...") for metrics_obj in gathered_metrics_objects: @@ -299,12 +300,10 @@ def _run_process(cfg, mode): def _spawn_wrapper(rank, world_size, cfg, mode): torch.cuda.set_device(rank) - - # FIXME(Qingwen): better to set these through command, since we might have more nodes to connected. os.environ['RANK'] = str(rank) os.environ['WORLD_SIZE'] = str(world_size) os.environ['MASTER_ADDR'] = 'localhost' - os.environ['MASTER_PORT'] = cfg.get('master_port', '12355') + os.environ['MASTER_PORT'] = str(cfg.get('master_port', 12355)) _run_process(cfg, mode) def launch_runner(cfg, mode): From 65d2f7cb5627c0f6c131ef6216e7a088e59480c9 Mon Sep 17 00:00:00 2001 From: Kin Date: Fri, 21 Nov 2025 12:48:41 +0100 Subject: [PATCH 9/9] revert: check the range_bucket is range of speed/distance. revert to previous fix. successfully run on all opt method. --- src/runner.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/runner.py b/src/runner.py index 518f3c6..d2cda83 100644 --- a/src/runner.py +++ b/src/runner.py @@ -268,13 +268,13 @@ def _run_process(cfg, mode): final_metrics.epe_3way[key].extend(val_list) for class_idx, class_name in enumerate(metrics_obj.bucketedMatrix.class_names): - for speed_idx, speed_bucket in enumerate(metrics_obj.bucketedMatrix.speed_buckets): - count = metrics_obj.bucketedMatrix.count_storage_matrix[class_idx, speed_idx] + for range_idx, range_bucket in enumerate(metrics_obj.bucketedMatrix.range_buckets): + count = metrics_obj.bucketedMatrix.count_storage_matrix[class_idx, range_idx] if count > 0: - avg_epe = metrics_obj.bucketedMatrix.epe_storage_matrix[class_idx, speed_idx] - avg_speed = metrics_obj.bucketedMatrix.speed_storage_matrix[class_idx, speed_idx] + avg_epe = metrics_obj.bucketedMatrix.epe_storage_matrix[class_idx, range_idx] + avg_range = metrics_obj.bucketedMatrix.range_storage_matrix[class_idx, range_idx] final_metrics.bucketedMatrix.accumulate_value( - class_name, speed_bucket, avg_epe, avg_speed, count + class_name, range_bucket, avg_epe, avg_range, count ) for class_idx, class_name in enumerate(metrics_obj.distanceMatrix.class_names): for range_idx, range_bucket in enumerate(metrics_obj.distanceMatrix.range_buckets):