release code

AlibabaResearch · Jul 22, 2022 · 6458503 · 6458503
commit 6458503
Show file tree

Hide file tree

Showing 261 changed files with 13,353 additions and 0 deletions.
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1,10 @@
+venv
+.idea
+__pycache__/
+.vscode
+checkpoints/*
+data/datasets/*
+lib/pointnet2/__pycache__
+lib/pointnet2/build
+lib/pointnet2/dist
+lib/pointnet2/pointnet2.egg-info
diff --git a/README.md b/README.md
@@ -0,0 +1,85 @@
+# RCP: Recurrent Closest Point for Scene Flow Estimation on 3D Point Clouds
+This is the official PyTorch implementation code for RCP. For technical details, please refer to:
+
+**RCP: Recurrent Closest Point for Scene Flow Estimation on 3D Point Clouds** <br />
+Xiaodong Gu, Chengzhou Tang, Weihao Yuan, Zuozhuo Dai, Siyu Zhu, Ping Tan <br />
+**[[Paper](https://arxiv.org/abs/2205.11028)]** <br />
+
+![frames](assets/rcp_arch_compressed.jpg)
+
+
+
+## Installation  
++ Install python dependencies lib:
+```
+pip install -r requirements.txt  
+```
++ Install PointNet2 CPP lib:
+```
+cd lib/pointnet2
+python3 setup.py install
+```
+
+## Datasets
+We follow [HPLFlowNet](https://github.com/laoreja/HPLFlowNet) preprocessing methods:
+
++ FlyingThings3D: Download and unzip the "Disparity", "Disparity Occlusions", "Disparity change", "Optical flow", "Flow Occlusions" for DispNet/FlowNet2.0 dataset subsets from the [FlyingThings3D website](https://lmb.informatik.uni-freiburg.de/resources/datasets/SceneFlowDatasets.en.html) (we used the paths from [this file](https://lmb.informatik.uni-freiburg.de/data/FlyingThings3D_subset/FlyingThings3D_subset_all_download_paths.txt), now they added torrent downloads) . They will be upzipped into the same directory, RAW_DATA_PATH. Then run the following script for 3D reconstruction:
+```
+python data/preprocess/process_flyingthings3d_subset.py --raw_data_path ${RAW_DATA_PATH} --save_path ${SAVE_PATH}/FlyingThings3D_subset_processed_35m --only_save_near_pts
+```
+
++ KITTI: Download and unzip [KITTI Scene Flow Evaluation 2015](http://www.cvlibs.net/download.php?file=data_scene_flow.zip) to directory RAW_DATA_PATH. Run the following script for 3D reconstruction:
+```
+python data/preprocess/process_kitti.py ${RAW_DATA_PATH} ${SAVE_PATH}/KITTI_processed_occ_final
+```
+
+## Training
+
++ Fully-supervised training:
+```bash
+python run.py -c configs/train/rcp_sup_pre.yaml
+python run.py -c configs/train/rcp_sup_ft.yaml --pre_ckpt ${pretrained_ckpt}
+```
+
++ Self-supervised training:
+```bash
+python run.py -c configs/train/rcp_self_pre.yaml
+python run.py -c configs/train/rcp_self_ft.yaml --pre_ckpt ${pretrained_ckpt}
+```
+
+## Evaluation
+
++ Evaluate on FlyingThings3D
+```
+python run.py -c configs/test/rcp_test.yaml --test_ckpt ${ft_ckpt}
+```
++ Evaluate on KITTI
+```
+python run.py -c configs/test/rcp_test_kitti.yaml --test_ckpt ${ft_ckpt}
+```
+
+## Pretrained Models
+[Dowload Link](https://virutalbuy-public.oss-cn-hangzhou.aliyuncs.com/share/rcp/rcp.ckpt)
+
+| Datasets | EPE3D | Acc3DS | AccDR | Outliers3D |
+| :--- | :---: | :---: | :---: |  :---: |
+| FlyingThings3D | 0.0403 | 0.8567 | 0.9635 | 0.1976 |
+| KITTI | 0.0481 | 0.8491 | 0.9448 | 0.1228 |
+
+
+## Citation
+If you find this code useful in your research, please cite:
+
+```
+@inproceedings{gu2022rcp,
+  title={RCP: Recurrent Closest Point for Point Cloud},
+  author={Gu, Xiaodong and Tang, Chengzhou and Yuan, Weihao and Dai, Zuozhuo and Zhu, Siyu and Tan, Ping},
+  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+  pages={8216--8226},
+  year={2022}
+}
+```
+
+## Acknowledgements
+Some code are borrowed from [Flowstep3d](https://github.com/yairkit/flowstep3d), [FLOT](https://github.com/valeoai/FLOT), [flownet3d_Pytorch](https://github.com/hyangwinter/flownet3d_pytorch) and [HPLFlowNet](https://github.com/laoreja/HPLFlowNet).
+Thanks for these great projects.
diff --git a/assets/rcp_arch.jpg b/assets/rcp_arch.jpg
diff --git a/assets/rcp_arch_compressed.jpg b/assets/rcp_arch_compressed.jpg
diff --git a/configs/test/rcp_test.yaml b/configs/test/rcp_test.yaml
@@ -0,0 +1,94 @@
+train: False
+
+task: sceneflow
+
+model_params: 
+  model_name: 'SF_RCP'
+  npoint: 8192
+  use_insrance_norm: False
+  use_curvature: True 
+
+exp_params:
+  disable_plots: True
+  batch_size: 8
+  num_workers: 32
+  train_iters: 8
+  val_iters: 8
+  test_iters: 8
+
+  data:
+    overfit_samples: null
+    train_dataset: null
+    train_data_root: null
+    test_dataset: flyingthings3d
+    test_data_root: /data/datasets/scene_flow/flyingthings3d
+    num_points: 8192
+    full: True
+    allow_less_points: False
+
+    data_process:
+      DEPTH_THRESHOLD: 35.
+      NO_CORR: True
+      IND_DIR: null
+
+    data_augmentation:
+      aug_together:
+        degree_range: 0.1745329252
+        shift_range: 1.
+        scale_low: 0.95
+        scale_high: 1.05
+        jitter_sigma: 0.01
+        jitter_clip: 0.00
+      aug_pc2:
+        degree_range: 0.
+        xy_shift_range: 0.
+        deg_max: False
+        shift_range: 0.2
+        jitter_sigma: 0.01
+        jitter_clip: 0.00
+      val_aug_pc2: null
+
+  loss:
+    loss_type: 'sv_l1_reg_v1'
+    w_data: [0.9]
+    w_smoothness: [0.1]
+    smoothness_loss_params:
+      w_knn: 3.
+      w_ball_q: 1.
+      knn_loss_params:
+        k: 16
+        radius: 0.5
+        loss_norm: 1
+      ball_q_loss_params:
+        k: 24
+        radius: 0.5
+        loss_norm: 1
+  optimizer:
+    type: 'AdamW'
+    lr: 0.000125
+    weight_decay: 0.01
+
+  scheduler:
+    type: 'MultiStepLR'
+    milestones: [1, 2, 3]
+    gamma: 0.5
+
+trainer_params:
+  deterministic: False # slower. reproduce.
+  strategy: 'ddp'
+  num_nodes: 1
+  gpus: 8
+  max_epochs: 90
+  num_sanity_val_steps: 3
+  check_val_every_n_epoch: 1
+  sync_batchnorm: True
+
+  resume_from_checkpoint: None
+
+logging_params:
+  ckpt:
+    save_dir: 'checkpoints/rcp_sup_ft'
+    top_k: 3
+    monitor: 'val_epe'
+
+  log: false
diff --git a/configs/test/rcp_test_kitti.yaml b/configs/test/rcp_test_kitti.yaml
@@ -0,0 +1,91 @@
+train: False
+task: "sceneflow"
+
+model_params:
+  model_name: 'SF_RCP'
+  npoint: 8192
+  use_insrance_norm: False 
+  use_curvature: True 
+
+exp_params:
+  disable_plots: True
+  batch_size: 1
+  num_workers: 32
+  train_iters: 16
+  val_iters: 16
+  test_iters: 16
+
+  data:
+    overfit_samples: null
+    train_dataset: null
+    train_data_root: null
+    test_dataset: kitti
+    test_data_root: /data/datasets/scene_flow/kitti
+    num_points: 8192
+    full: True
+    allow_less_points: False
+
+    data_process:
+      DEPTH_THRESHOLD: 35.
+      NO_CORR: True
+      IND_DIR: null
+
+    data_augmentation:
+      aug_together:
+        degree_range: 0.1745329252
+        shift_range: 1.
+        scale_low: 0.95
+        scale_high: 1.05
+        jitter_sigma: 0.01
+        jitter_clip: 0.00
+      aug_pc2:
+        degree_range: 0.
+        xy_shift_range: 0.
+        deg_max: False
+        shift_range: 0.2
+        jitter_sigma: 0.01
+        jitter_clip: 0.00
+      val_aug_pc2: null
+
+  loss:
+    loss_type: 'sv_l1_reg_v1'
+    w_data: [0.9]
+    w_smoothness: [0.1]
+    smoothness_loss_params:
+      w_knn: 3.
+      w_ball_q: 1.
+      knn_loss_params:
+        k: 16
+        radius: 0.5
+        loss_norm: 1
+      ball_q_loss_params:
+        k: 24
+        radius: 0.5
+        loss_norm: 1
+  optimizer:
+    type: 'AdamW'
+    lr: 0.000125
+    weight_decay: 0.01
+
+  scheduler:
+    type: 'MultiStepLR'
+    milestones: [1, 2, 3]
+    gamma: 0.5
+
+trainer_params:
+  deterministic: False # slower. reproduce.
+  strategy: 'ddp'
+  num_nodes: 1
+  gpus: 8
+  max_epochs: 90
+  num_sanity_val_steps: 3
+  check_val_every_n_epoch: 1
+  sync_batchnorm: True
+  resume_from_checkpoint: None
+
+logging_params:
+  ckpt:
+    save_dir: 'checkpoints/rcp_sup_ft'
+    top_k: 3
+    monitor: 'val_epe'
+  log: false
diff --git a/configs/train/rcp_self_ft.yaml b/configs/train/rcp_self_ft.yaml
@@ -0,0 +1,97 @@
+train: True
+
+task: 'sceneflow'
+model_params:
+  model_name: 'SF_RCP'
+  npoint: 8192
+  use_insrance_norm: False
+  use_curvature: True
+
+exp_params:
+  pre_trained_weights_checkpoint: None
+  disable_plots: True
+  batch_size: 1
+  num_workers: 32
+  train_iters: 8
+  val_iters: 8
+  loss_iters_w: [0.5, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2]
+
+  data:
+    overfit_samples: null
+    train_dataset: flyingthings3d
+    train_data_root: /data/datasets/scene_flow/flyingthings3d
+    test_dataset: null
+    test_data_root: null
+    num_points: 8192
+    full: True
+    allow_less_points: False
+
+    data_process:
+      DEPTH_THRESHOLD: 35.
+      NO_CORR: True
+      IND_DIR: null
+
+    data_augmentation:
+      aug_together:
+        degree_range: 0.1745329252
+        shift_range: 1.
+        scale_low: 0.95
+        scale_high: 1.05
+        jitter_sigma: 0.01
+        jitter_clip: 0.00
+      aug_pc2:
+        degree_range: 0.
+        xy_shift_range: 0.
+        deg_max: False
+        shift_range: 0.2
+        jitter_sigma: 0.01
+        jitter_clip: 0.00
+      val_aug_pc2: null
+
+  loss:
+    loss_type: 'unsup_l1'
+    w_data: [0.65, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7]
+    w_smoothness: [0.35, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3]
+    w_curvature:  [0.1, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01]
+    smoothness_loss_params:
+      w_knn: 4.
+      w_ball_q: 1.
+      knn_loss_params:
+        k: 16
+        radius: 0.25
+        loss_norm: 1
+      ball_q_loss_params:
+        k: 48
+        radius: 0.5
+        loss_norm: 1
+    chamfer_loss_params:
+      loss_norm: 2
+      k: 1
+  optimizer:
+    type: 'AdamW'
+    lr: 0.002
+    weight_decay: 0.01
+
+  scheduler:
+    type: 'MultiStepLR'
+    milestones: [3, 6, 9]
+    gamma: 0.5
+
+trainer_params:
+  deterministic: False # slower. reproduce.
+  strategy: 'ddp'
+  num_nodes: 1
+  gpus: 8
+  max_epochs: 15
+  num_sanity_val_steps: 3
+  check_val_every_n_epoch: 1
+  resume_from_checkpoint: null
+  sync_batchnorm: True
+
+logging_params:
+  ckpt:
+    save_dir: 'checkpoints/rcp_selfsup_ft'
+    top_k: 3
+    monitor: 'val_epe'
+
+  log: false