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

* feat: nus data initial haven't finished yet.
* env(mmcv): fix pip install
in case we need pip install -e rather than use setup.py directly.
* env(mmcv): remove ext in the load module for the setup changes.
in case old env cannot load, keep two options here.
* feat(nsfp, fastnsf): update two methods module files.
* chore(runner): for all optimization-based methods.
* try to speed up if we have more gpus, but I haven't test on multi-gpu and multi-node session. Will test later once I got the setup.
* upgrade training env to python 3.10 and with cuda 11.8 for afterward env upgrading etc.
* docs(readme): update readme and fix typo from copilot reviewers.
* hotfix: multi-gpu metric evaluation
* need set cuda device during spawn.
* apply coplit pull request review.
* finished checking on 2-gpu setting. (but not yet test on multi-nodes)
* env(opensf): update environment yaml to py310, add requirement.txt for if we already have torch.
* as dockerfile shown.
* fix(docker): torch is on the base env in the new dockerfile.
* hotfix(nsfp): we can remove pytorch3d dependence.
* update env file (check all good)
* update dockerfile
* fix(sftool): argoverse 2 need 0.2.1 for py38, revert to previous version.
* fix some issues on test on runner
* but haven't test it on multi-gpu.
* env(data): update two more datasets used in the himo projects.
* update header info about HiMo
* github(action): stale issues.
* fix(runner): save mode passed 1-gpu test now.
* visualization result also checked, corrected.
* todo tested on multi-gpu to see if it works.
* docs(README): update readme and running nsfp and fastnsf.
* docs: add contributing README for afterward users.
* style(header): update header info.

Author: Kin-Zhang

.github/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 }}

.github/workflows/issue_stale.yaml

@@ -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!

CONTRIBUTING.md

@@ -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://raw.githubusercontent.com/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

Dockerfile

@@ -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:

LICENSE

@@ -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

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/README.md

@@ -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/scatter_points.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/setup.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):