diff --git a/Dockerfile b/Dockerfile
index b80f234..1a3cda8 100644
--- a/Dockerfile
+++ b/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
diff --git a/README.md b/README.md
index 8c77302..782ed8a 100644
--- a/README.md
+++ b/README.md
@@ -1,22 +1,16 @@
-
-
-
-
-
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):
-
+[ Backbone ] [ Supervised ] - [ [arXiv](https://arxiv.org/abs/2407.07995) ] [ [Project](https://github.com/dgist-cvlab/Flow4D) ] → [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) ] → [here](#deflow)
+💞 If you find *OpenSceneFlow* useful to your research, please cite [our works 📖](#cite-us) and give a star 🌟 as encouragement. (੭ˊ꒳ˋ)੭✧
- 🎁 One repository, All methods!
+🎁 One repository, All methods!. 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
+ Summary of them:
+
+- [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.
-## 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:
+
## 0. Installation
@@ -97,33 +71,49 @@ cd assets/cuda/mmcv && python ./setup.py install && cd ../../..
cd assets/cuda/chamfer3D && python ./setup.py install && cd ../../..
```
-
+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
+
+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
+```
@@ -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}
+}
+```
-
+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) ...
\ No newline at end of file
+❤️: [BucketedSceneFlowEval](https://github.com/kylevedder/BucketedSceneFlowEval); [Pointcept](https://github.com/Pointcept/Pointcept); [ZeroFlow](https://github.com/kylevedder/zeroflow) ...
diff --git a/conf/config.yaml b/conf/config.yaml
index 2fd5932..1e34f76 100644
--- a/conf/config.yaml
+++ b/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}
diff --git a/conf/eval.yaml b/conf/eval.yaml
index cec0559..bc8c216 100644
--- a/conf/eval.yaml
+++ b/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]
diff --git a/conf/model/flow4d.yaml b/conf/model/flow4d.yaml
new file mode 100644
index 0000000..922fc15
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/environment.yaml b/environment.yaml
index 7fa83a8..8d99068 100644
--- a/environment.yaml
+++ b/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
diff --git a/envprocess.yaml b/envprocess.yaml
index 3203a90..b977c38 100644
--- a/envprocess.yaml
+++ b/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
\ No newline at end of file
+# 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.
\ No newline at end of file
diff --git a/process.py b/process.py
index 734c68a..076747d 100644
--- a/process.py
+++ b/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",
diff --git a/src/models/__init__.py b/src/models/__init__.py
index 3e7a139..347b5a6 100644
--- a/src/models/__init__.py
+++ b/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
\ No newline at end of file
+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}")
\ No newline at end of file
diff --git a/src/models/basic/__init__.py b/src/models/basic/__init__.py
index 7d46b47..244c04c 100644
--- a/src/models/basic/__init__.py
+++ b/src/models/basic/__init__.py
@@ -14,6 +14,50 @@ def cal_pose0to1(pose0: torch.Tensor, pose1: torch.Tensor):
pose_0to1 = pose1_inv @ pose0.type(torch.float64)
return pose_0to1.type(torch.float32)
+def wrap_batch_pcs(batch, num_frames=2):
+ batch_sizes = len(batch["pose0"])
+
+ pose_flows = []
+ transform_pc0s = []
+ transform_pc_m_frames = [[] for _ in range(num_frames - 2)]
+ # print(batch)
+ for batch_id in range(batch_sizes):
+ selected_pc0 = batch["pc0"][batch_id]
+ with torch.no_grad():
+ if 'ego_motion' in batch:
+ pose_0to1 = batch['ego_motion'][batch_id].type(torch.float32)
+ else:
+ pose_0to1 = cal_pose0to1(batch["pose0"][batch_id], batch["pose1"][batch_id])
+ if num_frames > 2:
+ past_poses = []
+ for i in range(1, num_frames - 1):
+ past_pose = cal_pose0to1(batch[f"poseh{i}"][batch_id], batch["pose1"][batch_id])
+ past_poses.append(past_pose)
+
+ transform_pc0 = selected_pc0 @ pose_0to1[:3, :3].T + pose_0to1[:3, 3] #t -> t+1 warping
+
+ pose_flows.append(transform_pc0 - selected_pc0)
+ transform_pc0s.append(transform_pc0)
+
+ for i in range(1, num_frames - 1):
+ selected_pc_m = batch[f"pch{i}"][batch_id]
+ transform_pc_m = selected_pc_m @ past_poses[i-1][:3, :3].T + past_poses[i-1][:3, 3]
+ transform_pc_m_frames[i-1].append(transform_pc_m)
+
+ pc_m_frames = [torch.stack(transform_pc_m_frames[i], dim=0) for i in range(num_frames - 2)]
+
+ pc0s = torch.stack(transform_pc0s, dim=0)
+ pc1s = batch["pc1"]
+ pcs_dict = {
+ 'pc0s': pc0s,
+ 'pc1s': pc1s,
+ 'pose_flows': pose_flows
+ }
+ for i in range(1, num_frames - 1):
+ pcs_dict[f'pch{i}s'] = pc_m_frames[i-1]
+
+ return pcs_dict
+
class ConvWithNorms(nn.Module):
def __init__(self, in_num_channels: int, out_num_channels: int,
diff --git a/src/models/basic/flow4d_module.py b/src/models/basic/flow4d_module.py
new file mode 100644
index 0000000..61827cb
--- /dev/null
+++ b/src/models/basic/flow4d_module.py
@@ -0,0 +1,348 @@
+"""
+# Updated: 2024-07-12 01:16
+# Please visit: https://github.com/dgist-cvlab/Flow4D to cite this network model.
+# Author: Jaeyeul Kim (jykim94@dgist.ac.kr), Qingwen Zhang (https://kin-zhang.github.io/)
+#
+#
+Changelog:
+2024/07/13 (Qingwen): spconv.ConvAlgo.Native needed to be added to all spconv layers, add some references here:
+- https://github.com/traveller59/spconv/issues/467
+- https://github.com/traveller59/spconv/issues/482
+
+I tried half hour to find the solution... and note down here for other readers:
+1. It's because of KeyError: ((16, 8, 8), float, float)
+2. Then I checked cumm/gemm file: # ((16, 8, 8), dtypes.float32, dtypes.float32): MmaM16N8K8F32((8, 0)),
+3. This line is commented that's why float32 float 32 is KeyError
+4. And I didn't find a solution to have cumm/gemm work out. But change all algorithm on layers to: spconv.ConvAlgo.Native
+
+This file is originally copied from: https://github.com/dgist-cvlab/Flow4D
+with some modifications to have unified format with all benchmark. Check above changelog I made.
+
+"""
+import torch, os
+import torch.nn as nn
+
+from .encoder import DynamicVoxelizer, DynamicPillarFeatureNet
+import spconv.pytorch as spconv
+import spconv as spconv_core
+spconv_core.constants.SPCONV_ALLOW_TF32 = True
+
+class DynamicEmbedder_4D(nn.Module):
+
+ def __init__(self, voxel_size, pseudo_image_dims, point_cloud_range,
+ feat_channels: int) -> None:
+ super().__init__()
+ self.voxelizer = DynamicVoxelizer(voxel_size=voxel_size,
+ point_cloud_range=point_cloud_range)
+ self.feature_net = DynamicPillarFeatureNet(
+ in_channels=3,
+ feat_channels=(feat_channels, ),
+ point_cloud_range=point_cloud_range,
+ voxel_size=voxel_size,
+ mode='avg')
+
+ self.voxel_spatial_shape = pseudo_image_dims
+
+ def forward(self, input_dict) -> torch.Tensor:
+ voxel_feats_list = []
+ voxel_coors_list = []
+ batch_index = 0
+
+ frame_keys = sorted([key for key in input_dict.keys() if key.startswith('pch')], reverse=True)
+ frame_keys += ['pc0s', 'pc1s']
+
+ pch1s_3dvoxel_infos_lst = None
+ pc0_point_feats_lst, pc1_point_feats_lst = [], []
+ for time_index, frame_key in enumerate(frame_keys):
+ pc = input_dict[frame_key]
+ voxel_info_list = self.voxelizer(pc)
+
+ voxel_feats_list_batch = []
+ voxel_coors_list_batch = []
+
+ for batch_index, voxel_info_dict in enumerate(voxel_info_list):
+ points = voxel_info_dict['points']
+ coordinates = voxel_info_dict['voxel_coords']
+ voxel_feats, voxel_coors, point_feats = self.feature_net(points, coordinates)
+
+ if frame_key == 'pc0s':
+ pc0_point_feats_lst.append(point_feats)
+ elif frame_key == 'pc1s':
+ pc1_point_feats_lst.append(point_feats)
+
+ batch_indices = torch.full((voxel_coors.size(0), 1), batch_index, dtype=torch.long, device=voxel_coors.device)
+ voxel_coors_batch = torch.cat([batch_indices, voxel_coors[:, [2, 1, 0]]], dim=1)
+
+ voxel_feats_list_batch.append(voxel_feats)
+ voxel_coors_list_batch.append(voxel_coors_batch)
+
+ voxel_feats_sp = torch.cat(voxel_feats_list_batch, dim=0)
+ coors_batch_sp = torch.cat(voxel_coors_list_batch, dim=0).to(dtype=torch.int32)
+
+ time_dimension = torch.full((coors_batch_sp.shape[0], 1), time_index, dtype=torch.int32, device='cuda')
+ coors_batch_sp_4d = torch.cat((coors_batch_sp, time_dimension), dim=1)
+
+ voxel_feats_list.append(voxel_feats_sp)
+ voxel_coors_list.append(coors_batch_sp_4d)
+
+ if frame_key == 'pc0s':
+ pc0s_3dvoxel_infos_lst = voxel_info_list
+ pc0s_num_voxels = voxel_feats_sp.shape[0]
+ elif frame_key == 'pc1s':
+ pc1s_3dvoxel_infos_lst = voxel_info_list
+ pc1s_num_voxels = voxel_feats_sp.shape[0]
+ elif frame_key == 'pch1s':
+ pch1s_3dvoxel_infos_lst = voxel_info_list
+
+ all_voxel_feats_sp = torch.cat(voxel_feats_list, dim=0)
+ all_coors_batch_sp_4d = torch.cat(voxel_coors_list, dim=0)
+
+ sparse_tensor_4d = spconv.SparseConvTensor(all_voxel_feats_sp.contiguous(), all_coors_batch_sp_4d.contiguous(), self.voxel_spatial_shape, int(batch_index + 1))
+ # dense shape: B C X Y Z T
+ output = {
+ '4d_tensor': sparse_tensor_4d,
+ 'pch1_3dvoxel_infos_lst': pch1s_3dvoxel_infos_lst,
+ 'pc0_3dvoxel_infos_lst': pc0s_3dvoxel_infos_lst,
+ 'pc0_point_feats_lst': pc0_point_feats_lst,
+ 'pc0_num_voxels': pc0s_num_voxels,
+ 'pc1_3dvoxel_infos_lst': pc1s_3dvoxel_infos_lst,
+ 'pc1_point_feats_lst': pc1_point_feats_lst,
+ 'pc1_num_voxels': pc1s_num_voxels
+ }
+
+ return output
+
+def conv1x1x1x3(in_planes, out_planes, stride=1, indice_key=None):
+ return spconv.SubMConv4d(in_planes, out_planes, kernel_size=(1,1,1,3), stride=stride,
+ padding=(0,0,0,1), bias=False, indice_key=indice_key, algo=spconv.ConvAlgo.Native)
+
+def conv3x3x3x1(in_planes, out_planes, stride=1, indice_key=None):
+ return spconv.SubMConv4d(in_planes, out_planes, kernel_size=(3,3,3,1), stride=stride,
+ padding=(1,1,1,0), bias=False, indice_key=indice_key, algo=spconv.ConvAlgo.Native)
+
+def conv1x1x1x1(in_planes, out_planes, stride=1, indice_key=None):
+ return spconv.SubMConv4d(in_planes, out_planes, kernel_size=(1,1,1,1), stride=stride,
+ padding=0, bias=False, indice_key=indice_key, algo=spconv.ConvAlgo.Native)
+
+
+def conv3x3x3x3(in_planes, out_planes, stride=1, indice_key=None):
+ return spconv.SubMConv4d(in_planes, out_planes, kernel_size=(3,3,3,3), stride=stride,
+ padding=(1,1,1,1), bias=False, indice_key=indice_key, algo=spconv.ConvAlgo.Native)
+
+
+class Seperate_to_3D(nn.Module):
+ def __init__(self, num_frames):
+ super(Seperate_to_3D, self).__init__()
+ self.num_frames = num_frames
+ #self.return_pc1 = return_pc1
+
+ def forward(self, sparse_4D_tensor):
+
+ indices_4d = sparse_4D_tensor.indices
+ features_4d = sparse_4D_tensor.features
+
+ pc0_time_value = self.num_frames-2
+
+ mask_pc0 = (indices_4d[:, -1] == pc0_time_value)
+
+ pc0_indices = indices_4d[mask_pc0][:, :-1]
+ pc0_features = features_4d[mask_pc0]
+
+ pc0_sparse_3D = sparse_4D_tensor.replace_feature(pc0_features)
+ pc0_sparse_3D.spatial_shape = sparse_4D_tensor.spatial_shape[:-1]
+ pc0_sparse_3D.indices = pc0_indices
+
+ return pc0_sparse_3D
+
+class SpatioTemporal_Decomposition_Block(nn.Module):
+ def __init__(self, in_filters, mid_filters, out_filters, indice_key=None, down_key = None, pooling=False, z_pooling=True, interact=False):
+ super(SpatioTemporal_Decomposition_Block, self).__init__()
+
+
+ self.pooling = pooling
+
+ self.act = nn.LeakyReLU()
+
+ self.spatial_conv_1 = conv3x3x3x1(in_filters, mid_filters, indice_key=indice_key + "bef")
+ self.bn_s_1 = nn.BatchNorm1d(mid_filters)
+
+ self.temporal_conv_1 = conv1x1x1x3(in_filters, mid_filters)
+ self.bn_t_1 = nn.BatchNorm1d(mid_filters)
+
+ self.fusion_conv_1 = conv1x1x1x1(mid_filters*2+in_filters, mid_filters, indice_key=indice_key + "1D")
+ self.bn_fusion_1 = nn.BatchNorm1d(mid_filters)
+
+
+ self.spatial_conv_2 = conv3x3x3x1(mid_filters, mid_filters, indice_key=indice_key + "bef")
+ self.bn_s_2 = nn.BatchNorm1d(mid_filters)
+
+ self.temporal_conv_2 = conv1x1x1x3(mid_filters, mid_filters)
+ self.bn_t_2 = nn.BatchNorm1d(mid_filters)
+
+ self.fusion_conv_2 = conv1x1x1x1(mid_filters*3, out_filters, indice_key=indice_key + "1D")
+ self.bn_fusion_2 = nn.BatchNorm1d(out_filters)
+
+
+ if self.pooling:
+ if z_pooling == True:
+ self.pool = spconv.SparseConv4d(out_filters, out_filters, kernel_size=(2,2,2,1), stride=(2,2,2,1), indice_key=down_key, bias=False, algo=spconv.ConvAlgo.Native)
+ else:
+ self.pool = spconv.SparseConv4d(out_filters, out_filters, kernel_size=(2,2,1,1), stride=(2,2,1,1), indice_key=down_key, bias=False, algo=spconv.ConvAlgo.Native)
+
+ self.weight_initialization()
+
+ def weight_initialization(self):
+ for m in self.modules():
+ if isinstance(m, nn.BatchNorm1d):
+ nn.init.constant_(m.weight, 1)
+ nn.init.constant_(m.bias, 0)
+
+ def forward(self, x):
+
+ #ST block
+ S_feat_1 = self.spatial_conv_1(x)
+ S_feat_1 = S_feat_1.replace_feature(self.bn_s_1(S_feat_1.features))
+ S_feat_1 = S_feat_1.replace_feature(self.act(S_feat_1.features))
+
+ T_feat_1 = self.temporal_conv_1(x)
+ T_feat_1 = T_feat_1.replace_feature(self.bn_t_1(T_feat_1.features))
+ T_feat_1 = T_feat_1.replace_feature(self.act(T_feat_1.features))
+
+ ST_feat_1 = x.replace_feature(torch.cat([S_feat_1.features, T_feat_1.features, x.features], 1)) #residual까지 concate
+
+ ST_feat_1 = self.fusion_conv_1(ST_feat_1)
+ ST_feat_1 = ST_feat_1.replace_feature(self.bn_fusion_1(ST_feat_1.features))
+ ST_feat_1 = ST_feat_1.replace_feature(self.act(ST_feat_1.features))
+
+ #TS block
+ S_feat_2 = self.spatial_conv_2(ST_feat_1)
+ S_feat_2 = S_feat_2.replace_feature(self.bn_s_2(S_feat_2.features))
+ S_feat_2 = S_feat_2.replace_feature(self.act(S_feat_2.features))
+
+ T_feat_2 = self.temporal_conv_2(ST_feat_1)
+ T_feat_2 = T_feat_2.replace_feature(self.bn_t_2(T_feat_2.features))
+ T_feat_2 = T_feat_2.replace_feature(self.act(T_feat_2.features))
+
+ ST_feat_2 = x.replace_feature(torch.cat([S_feat_2.features, T_feat_2.features, ST_feat_1.features], 1)) #residual까지 concate
+
+ ST_feat_2 = self.fusion_conv_2(ST_feat_2)
+ ST_feat_2 = ST_feat_2.replace_feature(self.bn_fusion_2(ST_feat_2.features))
+ ST_feat_2 = ST_feat_2.replace_feature(self.act(ST_feat_2.features))
+
+ if self.pooling:
+ pooled = self.pool(ST_feat_2)
+ return pooled, ST_feat_2
+ else:
+ return ST_feat_2
+
+
+class Network_4D(nn.Module):
+ def __init__(self, in_channel=16, out_channel=16, model_size = 16):
+ super().__init__()
+
+ SpatioTemporal_Block = SpatioTemporal_Decomposition_Block
+
+ self.model_size = model_size
+
+
+ self.STDB_1_1_1 = SpatioTemporal_Block(in_channel, model_size, model_size, indice_key="st1_1", down_key='floor1')
+ self.STDB_1_1_2 = SpatioTemporal_Block(model_size, model_size, model_size*2, indice_key="st1_1", down_key='floor1', pooling=True) #512 512 32 -> 256 256 16
+
+ self.STDB_2_1_1 = SpatioTemporal_Block(model_size*2, model_size*2, model_size*2, indice_key="st2_1", down_key='floor2')
+ self.STDB_2_1_2 = SpatioTemporal_Block(model_size*2, model_size*2, model_size*4, indice_key="st2_1", down_key='floor2', pooling=True) #256 256 16 -> 128 128 8
+
+ self.STDB_3_1_1 = SpatioTemporal_Block(model_size*4, model_size*4, model_size*4, indice_key="st3_1", down_key='floor3')
+ self.STDB_3_1_2 = SpatioTemporal_Block(model_size*4, model_size*4, model_size*4, indice_key="st3_1", down_key='floor3', pooling=True) #128 128 8 -> 64 64 4
+
+ self.STDB_4_1_1 = SpatioTemporal_Block(model_size*4, model_size*4, model_size*4, indice_key="st4_1", down_key='floor4')
+ self.STDB_4_1_2 = SpatioTemporal_Block(model_size*4, model_size*4, model_size*4, indice_key="st4_1", down_key='floor4', pooling=True, z_pooling=False) #64 64 4 -> 64 64 4
+
+ self.STDB_5_1_1 = SpatioTemporal_Block(model_size*4, model_size*4, model_size*4, indice_key="st5_1")
+ self.STDB_5_1_2 = SpatioTemporal_Block(model_size*4, model_size*4, model_size*4, indice_key="st5_1")
+ self.up_subm_5 = spconv.SparseInverseConv4d(model_size*4, model_size*4, kernel_size=(2,2,1,1), indice_key='floor4', bias=False, algo=spconv.ConvAlgo.Native) #zpooling false
+
+ self.STDB_4_2_1 = SpatioTemporal_Block(model_size*8, model_size*8, model_size*4, indice_key="st4_2")
+ self.up_subm_4 = spconv.SparseInverseConv4d(model_size*4, model_size*4, kernel_size=(2,2,2,1), indice_key='floor3', bias=False, algo=spconv.ConvAlgo.Native)
+
+ self.STDB_3_2_1 = SpatioTemporal_Block(model_size*8, model_size*8, model_size*4, indice_key="st3_2")
+ self.up_subm_3 = spconv.SparseInverseConv4d(model_size*4, model_size*4, kernel_size=(2,2,2,1), indice_key='floor2', bias=False, algo=spconv.ConvAlgo.Native)
+
+ self.STDB_2_2_1 = SpatioTemporal_Block(model_size*8, model_size*4, model_size*4, indice_key="st_2_2")
+ self.up_subm_2 = spconv.SparseInverseConv4d(model_size*4, model_size*2, kernel_size=(2,2,2,1), indice_key='floor1', bias=False, algo=spconv.ConvAlgo.Native)
+
+ self.STDB_1_2_1 = SpatioTemporal_Block(model_size*4, model_size*2, out_channel, indice_key="st_1_2")
+
+
+ def forward(self, sp_tensor):
+
+ sp_tensor = self.STDB_1_1_1(sp_tensor) # [B, C, 512, 512, 32, T] [512,512,32] = [coor_x, coor_y, coor_z]
+ down_2, skip_1 = self.STDB_1_1_2(sp_tensor)
+
+ down_2 = self.STDB_2_1_1(down_2)
+ down_3, skip_2 = self.STDB_2_1_2(down_2)
+
+ down_3 = self.STDB_3_1_1(down_3)
+ down_4, skip_3 = self.STDB_3_1_2(down_3)
+
+ down_4 = self.STDB_4_1_1(down_4)
+ down_5, skip_4 = self.STDB_4_1_2(down_4)
+
+ down_5 = self.STDB_5_1_1(down_5)
+ down_5 = self.STDB_5_1_2(down_5)
+
+ up_4 = self.up_subm_5(down_5)
+ up_4 = up_4.replace_feature(torch.cat((up_4.features, skip_4.features), 1))
+ up_4 = self.STDB_4_2_1(up_4)
+
+ up_3 = self.up_subm_4(up_4)
+ up_3 = up_3.replace_feature(torch.cat((up_3.features, skip_3.features), 1))
+ up_3 = self.STDB_3_2_1(up_3)
+
+ up_2 = self.up_subm_3(up_3)
+ up_2 = up_2.replace_feature(torch.cat((up_2.features, skip_2.features), 1))
+ up_2 = self.STDB_2_2_1(up_2)
+
+ up_1 = self.up_subm_2(up_2)
+ up_1 = up_1.replace_feature(torch.cat((up_1.features, skip_1.features), 1))
+ up_1 = self.STDB_1_2_1(up_1)
+
+ return up_1
+
+class Point_head(nn.Module):
+ def __init__(self, voxel_feat_dim: int = 96, point_feat_dim: int = 32):
+ super().__init__()
+
+ self.input_dim = voxel_feat_dim + point_feat_dim
+
+ self.PPmodel_flow = nn.Sequential(
+ nn.Linear(self.input_dim, 32),
+ nn.BatchNorm1d(32),
+ nn.ReLU(),
+
+ nn.Linear(32, 3)
+ )
+
+ def forward_single(self, voxel_feat, voxel_coords, point_feat):
+
+ voxel_to_point_feat = voxel_feat[:, voxel_coords[:,2], voxel_coords[:,1], voxel_coords[:,0]].T
+ concated_point_feat = torch.cat([voxel_to_point_feat, point_feat],dim=-1)
+
+ flow = self.PPmodel_flow(concated_point_feat)
+
+ return flow
+
+ def forward(self, sparse_tensor, voxelizer_infos, pc0_point_feats_lst):
+
+ voxel_feats = sparse_tensor.dense()
+
+ flow_outputs = []
+ batch_idx = 0
+ for voxelizer_info in voxelizer_infos:
+ voxel_coords = voxelizer_info["voxel_coords"]
+ point_feat = pc0_point_feats_lst[batch_idx]
+ voxel_feat = voxel_feats[batch_idx, :]
+ flow = self.forward_single(voxel_feat, voxel_coords, point_feat)
+ batch_idx += 1
+ flow_outputs.append(flow)
+
+ return flow_outputs
\ No newline at end of file
diff --git a/src/models/flow4d.py b/src/models/flow4d.py
new file mode 100644
index 0000000..a9c2fdf
--- /dev/null
+++ b/src/models/flow4d.py
@@ -0,0 +1,98 @@
+
+"""
+This file is directly copied from:
+https://github.com/dgist-cvlab/Flow4D
+
+with slightly modification to have unified format with all benchmark.
+"""
+
+import torch.nn as nn
+import dztimer, torch
+
+from .basic import wrap_batch_pcs
+from .basic.flow4d_module import DynamicEmbedder_4D
+from .basic.flow4d_module import Network_4D, Seperate_to_3D, Point_head
+
+class Flow4D(nn.Module):
+ def __init__(self, voxel_size = [0.2, 0.2, 0.2],
+ point_cloud_range = [-51.2, -51.2, -2.2, 51.2, 51.2, 4.2],
+ grid_feature_size = [512, 512, 32],
+ num_frames = 5):
+ super().__init__()
+
+ point_output_ch = 16
+ voxel_output_ch = 16
+
+ self.num_frames = num_frames
+ print('voxel_size = {}, pseudo_dims = {}, input_num_frames = {}'.format(voxel_size, grid_feature_size, self.num_frames))
+
+ self.embedder_4D = DynamicEmbedder_4D(voxel_size=voxel_size,
+ pseudo_image_dims=[grid_feature_size[0], grid_feature_size[1], grid_feature_size[2], num_frames],
+ point_cloud_range=point_cloud_range,
+ feat_channels=point_output_ch)
+
+ self.network_4D = Network_4D(in_channel=point_output_ch, out_channel=voxel_output_ch)
+ self.seperate_feat = Seperate_to_3D(num_frames)
+ self.pointhead_3D = Point_head(voxel_feat_dim=voxel_output_ch, point_feat_dim=point_output_ch)
+
+ self.timer = dztimer.Timing()
+ self.timer.start("Total")
+
+ def load_from_checkpoint(self, ckpt_path):
+ ckpt = torch.load(ckpt_path, map_location="cpu")["state_dict"]
+ state_dict = {
+ k[len("model.") :]: v for k, v in ckpt.items() if k.startswith("model.")
+ }
+ print("\nLoading... model weight from: ", ckpt_path, "\n")
+ return self.load_state_dict(state_dict=state_dict, strict=False)
+
+ def forward(self, batch):
+ #t_deflow_start = time.time()
+ """
+ input: using the batch from dataloader, which is a dict
+ Detail: [pc0, pc1, pose0, pose1]
+ output: the predicted flow, pose_flow, and the valid point index of pc0
+ """
+
+ self.timer[0].start("Data Preprocess")
+ pcs_dict = wrap_batch_pcs(batch, self.num_frames)
+ self.timer[0].stop()
+
+ self.timer[1].start("4D_voxelization")
+ dict_4d = self.embedder_4D(pcs_dict)
+ pc01_tesnor_4d = dict_4d['4d_tensor']
+ pch1_3dvoxel_infos_lst = dict_4d['pch1_3dvoxel_infos_lst']
+ pc0_3dvoxel_infos_lst =dict_4d['pc0_3dvoxel_infos_lst']
+
+ pc0_point_feats_lst =dict_4d['pc0_point_feats_lst']
+ pc0_num_voxels = dict_4d['pc0_num_voxels']
+
+ pc1_3dvoxel_infos_lst =dict_4d['pc1_3dvoxel_infos_lst']
+ self.timer[1].stop()
+
+ self.timer[2].start("4D_backbone")
+ pc_all_output_4d = self.network_4D(pc01_tesnor_4d) #all = past, current, next 다 합친것
+ self.timer[2].stop()
+
+ self.timer[3].start("4D pc01 to 3D pc0")
+ pc0_last = self.seperate_feat(pc_all_output_4d)
+ assert pc0_last.features.shape[0] == pc0_num_voxels, 'voxel number mismatch'
+ self.timer[3].stop()
+
+ self.timer[4].start("3D_sparsetensor_to_point and head")
+ flows = self.pointhead_3D(pc0_last, pc0_3dvoxel_infos_lst, pc0_point_feats_lst)
+ self.timer[4].stop()
+
+ model_res = {
+ "flow": flows,
+ 'pose_flow': pcs_dict['pose_flows'],
+
+ "pc0_valid_point_idxes": [e["point_idxes"] for e in pc0_3dvoxel_infos_lst],
+ "pc0_points_lst": [e["points"] for e in pc0_3dvoxel_infos_lst] ,
+
+ "pc1_valid_point_idxes": [e["point_idxes"] for e in pc1_3dvoxel_infos_lst],
+ "pc1_points_lst": [e["points"] for e in pc1_3dvoxel_infos_lst],
+
+ 'pch1_valid_point_idxes': [e["point_idxes"] for e in pch1_3dvoxel_infos_lst] if pch1_3dvoxel_infos_lst != None else None,
+ }
+ return model_res
\ No newline at end of file
diff --git a/src/trainer.py b/src/trainer.py
index 97dc1e7..b05f592 100644
--- a/src/trainer.py
+++ b/src/trainer.py
@@ -268,7 +268,7 @@ def run_model_wo_ground_data(self, batch):
# NOTE (Qingwen): Since val and test, we will force set batch_size = 1
batch = {key: batch[key][0] for key in batch if len(batch[key])>0}
- res_dict = {key: res_dict[key][0] for key in res_dict if len(res_dict[key])>0}
+ res_dict = {key: res_dict[key][0] for key in res_dict if res_dict[key]!=None and len(res_dict[key])>0}
return batch, res_dict
def validation_step(self, batch, batch_idx):
diff --git a/src/utils/eval_metric.py b/src/utils/eval_metric.py
index 878e431..0f98ae1 100644
--- a/src/utils/eval_metric.py
+++ b/src/utils/eval_metric.py
@@ -13,13 +13,15 @@
import torch
import os, sys
import numpy as np
-from typing import Dict, Final, List, Tuple
+from typing import List, Tuple
from tabulate import tabulate
BASE_DIR = os.path.abspath(os.path.join( os.path.dirname( __file__ ), '../..' ))
sys.path.append(BASE_DIR)
from src.utils.av2_eval import compute_metrics, compute_bucketed_epe, CLOSE_DISTANCE_THRESHOLD
-
+import warnings
+warnings.simplefilter(action='ignore', category=FutureWarning)
+warnings.simplefilter(action='ignore', category=RuntimeWarning)
# EPE Three-way: Foreground Dynamic, Background Dynamic, Background Static
# leaderboard link: https://eval.ai/web/challenges/challenge-page/2010/evaluation