Merge two optimization-based (NSFP, FastNSF) methods into codebase from HiMo project

.github/workflows/issue_stale.yaml

@@ -0,0 +1,30 @@
+name: Close inactive issues
+on:
+  schedule:
+    - cron: "35 11 * * 5"
+
+env:
+  DAYS_BEFORE_ISSUE_STALE: 30
+  DAYS_BEFORE_ISSUE_CLOSE: 14
+
+jobs:
+  close-issues:
+    runs-on: ubuntu-latest
+    permissions:
+      issues: write
+      pull-requests: write
+    steps:
+      - uses: actions/stale@v5
+        with:
+          days-before-issue-stale: ${{ env.DAYS_BEFORE_ISSUE_STALE }}
+          days-before-issue-close: ${{ env.DAYS_BEFORE_ISSUE_CLOSE }}
+          stale-issue-label: "stale"
+          stale-issue-message: |
+            This issue is stale because it has been open for ${{ env.DAYS_BEFORE_ISSUE_STALE }} days with no activity. 
+            It will be closed if no further activity occurs. Let us know if you still need help!
+          close-issue-message: |
+            This issue is being closed because it has been stale for ${{ env.DAYS_BEFORE_ISSUE_CLOSE }} days with no activity. 
+            If you still need help, please feel free to leave comments.
+          days-before-pr-stale: -1
+          days-before-pr-close: -1
+          repo-token: ${{ secrets.GITHUB_TOKEN }}

CONTRIBUTING.md

@@ -0,0 +1,43 @@
+# Contributing to OpenSceneFlow
+
+We want to make contributing to this project as easy and transparent as possible. We welcome any contributions, from bug fixes to new features. If you're interested in adding your own scene flow method, this guide will walk you through the process.
+
+## Adding a New Method
+
+Here is a quick guide to integrating a new method into the OpenSceneFlow codebase.
+
+### 1. Data Preparation
+
+All data is expected to be processed into the `.h5` format. Each file represents a scene, and within the file, each data sample is indexed by a `timestamp` key.
+
+For more details on the data processing pipeline, please see the [Data Processing README](./dataprocess/README.md#process).
+
+### 2. Model Implementation
+
+All model source files are located in [`src/models`](./src/models). When adding your model, please remember to import your new model class in the [`src/models/__init__.py`](./src/models/__init__.py) file. Don't forget to add your model conf files in [`conf/model`](./conf/model).
+
+  * **For Feed-Forward Methods:** You can use `deflow` and `fastflow3d` as implementation examples.
+  * **For Optimization-Based Methods:** Please refer to `nsfp` and `fastnsf` for guidance on structure and integration. A detailed example can be found in the [NSFP model file](./src/models/nsfp.py).
+
+### 3. Custom Loss Functions
+
+All loss functions are defined in [`src/lossfuncs.py`](./src/lossfuncs.py). If your model requires a new loss function, you can add it to this file by following the pattern of the existing functions. SeFlow provided a self-supervised loss function example for all feed-forward methods. Feel free to check.
+
+### 4.1 Training a Feed-Forward Model
+
+1.  Add a configuration file for your new model in the [`conf/model`](./conf/model) directory.
+2.  Begin training by running the following command:
+    ```bash
+    python train.py model=your_model_name
+    ```
+3.  **Note:** If your model's output dictionary (`res_dict`) has a different structure from the existing models, you may need to add a new pattern in the `training_step` and `validation_step` methods in the main training script.
+
+### 4.2 Running an Optimization-Based Model
+
+Our framework supports multi-GPU execution for optimization-based methods out of the box. You can follow the structure of existing methods like NSFP to run and evaluate your model.
+
+-----
+
+Once the steps above are completed, other parts of the framework, such as evaluation (`eval`) and visualization (`save`), should integrate with your new model accordingly.
+
+Thank you for your contribution!

Dockerfile

@@ -1,14 +1,8 @@
-# check more: https://hub.docker.com/r/nvidia/cuda
-FROM nvidia/cuda:11.7.1-devel-ubuntu20.04
+FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-devel
 ENV DEBIAN_FRONTEND noninteractive
+LABEL maintainer="Qingwen Zhang <https://kin-zhang.github.io/>"
 
-RUN apt update && apt install -y git curl vim rsync htop
-
-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 apt update && apt install -y git tmux curl vim rsync libgl1 libglib2.0-0 ca-certificates
 
 # install zsh and oh-my-zsh
 RUN apt update && apt install -y wget git zsh tmux vim g++
@@ -20,15 +14,23 @@ 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://github.com/ghraw/Kin-Zhang/Kin-Zhang/main/scripts/setup_ohmyzsh.sh)"
-RUN /opt/conda/bin/conda init zsh && /opt/conda/bin/mamba init zsh
+RUN /opt/conda/bin/conda init zsh
 
 # change to conda env
 ENV PATH /opt/conda/bin:$PATH
+RUN /opt/conda/bin/conda config --set solver libmamba
 
-RUN mkdir -p /home/kin/workspace && cd /home/kin/workspace && git clone https://github.com/KTH-RPL/OpenSceneFlow.git
+RUN mkdir -p /home/kin/workspace && cd /home/kin/workspace && git clone https://github.com/Kin-Zhang/OpenSceneFlow
 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/OpenSceneFlow && /opt/conda/bin/mamba env create -f environment.yaml
-# environment for dataprocessing inlucdes data-api
-RUN cd /home/kin/workspace/OpenSceneFlow && /opt/conda/bin/mamba env create -f envsftool.yaml
+RUN /opt/conda/bin/pip install -r /home/kin/workspace/OpenSceneFlow/requirements.txt
+RUN /opt/conda/bin/pip install FastGeodis --no-build-isolation
+RUN /opt/conda/bin/pip install --no-cache-dir -e ./assets/cuda/chamfer3D && /opt/conda/bin/pip install --no-cache-dir -e ./assets/cuda/mmcv
+
+# environment for dataprocessing includes data-api
+RUN /opt/conda/bin/conda env create -f envsftool.yaml
+RUN /opt/conda/envs/sftool/bin/pip install numpy==1.22
+
+# clean up apt cache
+RUN rm -rf /var/lib/apt/lists/* && rm -rf /root/.cache/pip

LICENSE

@@ -1,6 +1,6 @@
 BSD 3-Clause License
 
-Copyright (c) 2024, Robotics, Perception and Learning @KTH
+Copyright (c) 2024, Qingwen Zhang, Robotics, Perception and Learning @KTH
 
 Redistribution and use in source and binary forms, with or without
 modification, are permitted provided that the following conditions are met:

README.md

@@ -43,13 +43,13 @@ 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).
-- [ ] [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. SSL optimization-based. Done coding, public after review.
+- [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.
 
 </details>
 
-💡: Want to learn how to add your own network in this structure? Check [Contribute section](assets/README.md#contribute) and know more about the code. Fee free to pull request and your bibtex [here](#cite-us).
+💡: Want to learn how to add your own network in this structure? Check [Contribute section](CONTRIBUTING.md#adding-a-new-method) and know more about the code. Fee free to pull request and your bibtex [here](#cite-us).
 
 ---
 
@@ -76,6 +76,7 @@ cd OpenSceneFlow && mamba env create -f environment.yaml
 # You may need export your LD_LIBRARY_PATH with env lib
 # export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/kin/mambaforge/lib
 ```
+We also provide [requirements.txt](requirements.txt), please check usage through [Dockerfile](Dockerfile).
 
 ### Docker (Recommended for Isolation)
 
@@ -86,11 +87,13 @@ You always can choose [Docker](https://en.wikipedia.org/wiki/Docker_(software))
 docker pull zhangkin/opensf
 
 # run container
-docker run -it --gpus all -v /dev/shm:/dev/shm -v /home/kin/data:/home/kin/data --name opensceneflow zhangkin/opensf /bin/zsh
+docker run -it --net=host --gpus all -v /dev/shm:/dev/shm -v /home/kin/data:/home/kin/data --name opensf zhangkin/opensf /bin/zsh
+
 # and better to read your own gpu device info to compile the cuda extension again:
+cd /home/kin/workspace/OpenSceneFlow && git pull
 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
-
+cd /home/kin/workspace/OpenSceneFlow
 mamba activate opensf
 ```
 
@@ -119,7 +122,7 @@ Some tips before running the code:
 * If you want to use [wandb](wandb.ai), replace all `entity="kth-rpl",` to your own entity otherwise tensorboard will be used locally.
 * Set correct data path by passing the config, e.g. `train_data=/home/kin/data/av2/h5py/demo/train val_data=/home/kin/data/av2/h5py/demo/val`.
 
-And free yourself from trainning, you can download the pretrained weight from [HuggingFace](https://huggingface.co/kin-zhang/OpenSceneFlow) and we provided the detail `wget` command in each model section.
+And free yourself from trainning, you can download the pretrained weight from [HuggingFace](https://huggingface.co/kin-zhang/OpenSceneFlow) and we provided the detail `wget` command in each model section. For optimization-based method, it's train-free so you can directly run with [3. Evaluation](#3-evaluation) (check more in the evaluation section).
 
 ```bash
 mamba activate opensf
@@ -143,6 +146,8 @@ wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/flow4d_best.ckp
 Extra pakcges needed for SSF model:
 ```bash
 pip install mmengine-lite torch-scatter
+# torch-scatter might not working, then reinstall by:
+pip install https://data.pyg.org/whl/torch-2.0.0%2Bcu118/torch_scatter-2.1.2%2Bpt20cu118-cp310-cp310-linux_x86_64.whl
 ```
 
 Train SSF with the leaderboard submit config. [Runtime: Around 6 hours in 8x A100 GPUs.]
@@ -194,9 +199,12 @@ You can view Wandb dashboard for the training and evaluation results or upload r
 Since in training, we save all hyper-parameters and model checkpoints, the only thing you need to do is to specify the checkpoint path. Remember to set the data path correctly also.
 
 ```bash
-# it will directly prints all metric
+# (feed-forward): load ckpt and run it, it will directly prints all metric
 python eval.py checkpoint=/home/kin/seflow_best.ckpt av2_mode=val
 
+# (optimization-based): it might need take really long time, maybe tmux for run it.
+python eval.py model=nsfp
+
 # it will output the av2_submit.zip or av2_submit_v2.zip for you to submit to leaderboard
 python eval.py checkpoint=/home/kin/seflow_best.ckpt av2_mode=test leaderboard_version=1
 python eval.py checkpoint=/home/kin/seflow_best.ckpt av2_mode=test leaderboard_version=2
@@ -238,7 +246,10 @@ evalai challenge 2210 phase 4396 submit --file av2_submit_v2.zip --large --priva
 We provide a script to visualize the results of the model also. You can specify the checkpoint path and the data path to visualize the results. The step is quite similar to evaluation.
 
 ```bash
+# (feed-forward): load ckpt
 python save.py checkpoint=/home/kin/seflow_best.ckpt dataset_path=/home/kin/data/av2/preprocess_v2/sensor/vis
+# (optimization-based): change another model by passing model name.
+python eval.py model=nsfp dataset_path=/home/kin/data/av2/h5py/demo/val
 
 # The output of above command will be like:
 Model: DeFlow, Checkpoint from: /home/kin/model_zoo/v2/seflow_best.ckpt
@@ -252,6 +263,11 @@ python tools/visualization.py --res_name 'seflow_best' --data_dir /home/kin/data
 
 https://github.com/user-attachments/assets/f031d1a2-2d2f-4947-a01f-834ed1c146e6
 
+For exporting easy comparsion with ground truth and other methods, we also provided multi-visulization open3d window:
+```bash
+python tools/visualization.py --mode mul --res_name "['flow', 'seflow_best']" --data_dir /home/kin/data/av2/preprocess_v2/sensor/vis
+```
+
 Or another way to interact with [rerun](https://github.com/rerun-io/rerun) but please only vis scene by scene, not all at once.
 
 ```bash

assets/README.md

@@ -49,7 +49,7 @@ Then follow [this stackoverflow answers](https://stackoverflow.com/questions/596
 
 3. Then you can build the docker image:
    ```bash
-   cd OpenSceneFlow && docker build -t zhangkin/OpenSceneFlow .
+   cd OpenSceneFlow && docker build -f Dockerfile -t zhangkin/opensf .
    ```
 
 ## Installation
@@ -98,12 +98,5 @@ python -c "from assets.cuda.chamfer3D import nnChamferDis;print('successfully im
     Solved by `export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/proj/berzelius-2023-154/users/x_qinzh/mambaforge/lib`
 
 
-## Contribute
-
-If you want to contribute to new model, here are tips you can follow:
-1. Dataloader: we believe all data could be process to `.h5`, we named as different scene and inside a scene, the key of each data is timestamp. Check [dataprocess/README.md](../dataprocess/README.md#process) for more details.
-2. Model: All model files can be found [here: src/models](../src/models). You can view deflow and fastflow3d to know how to implement a new model. Don't forget to add to the `__init__.py` [file to import class](../src/models/__init__.py).
-3. Loss: All loss files can be found [here: src/lossfuncs.py](../src/lossfuncs.py). There are three loss functions already inside the file, you can add a new one following the same pattern.
-4. Training: Once you have implemented the model, you can add the model to the config file [here: conf/model](../conf/model) and train the model using the command `python train.py model=your_model_name`. One more note here may: if your res_dict from model output is different, you may need add one pattern in `def training_step` and `def validation_step`.
-
-All others like eval and vis will be changed according to the model you implemented as you follow the above steps.
+3. torch_scatter problem: `OSError: /home/kin/mambaforge/envs/opensf-v2/lib/python3.10/site-packages/torch_scatter/_version_cpu.so: undefined symbol: _ZN5torch3jit17parseSchemaOrNameERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE`
+   Solved by install the torch-cuda version: `pip install https://data.pyg.org/whl/torch-2.0.0%2Bcu118/torch_scatter-2.1.2%2Bpt20cu118-cp310-cp310-linux_x86_64.whl`

assets/cuda/mmcv/scatter_points.py

@@ -8,14 +8,24 @@
 
 # from utils import ext_loader
 import importlib
-def load_ext(name, funcs):
-    ext = importlib.import_module('mmcv.' + name)
-    for fun in funcs:
-        assert hasattr(ext, fun), f'{fun} miss in module {name}'
-    return ext
 
+def load_ext(possible_names, funcs):
+    """Try loading module from list of possible names, return first matching."""
+    for name in possible_names:
+        try:
+            ext = importlib.import_module('mmcv' + name)
+            missing = [f for f in funcs if not hasattr(ext, f)]
+            if missing:
+                print(f"Missing functions in 'mmcv{name}': {missing}")
+                continue
+            return ext  # success
+        except (ModuleNotFoundError, ImportError) as e:
+            print(f"Failed to import mmcv{name}: {e}")
+    raise ImportError(f"Could not load mmcv extension with functions: {funcs}")
+
+# Usage
 ext_module = load_ext(
-    '_ext',
+    ['', '._ext'],
     ['dynamic_point_to_voxel_forward', 'dynamic_point_to_voxel_backward'])
 
 

assets/cuda/mmcv/setup.py

@@ -8,7 +8,7 @@
     version='1.0.1',
     ext_modules=[
         CUDAExtension(
-            name='mmcv._ext',
+            name='mmcv',
             sources=[
                 "/".join(__file__.split("/")[:-1] + ["scatter_points_cuda.cu"]),
                 "/".join(__file__.split("/")[:-1] + ["scatter_points.cpp"]),

assets/cuda/mmcv/voxelize.py

@@ -8,14 +8,25 @@
 
 # from utils import ext_loader
 import importlib
-def load_ext(name, funcs):
-    ext = importlib.import_module('mmcv.' + name)
-    for fun in funcs:
-        assert hasattr(ext, fun), f'{fun} miss in module {name}'
-    return ext
 
+def load_ext(possible_names, funcs):
+    """Try loading module from list of possible names, return first matching."""
+    for name in possible_names:
+        try:
+            ext = importlib.import_module('mmcv' + name)
+            missing = [f for f in funcs if not hasattr(ext, f)]
+            if missing:
+                print(f"Missing functions in 'mmcv{name}': {missing}")
+                continue
+            return ext  # success
+        except (ModuleNotFoundError, ImportError) as e:
+            print(f"Failed to import mmcv{name}: {e}")
+    raise ImportError(f"Could not load mmcv extension with functions: {funcs}")
+
+# Usage
 ext_module = load_ext(
-    '_ext', ['dynamic_voxelize_forward', 'hard_voxelize_forward'])
+    ['', '._ext'],
+    ['dynamic_voxelize_forward', 'hard_voxelize_forward'])
 
 
 class _Voxelization(Function):

conf/eval.yaml

@@ -8,6 +8,8 @@ leaderboard_version: 1 # [1, 2]
 supervised_flag: True # [True, False], whether you use any label from the dataset
 
 # no need to change
+save_res_path: 
+num_frames: 2
 slurm_id: 00000
 output: ${model.name}-${slurm_id}
 gpus: 1

conf/model/fastnsf.yaml

@@ -0,0 +1,17 @@
+name: fastnsf
+
+target:
+  _target_: src.models.FastNSF
+  filter_size: 128
+  act_fn: relu
+  layer_size: 8
+  itr_num: 5000
+  lr: 8e-3
+  min_delta: 5e-5
+  grid_factor: 10
+  init_weight: True
+  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

conf/model/nsfp.yaml

@@ -0,0 +1,15 @@
+name: nsfp
+
+target:
+  _target_: src.models.NSFP
+  filter_size: 128
+  act_fn: relu
+  layer_size: 8
+  itr_num: 5000
+  lr: 8e-3
+  min_delta: 5e-5
+  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

conf/save.yaml

@@ -4,6 +4,7 @@ res_name: # if None will directly be the `model_name.ckpt` in checkpoint path
 
 
 # no need to change
+num_frames: 2
 defaults:
   - hydra: default
   - model: deflow