Merge Flow4D model into codebase (#1)

* feat(flow4d): update flow4d model.

docs(conf): update flow4d README and conf

* Add Flow4D details to README.md

* docs: add flow4d bib also.

* fix(trainer): validation setp with res may None, fix with if.

* docs(dockerfile): update dockerfile for convenient env setup.

* update README and delete useless info
* update bib on the end of readme but with link to jump at the beginning.

* docs(REDME): update print message

and ignore warning info for np in eval script.

---------

Co-authored-by: jykim94 <89293559+jykim94@users.noreply.github.com>

Author: Kin-Zhang

Dockerfile

@@ -2,22 +2,16 @@
 FROM nvidia/cuda:11.7.1-devel-ubuntu20.04
 ENV DEBIAN_FRONTEND noninteractive
 
-RUN apt update && apt install -y --no-install-recommends \
-    git curl vim rsync htop
+RUN apt update && apt install -y git curl vim rsync htop
 
-RUN curl -o ~/miniconda.sh -LO https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh  && \
-    chmod +x ~/miniconda.sh && \
-    ~/miniconda.sh -b -p /opt/conda && \
-    rm ~/miniconda.sh && \
+RUN curl -o ~/miniforge3.sh -LO https://github.com/conda-forge/miniforge/releases/latest/download/miniforge3-Linux-x86_64.sh  && \
+    chmod +x ~/miniforge3.sh && \
+    ~/miniforge3.sh -b -p /opt/conda && \
+    rm ~/miniforge3.sh && \
     /opt/conda/bin/conda clean -ya && /opt/conda/bin/conda init bash
 
-RUN curl -o ~/mamba.sh -LO https://github.com/conda-forge/miniforge/releases/latest/download/Mambaforge-Linux-x86_64.sh && \
-    chmod +x ~/mamba.sh && \
-    ~/mamba.sh -b -p /opt/mambaforge && \
-    rm ~/mamba.sh && /opt/mambaforge/bin/mamba init bash
-
 # install zsh and oh-my-zsh
-RUN apt install -y wget git zsh tmux vim g++
+RUN apt update && apt install -y wget git zsh tmux vim g++
 RUN sh -c "$(wget -O- https://github.com/deluan/zsh-in-docker/releases/download/v1.1.5/zsh-in-docker.sh)" -- \
     -t robbyrussell -p git \
     -p https://github.com/agkozak/zsh-z \
@@ -26,18 +20,16 @@ RUN sh -c "$(wget -O- https://github.com/deluan/zsh-in-docker/releases/download/
     -p https://github.com/zsh-users/zsh-syntax-highlighting
 
 RUN printf "y\ny\ny\n\n" | bash -c "$(curl -fsSL https://raw.githubusercontent.com/Kin-Zhang/Kin-Zhang/main/scripts/setup_ohmyzsh.sh)"
-RUN /opt/conda/bin/conda init zsh && /opt/mambaforge/bin/mamba init zsh
+RUN /opt/conda/bin/conda init zsh && /opt/conda/bin/mamba init zsh
 
 # change to conda env
 ENV PATH /opt/conda/bin:$PATH
-ENV PATH /opt/mambaforge/bin:$PATH
 
-RUN mkdir -p /home/kin/workspace && cd /home/kin/workspace && git clone https://github.com/KTH-RPL/SeFlow.git
-WORKDIR /home/kin/workspace/SeFlow
+RUN mkdir -p /home/kin/workspace && cd /home/kin/workspace && git clone https://github.com/KTH-RPL/OpenSceneFlow.git
+WORKDIR /home/kin/workspace/OpenSceneFlow
 RUN apt-get update && apt-get install libgl1 -y
 # need read the gpu device info to compile the cuda extension
-RUN cd /home/kin/workspace/SeFlow && /opt/mambaforge/bin/mamba env create -f environment.yaml
-RUN cd /home/kin/workspace/SeFlow/assets/cuda/mmcv && /opt/mambaforge/envs/seflow/bin/python ./setup.py install
-RUN cd /home/kin/workspace/SeFlow/assets/cuda/chamfer3D && /opt/mambaforge/envs/seflow/bin/python ./setup.py install
-
+RUN cd /home/kin/workspace/OpenSceneFlow && /opt/conda/bin/mamba env create -f environment.yaml
+RUN cd /home/kin/workspace/OpenSceneFlow/assets/cuda/mmcv && /opt/conda/envs/opensf/bin/python ./setup.py install
+RUN cd /home/kin/workspace/OpenSceneFlow/assets/cuda/chamfer3D && /opt/conda/envs/opensf/bin/python ./setup.py install
 

README.md

@@ -1,22 +1,16 @@
 <p align="center">
-    <!-- pypi-strip -->
     <picture>
-    <!-- <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/Pointcept/Pointcept/main/docs/logo_dark.png">
-    <source media="(prefers-color-scheme: light)" srcset="https://raw.githubusercontent.com/Pointcept/Pointcept/main/docs/logo.png"> -->
-    <!-- /pypi-strip -->
     <img alt="opensceneflow" src="assets/docs/logo.png" width="600">
-    <!-- pypi-strip -->
     </picture><br>
-    <!-- /pypi-strip -->
 </p>
 
 OpenSceneFlow is a codebase for point cloud scene flow estimation. 
 It is also an official implementation of the following paper (sored by the time of publication):
 
-<!-- - **Flow4D: Leveraging 4D Voxel Network for LiDAR Scene Flow Estimation**  
+- **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  
-[ Backbone ] [ Supervised ] - [ [arXiv](https://arxiv.org/abs/2407.07995) ] [ [Project](https://github.com/dgist-cvlab/Flow4D) ] &rarr; [here](#flow4d) -->
+[ Backbone ] [ Supervised ] - [ [arXiv](https://arxiv.org/abs/2407.07995) ] [ [Project](https://github.com/dgist-cvlab/Flow4D) ] &rarr; [here](#flow4d)
 
 - **SSF: Sparse Long-Range Scene Flow for Autonomous Driving**  
 *Ajinkya Khoche, Qingwen Zhang, Laura Pereira Sánchez, Aron Asefaw, Sina Sharif Mansouri and Patric Jensfelt*  
@@ -34,51 +28,31 @@ International Conference on Robotics and Automation (**ICRA**) 2024
 [ Backbone ] [ Supervised ] - [ [arXiv](https://arxiv.org/abs/2401.16122) ] [ [Project](https://github.com/KTH-RPL/DeFlow) ] &rarr; [here](#deflow)
 
 
+💞 If you find *OpenSceneFlow* useful to your research, please cite [our works 📖](#cite-us) and give a star 🌟 as encouragement. (੭ˊ꒳​ˋ)੭✧
 
-<details> <summary>🎁 <b>One repository, All methods!</b> </summary>
+🎁 <b>One repository, All methods!</b>. Additionally, *OpenSceneFlow* integrates the following excellent work: [ICLR'24 ZeroFlow](https://arxiv.org/abs/2305.10424), [ICCV'23 FastNSF](https://arxiv.org/abs/2304.09121), [RA-L'21 FastFlow](https://arxiv.org/abs/2103.01306), [NeurIPS'21 NSFP](https://arxiv.org/abs/2111.01253), 
 
-- [x] [FastFlow3d](https://arxiv.org/abs/2103.01306): RA-L 2021
+<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).
 - [ ] [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. Done coding, public after review.
 - [ ] [FastNSF](https://arxiv.org/abs/2304.09121): ICCV 2023. Done coding, public after review.
-- [ ] [Flow4D](https://arxiv.org/abs/2407.07995): Under Review. Done coding, public after review.
-- [ ] ... more on the way
+- [ ] [ICP-Flow](https://arxiv.org/abs/2402.17351): CVPR 2024. Done coding, public after review.
 
 </details>
 
-## Citation
-
-If you find *OpenSceneFlow* useful to your research, please cite our work as encouragement. (੭ˊ꒳​ˋ)੭✧
-
-```
-@inproceedings{zhang2024seflow,
-  author={Zhang, Qingwen and Yang, Yi and Li, Peizheng and Andersson, Olov and Jensfelt, Patric},
-  title={{SeFlow}: A Self-Supervised Scene Flow Method in Autonomous Driving},
-  booktitle={European Conference on Computer Vision (ECCV)},
-  year={2024},
-  pages={353–369},
-  organization={Springer},
-  doi={10.1007/978-3-031-73232-4_20},
-}
-@inproceedings{zhang2024deflow,
-  author={Zhang, Qingwen and Yang, Yi and Fang, Heng and Geng, Ruoyu and Jensfelt, Patric},
-  booktitle={2024 IEEE International Conference on Robotics and Automation (ICRA)}, 
-  title={{DeFlow}: Decoder of Scene Flow Network in Autonomous Driving}, 
-  year={2024},
-  pages={2105-2111},
-  doi={10.1109/ICRA57147.2024.10610278}
-}
-```
+💡: Want to learn how to add your own network in this structure? Check [Contribute section] and know more about the code. Fee free to pull request and your bibtex [here](#cite-us) by pull request.
 
 ---
 
-📜 Changelog:
+<!-- 📜 Changelog:
 
 - 🎁 2025/1/28 14:58: Update the codebase to collect all methods in one repository reference [Pointcept](https://github.com/Pointcept/Pointcept) repo.
 - 🤗 2024/11/18 16:17: Update model and demo data download link through HuggingFace, Personally I found `wget` from HuggingFace link is much faster than Zenodo.
 - 2024/09/26 16:24: All codes already uploaded and tested. You can to try training directly by downloading (through [HuggingFace](https://huggingface.co/kin-zhang/OpenSceneFlow)/[Zenodo](https://zenodo.org/records/13744999)) demo data or pretrained weight for evaluation. 
 - 2024/07/24: Merging SeFlow & DeFlow code together, lighter setup and easier running.
-- 🔥 2024/07/02: Check the self-supervised version in our new ECCV'24 [SeFlow](https://github.com/KTH-RPL/SeFlow). The 1st ranking in new leaderboard among self-supervise methods.
+- 🔥 2024/07/02: Check the self-supervised version in our new ECCV'24 [SeFlow](https://github.com/KTH-RPL/SeFlow). The 1st ranking in new leaderboard among self-supervise methods. -->
 
 ## 0. Installation
 
@@ -97,33 +71,49 @@ cd assets/cuda/mmcv && python ./setup.py install && cd ../../..
 cd assets/cuda/chamfer3D && python ./setup.py install && cd ../../..
 ```
 
-<!-- Or you always can choose [Docker](https://en.wikipedia.org/wiki/Docker_(software)) which isolated environment and free yourself from installation, you can pull it by. 
-If you have different arch, please build it by yourself `cd OpenSceneFlow && docker build -t zhangkin/opensf` by going through [build-docker-image](assets/README.md/#build-docker-image) section.
+Or you always can choose [Docker](https://en.wikipedia.org/wiki/Docker_(software)) which isolated environment and free yourself from installation, you can pull it by. 
+If you have different arch, please build it by yourself `cd OpenSceneFlow && docker build -t zhangkin/opensf` by going through [build-docker-image](assets/README.md#build-docker-image) section.
+
 ```bash
 # option 1: pull from docker hub
-docker pull zhangkin/seflow
+docker pull zhangkin/opensf
 
 # run container
-docker run -it --gpus all -v /dev/shm:/dev/shm -v /home/kin/data:/home/kin/data --name seflow zhangkin/seflow /bin/zsh
-``` -->
+docker run -it --gpus all -v /dev/shm:/dev/shm -v /home/kin/data:/home/kin/data --name opensceneflow zhangkin/opensf /bin/zsh
+# and better to read your own gpu device info to compile the cuda extension again:
+cd /home/kin/workspace/OpenSceneFlow/assets/cuda/mmcv && /opt/conda/envs/opensf/bin/python ./setup.py install
+cd /home/kin/workspace/OpenSceneFlow/assets/cuda/chamfer3D && /opt/conda/envs/opensf/bin/python ./setup.py install
+```
 
 
 ## 1. Data Preparation
 
-Check [dataprocess/README.md](dataprocess/README.md#argoverse-20) for downloading tips for the raw Argoverse 2 dataset and [data preprocessed to h5 files commands](dataprocess/README.md#process).
-Another good way to try code quickly is using **mini processed dataset**, we directly provide one scene inside `train` and `val`. 
-It already converted to `.h5` format and processed with the label data. 
-You can download it from [Zenodo](https://zenodo.org/records/13744999/files/demo_data.zip)/[HuggingFace](https://huggingface.co/kin-zhang/OpenSceneFlow/blob/main/demo_data.zip) and extract it to the data folder. 
-Then you can directly use this mini processed demo data to run the [training script](#2-quick-start).
+Refer to [dataprocess/README.md](dataprocess/README.md) for dataset download instructions. Currently, we support **Argoverse 2**, **Waymo**, and **custom datasets** (more datasets will be added in the future). 
+
+After downloading, convert the raw data to `.h5` format for easy training, evaluation, and visualization. Follow the steps in [dataprocess/README.md#process](dataprocess/README.md#process). For a quick start, use our **mini processed dataset**, which includes one scene in `train` and `val`. It is pre-converted to `.h5` format with label data ([Zenodo](https://zenodo.org/records/13744999/files/demo_data.zip)/[HuggingFace](https://huggingface.co/kin-zhang/OpenSceneFlow/blob/main/demo_data.zip)).
+
 
 ```bash
 wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/demo_data.zip
 unzip demo_data.zip -p /home/kin/data/av2
 ```
 
+Once extracted, you can directly use this dataset to run the [training script](#2-quick-start) without further processing.
+
 ## 2. Quick Start
 
-<!-- ### Flow4D -->
+### Flow4D
+
+Train Flow4D with the leaderboard submit config. [Runtime: Around 18 hours in 4x RTX 3090 GPUs.]
+
+```bash
+python train.py model=flow4d lr=1e-3 epochs=15 batch_size=8 num_frames=5 loss_fn=deflowLoss "voxel_size=[0.2, 0.2, 0.2]" "point_cloud_range=[-51.2, -51.2, -3.2, 51.2, 51.2, 3.2]"
+```
+
+Pretrained weight can be downloaded through:
+```bash
+wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/flow4d_best.ckpt
+```
 
 <!-- ### SSF -->
 
@@ -132,7 +122,7 @@ unzip demo_data.zip -p /home/kin/data/av2
 Train SeFlow needed to specify the loss function, we set the config of our best model in the leaderboard. [Runtime: Around 11 hours in 4x A100 GPUs.]
 
 ```bash
-python train.py model=deflow lr=2e-4 epochs=9 batch_size=16 loss_fn=seflowLoss "add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}" "model.target.num_iters=2" "model.val_monitor=val/Dynamic/Mean"
+python train.py model=deflow lr=2e-4 epochs=9 batch_size=16 loss_fn=seflowLoss "add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}" "model.target.num_iters=2"
 ```
 
 Pretrained weight can be downloaded through:
@@ -213,11 +203,51 @@ python tools/visualization_rerun.py --data_dir /home/kin/data/av2/h5py/demo/trai
 https://github.com/user-attachments/assets/07e8d430-a867-42b7-900a-11755949de21
 
 
-## Acknowledgement
+## Cite Us
 
-These work were partially supported by the Wallenberg AI, Autonomous Systems and Software Program (WASP) funded by the Knut and Alice Wallenberg Foundation and Prosense (2020-02963) funded by Vinnova. 
-The computations were enabled by the supercomputing resource Berzelius provided by National Supercomputer Centre at Linköping University and the Knut and Alice Wallenberg Foundation, Sweden.
+*OpenSceneFlow* is designed by [Qingwen Zhang](https://kin-zhang.github.io/) from DeFlow and SeFlow project. If you find it useful, please cite our works:
+
+```bibtex
+@inproceedings{zhang2024seflow,
+  author={Zhang, Qingwen and Yang, Yi and Li, Peizheng and Andersson, Olov and Jensfelt, Patric},
+  title={{SeFlow}: A Self-Supervised Scene Flow Method in Autonomous Driving},
+  booktitle={European Conference on Computer Vision (ECCV)},
+  year={2024},
+  pages={353–369},
+  organization={Springer},
+  doi={10.1007/978-3-031-73232-4_20},
+}
+@inproceedings{zhang2024deflow,
+  author={Zhang, Qingwen and Yang, Yi and Fang, Heng and Geng, Ruoyu and Jensfelt, Patric},
+  booktitle={2024 IEEE International Conference on Robotics and Automation (ICRA)}, 
+  title={{DeFlow}: Decoder of Scene Flow Network in Autonomous Driving}, 
+  year={2024},
+  pages={2105-2111},
+  doi={10.1109/ICRA57147.2024.10610278}
+}
+```
+
+And our excellent collaborators works as followings:
+
+```bibtex
+@article{kim2025flow4d,
+  author={Kim, Jaeyeul and Woo, Jungwan and Shin, Ukcheol and Oh, Jean and Im, Sunghoon},
+  journal={IEEE Robotics and Automation Letters}, 
+  title={Flow4D: Leveraging 4D Voxel Network for LiDAR Scene Flow Estimation}, 
+  year={2025},
+  volume={10},
+  number={4},
+  pages={3462-3469},
+  doi={10.1109/LRA.2025.3542327}
+}
+@article{khoche2025ssf,
+  title={SSF: Sparse Long-Range Scene Flow for Autonomous Driving},
+  author={Khoche, Ajinkya and Zhang, Qingwen and Sanchez, Laura Pereira and Asefaw, Aron and Mansouri, Sina Sharif and Jensfelt, Patric},
+  journal={arXiv preprint arXiv:2501.17821},
+  year={2025}
+}
+```
 
-<!-- *OpenSceneFlow* is designed by [Qingwen Zhang](https://kin-zhang.github.io/). It  -->
+Feel free to contribute your method and add your bibtex here by pull request!
 
-❤️: Evaluation Metric from [BucketedSceneFlowEval](https://github.com/kylevedder/BucketedSceneFlowEval); README reference from [Pointcept](https://github.com/Pointcept/Pointcept); Many thanks to [ZeroFlow](https://github.com/kylevedder/zeroflow) ...
+❤️: [BucketedSceneFlowEval](https://github.com/kylevedder/BucketedSceneFlowEval); [Pointcept](https://github.com/Pointcept/Pointcept); [ZeroFlow](https://github.com/kylevedder/zeroflow) ...

conf/config.yaml

@@ -5,11 +5,11 @@ defaults:
 
 slurm_id: 00000
 
-wandb_mode: offline # [offline, disabled, online]
+wandb_mode: disabled # [offline, disabled, online]
 wandb_project_name: seflow
 
-train_data: /home/kin/data/av2/preprocess_v2/demo/sensor/train
-val_data: /home/kin/data/av2/preprocess_v2/demo/sensor/val
+train_data: /home/kin/data/av2/h5py/demo/train
+val_data: /home/kin/data/av2/h5py/demo/val
 
 output: ${model.name}-${slurm_id}
 

conf/eval.yaml

@@ -1,5 +1,5 @@
 
-dataset_path: /home/kin/data/av2/preprocess_v2/sensor
+dataset_path: /home/kin/data/av2/h5py/sensor
 checkpoint: /home/kin/model_zoo/deflow.ckpt
 av2_mode: val # [val, test]
 save_res: False # [True, False]

conf/model/flow4d.yaml

@@ -0,0 +1,9 @@
+name: flow4d
+
+target:
+  _target_: src.models.Flow4D
+  voxel_size: ${voxel_size}
+  point_cloud_range: ${point_cloud_range}
+  num_frames: ${num_frames}
+
+val_monitor: val/Dynamic/Mean

environment.yaml

@@ -31,6 +31,7 @@ dependencies:
     - dztimer
     - av2==0.2.1
     - dufomap==1.0.0
+    - spconv-cu117
 
 # Reason about the version fixed:
 # setuptools==68.5.1: https://github.com/aws-neuron/aws-neuron-sdk/issues/893

envprocess.yaml

@@ -6,13 +6,13 @@ dependencies:
   - python=3.8
   - pytorch::pytorch=2.0.0
   - pytorch::torchvision
-  - mkl==2024.0.0
   - numba
-  - numpy
+  - numpy==1.22
   - pandas
   - pip
   - scipy
   - tqdm
+  - scikit-learn
   - fire
   - hdbscan
   - s5cmd
@@ -21,10 +21,13 @@ dependencies:
     - nuscenes-devkit
     - av2==0.2.1
     - waymo-open-dataset-tf-2.11.0==1.5.0
-    - dufomap==1.0.0
+    - open3d==0.18.0
     - linefit
     - dztimer
+    - dufomap==1.0.0
+    - evalai
 
 # Reason about the version fixed:
 # numpy==1.22: package conflicts, need numpy higher or same 1.22
-# mkl==2024.0.0: https://github.com/pytorch/pytorch/issues/123097
+# open3d==0.18.0: because 0.17.0 have bug on set the view json file
+# dufomap==1.0.0: in case later updating may not compatible with the code.

process.py

@@ -75,7 +75,7 @@ def run_cluster(
                     del f[key]['label']
                 f[key].create_dataset('label', data=np.array(cluster_label).astype(np.int16))
         print(f"==> Scene {scene_id} finished, used: {(time.time() - start_time)/60:.2f} mins")
-    print(f"Data inside {str(data_path)} finished. Check the result with vis() function if you want to visualize them.")
+    print(f"Data inside {str(data_path)} finished. Check the result with tools/visulization.py if you want to visualize them.")
 
 def run_dufo(
     data_dir: str ="/home/kin/data/av2/preprocess/sensor/train",

src/models/__init__.py

@@ -1,2 +1,20 @@
+"""
+# Created: 2024-11-21 20:12
+# Copyright (C) 2023-now, RPL, KTH Royal Institute of Technology
+# Author: Qingwen Zhang  (https://kin-zhang.github.io/)
+#
+# This file is part of OpenSceneFlow (https://github.com/KTH-RPL/OpenSceneFlow)
+# If you find this repo helpful, please cite the respective publication as 
+# listed on the above website.
+"""
+
 from .deflow import DeFlow
-from .fastflow3d import FastFlow3D
+from .fastflow3d import FastFlow3D
+
+# following need install extra package: 
+# * pip install spconv-cu117
+try:
+    from .flow4d import Flow4D
+except ImportError as e:
+    print("\033[93m--- WARNING [model]: Model with SparseConv is not imported, as it requires spconv lib which is not installed.")
+    print(f"\033[91m--- Detail error message\033[0m: {e}")