Merge SSF into codebase (#4)

* update ssf model
* add env dependencies, fix bug
* fix bug in ssf metric

* added details in readme
* docs: move cite at the end of README.

* style: change print style.
* think about whether have better way to print #number of pts.
* add SSF into model init file.

* env: strict spconv version to avoid error on sparse lib.
* update readme with link and add OpenPCSeg also.
* update rerun into the main env: #3

* hotfix(compile): split ssf_module with mmengine and torch scatter.
* so that users can run other models even without engine or scatter.
* move two packages into README for the extra package at least for now.... in case people struggling with setup env again.
* docs update ssf model link.

* docs(readme): update readme.

---------

Co-authored-by: Kin <kinzhangglimmer@gmail.com>

Author: ajinkyakhoche

README.md

@@ -1,9 +1,13 @@
 <p align="center">
+    <a href="https://github.com/KTH-RPL/OpenSceneFlow">
     <picture>
     <img alt="opensceneflow" src="assets/docs/logo.png" width="600">
     </picture><br>
+    </a>
 </p>
 
+💞 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):
 
@@ -27,10 +31,7 @@ European Conference on Computer Vision (**ECCV**) 2024
 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. (੭ˊ꒳​ˋ)੭✧
-
-🎁 <b>One repository, All methods!</b>
+🎁 <b>One repository, All methods!</b> 
 Additionally, *OpenSceneFlow* integrates following excellent works: [ICLR'24 ZeroFlow](https://arxiv.org/abs/2305.10424), [ICCV'23 FastNSF](https://arxiv.org/abs/2304.09121), [RA-L'21 FastFlow](https://arxiv.org/abs/2103.01306), [NeurIPS'21 NSFP](https://arxiv.org/abs/2111.01253). (More on the way...)
 
 <details> <summary> Summary of them:</summary>
@@ -43,7 +44,7 @@ Additionally, *OpenSceneFlow* integrates following excellent works: [ICLR'24 Zer
 
 </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) by pull request.
+💡: 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).
 
 ---
 
@@ -102,7 +103,7 @@ Refer to [dataprocess/README.md](dataprocess/README.md) for dataset download ins
 
 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)).
+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 ([HuggingFace](https://huggingface.co/kin-zhang/OpenSceneFlow/blob/main/demo_data.zip)/[Zenodo](https://zenodo.org/records/13744999/files/demo_data.zip)).
 
 
 ```bash
@@ -114,7 +115,9 @@ Once extracted, you can directly use this dataset to run the [training script](#
 
 ## 2. Quick Start
 
-Don't forget to active Python environment before running the code.
+Don't forget to active Python environment 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.
+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.
 
 ```bash
 mamba activate opensf
@@ -133,7 +136,28 @@ Pretrained weight can be downloaded through:
 wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/flow4d_best.ckpt
 ```
 
-<!-- ### SSF -->
+### SSF
+
+Extra pakcges needed for SSF model:
+```bash
+pip install mmengine-lite torch-scatter
+```
+
+Train SSF with the leaderboard submit config. [Runtime: Around 6 hours in 8x A100 GPUs.]
+
+```bash
+python train.py model=ssf lr=8e-3 epochs=25 batch_size=64 loss_fn=deflowLoss "voxel_size=[0.2, 0.2, 6]" "point_cloud_range=[-51.2, -51.2, -3, 51.2, 51.2, 3]"
+```
+
+Pretrained weight can be downloaded through:
+```bash
+# the leaderboard weight
+wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/ssf_best.ckpt
+
+# the long-range weight:
+wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/ssf_long.ckpt
+```
+
 
 ### SeFlow
 
@@ -223,7 +247,8 @@ https://github.com/user-attachments/assets/07e8d430-a867-42b7-900a-11755949de21
 
 ## Cite Us
 
-*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:
+[*OpenSceneFlow*](https://github.com/KTH-RPL/OpenSceneFlow) is originally designed by [Qingwen Zhang](https://kin-zhang.github.io/) from DeFlow and SeFlow. 
+If you find it useful, please cite our works:
 
 ```bibtex
 @inproceedings{zhang2024seflow,
@@ -251,7 +276,7 @@ https://github.com/user-attachments/assets/07e8d430-a867-42b7-900a-11755949de21
 }
 ```
 
-And our excellent collaborators works as followings:
+And our excellent collaborators works contributed to this codebase also:
 
 ```bibtex
 @article{kim2025flow4d,
@@ -272,6 +297,7 @@ And our excellent collaborators works as followings:
 }
 ```
 
+Thank you for your support! ❤️
 Feel free to contribute your method and add your bibtex here by pull request!
 
-❤️: [BucketedSceneFlowEval](https://github.com/kylevedder/BucketedSceneFlowEval); [Pointcept](https://github.com/Pointcept/Pointcept); [ZeroFlow](https://github.com/kylevedder/zeroflow) ...
+❤️: [BucketedSceneFlowEval](https://github.com/kylevedder/BucketedSceneFlowEval); [Pointcept](https://github.com/Pointcept/Pointcept); [OpenPCSeg](https://github.com/BAI-Yeqi/OpenPCSeg); [ZeroFlow](https://github.com/kylevedder/zeroflow) ...

assets/README.md

@@ -69,6 +69,7 @@ Create base env: [~5 mins]
 
 ```bash
 git clone https://github.com/KTH-RPL/OpenSceneFlow.git
+cd OpenSceneFlow
 mamba env create -f assets/environment.yml
 ```
 
@@ -92,16 +93,16 @@ python -c "from assets.cuda.chamfer3D import nnChamferDis;print('successfully im
 
 ### Other issues
 
-1. looks like open3d and fire package conflict, not sure
+<!-- 1. looks like open3d and fire package conflict, not sure
    -  need install package like `pip install --ignore-installed`, ref: [pip cannot install distutils installed project](https://stackoverflow.com/questions/53807511/pip-cannot-uninstall-package-it-is-a-distutils-installed-project), my error: `ERROR: Cannot uninstall 'blinker'.`
-   -  need specific werkzeug version for open3d 0.16.0, otherwise error: `ImportError: cannot import name 'url_quote' from 'werkzeug.urls'`. But need update to solve the problem: `pip install --upgrade Flask` [ref](https://stackoverflow.com/questions/77213053/why-did-flask-start-failing-with-importerror-cannot-import-name-url-quote-fr)
+   -  need specific werkzeug version for open3d 0.16.0, otherwise error: `ImportError: cannot import name 'url_quote' from 'werkzeug.urls'`. But need update to solve the problem: `pip install --upgrade Flask` [ref](https://stackoverflow.com/questions/77213053/why-did-flask-start-failing-with-importerror-cannot-import-name-url-quote-fr) -->
 
 
-2. `ImportError: libtorch_cuda.so: undefined symbol: cudaGraphInstantiateWithFlags, version libcudart.so.11.0`
+1. `ImportError: libtorch_cuda.so: undefined symbol: cudaGraphInstantiateWithFlags, version libcudart.so.11.0`
    The cuda version: `pytorch::pytorch-cuda` and `nvidia::cudatoolkit` need be same. [Reference link](https://github.com/pytorch/pytorch/issues/90673#issuecomment-1563799299)
 
 
-3. In cluster have error: `pandas ImportError: /lib64/libstdc++.so.6: version 'GLIBCXX_3.4.29' not found`
+2. In cluster have error: `pandas ImportError: /lib64/libstdc++.so.6: version 'GLIBCXX_3.4.29' not found`
     Solved by `export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/proj/berzelius-2023-154/users/x_qinzh/mambaforge/lib`
 
 

assets/cuda/mmdet/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from .conv import *
+from .norm import *
+from .plugin import *
+from .resnet import BasicBlock, Bottleneck

assets/cuda/mmdet/conv.py

@@ -0,0 +1,51 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import inspect
+from typing import Dict, Optional
+
+from mmengine.registry import MODELS
+from torch import nn
+
+MODELS.register_module('Conv1d', module=nn.Conv1d)
+MODELS.register_module('Conv2d', module=nn.Conv2d)
+MODELS.register_module('Conv3d', module=nn.Conv3d)
+MODELS.register_module('Conv', module=nn.Conv2d)
+
+
+def build_conv_layer(cfg: Optional[Dict], *args, **kwargs) -> nn.Module:
+    """Build convolution layer.
+
+    Args:
+        cfg (None or dict): The conv layer config, which should contain:
+            - type (str): Layer type.
+            - layer args: Args needed to instantiate an conv layer.
+        args (argument list): Arguments passed to the `__init__`
+            method of the corresponding conv layer.
+        kwargs (keyword arguments): Keyword arguments passed to the `__init__`
+            method of the corresponding conv layer.
+
+    Returns:
+        nn.Module: Created conv layer.
+    """
+    if cfg is None:
+        cfg_ = dict(type='Conv2d')
+    else:
+        if not isinstance(cfg, dict):
+            raise TypeError('cfg must be a dict')
+        if 'type' not in cfg:
+            raise KeyError('the cfg dict must contain the key "type"')
+        cfg_ = cfg.copy()
+
+    layer_type = cfg_.pop('type')
+    if inspect.isclass(layer_type):
+        return layer_type(*args, **kwargs, **cfg_)  # type: ignore
+    # Switch registry to the target scope. If `conv_layer` cannot be found
+    # in the registry, fallback to search `conv_layer` in the
+    # mmengine.MODELS.
+    with MODELS.switch_scope_and_registry(None) as registry:
+        conv_layer = registry.get(layer_type)
+    if conv_layer is None:
+        raise KeyError(f'Cannot find {conv_layer} in registry under scope '
+                       f'name {registry.scope}')
+    layer = conv_layer(*args, **kwargs, **cfg_)
+
+    return layer