Initial

jmwang0117 · Jun 11, 2024 · f63627f · f63627f
1 parent cee8b0f
commit f63627f
Show file tree

Hide file tree

Showing 50 changed files with 3,555 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -1 +1,89 @@
-# OccRWKV
+# OccRWKV
+
+## OccRWKV: Rethinking Sparse Latent Representation for 3D Semantic Occupancy Prediction
+
+
+## Preperation
+
+### Prerequisites
+```
+conda create -n occ_rwkv python=3.10 -y
+conda install pytorch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 pytorch-cuda=12.1 -c pytorch -c nvidia -y
+pip install spconv-cu120
+pip install tensorboardX
+pip install dropblock
+pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.1.0+cu121.html
+```
+
+### Dataset
+
+Please download the Semantic Scene Completion dataset (v1.1) from the [SemanticKITTI website](http://www.semantic-kitti.org/dataset.html) and extract it.
+
+Or you can use [voxelizer](https://github.com/jbehley/voxelizer) to generate ground truths of semantic scene completion.
+
+The dataset folder should be organized as follows.
+```angular2
+SemanticKITTI
+├── dataset
+│   ├── sequences
+│   │  ├── 00
+│   │  │  ├── labels
+│   │  │  ├── velodyne
+│   │  │  ├── voxels
+│   │  │  ├── [OTHER FILES OR FOLDERS]
+│   │  ├── 01
+│   │  ├── ... ...
+```
+
+## Getting Start
+Clone the repository:
+```
+https://github.com/jmwang0117/OccRWKV.git
+```
+
+We provide training routine examples in the `cfgs` folder. Make sure to change the dataset path to your extracted dataset location in such files if you want to use them for training. Additionally, you can change the folder where the performance and states will be stored.
+* `config_dict['DATASET']['DATA_ROOT']` should be changed to the root directory of the SemanticKITTI dataset (`/.../SemanticKITTI/dataset/sequences`)
+* `config_dict['OUTPUT']['OUT_ROOT'] ` should be changed to desired output folder.
+
+### Train SSC-RS Net
+
+```
+$ cd <root dir of this repo>
+$ python train.py --cfg cfgs/DSC-Base.yaml --dset_root <path/dataset/root>
+```
+### Validation
+
+Validation passes are done during training routine. Additional pass in the validation set with saved model can be done by using the `validate.py` file. You need to provide the path to the saved model and the dataset root directory.
+
+```
+$ cd <root dir of this repo>
+$ python validate.py --weights </path/to/model.pth> --dset_root <path/dataset/root>
+```
+### Test
+
+Since SemantiKITTI contains a hidden test set, we provide test routine to save predicted output in same format of SemantiKITTI, which can be compressed and uploaded to the [SemanticKITTI Semantic Scene Completion Benchmark](http://www.semantic-kitti.org/tasks.html#ssc).
+
+We recommend to pass compressed data through official checking script provided in the [SemanticKITTI Development Kit](http://www.semantic-kitti.org/resources.html#devkit) to avoid any issue.
+
+You can provide which checkpoints you want to use for testing. We used the ones that performed best on the validation set during training. For testing, you can use the following command.
+
+```
+$ cd <root dir of this repo>
+$ python test.py --weights </path/to/model.pth> --dset_root <path/dataset/root> --out_path <predictions/output/path>
+```
+### Pretrained Model
+
+You can download the models with the scores below from this [Google drive link](https://drive.google.com/file/d/1-b3O7QS6hBQIGFTO-7qSG7Zb9kbQuxdO/view?usp=sharing),
+
+| Model  | Segmentation | Completion |
+|--|--|--|
+| SSC-RS | 24.2 | 59.7 |
+
+<sup>*</sup> Results reported to SemanticKITTI: Semantic Scene Completion leaderboard ([link](https://codalab.lisn.upsaclay.fr/competitions/7170\#results)).
+
+## Acknowledgement
+This project is not possible without multiple great opensourced codebases.
+* [spconv](https://github.com/traveller59/spconv)
+* [LMSCNet](https://github.com/cv-rits/LMSCNet)
+* [SSA-SC](https://github.com/jokester-zzz/SSA-SC)
+* [GASN](https://github.com/ItIsFriday/PcdSeg)
diff --git a/cfgs/2024.6.11.yaml b/cfgs/2024.6.11.yaml
@@ -0,0 +1,93 @@
+DATALOADER:
+  NUM_WORKERS: 4
+DATASET:
+  CONFIG_FILE: data/semantic-kitti.yaml
+  DATA_ROOT: data/SemanticKITTI/dataset/sequences
+  GRID_METERS:
+  - 0.2
+  - 0.2
+  - 0.2
+  LIMS:
+  - - 0
+    - 51.2
+  - - -25.6
+    - 25.6
+  - - -2
+    - 4.4
+  NCLASS: 20
+  SC_CLASS_FREQ:
+  - 7632350044
+  - 15783539
+  - 125136
+  - 118809
+  - 646799
+  - 821951
+  - 262978
+  - 283696
+  - 204750
+  - 61688703
+  - 4502961
+  - 44883650
+  - 2269923
+  - 56840218
+  - 15719652
+  - 158442623
+  - 2061623
+  - 36970522
+  - 1151988
+  - 334146
+  SIZES:
+  - 256
+  - 256
+  - 32
+  SS_CLASS_FREQ:
+  - 55437630
+  - 320797
+  - 541736
+  - 2578735
+  - 3274484
+  - 552662
+  - 184064
+  - 78858
+  - 240942562
+  - 17294618
+  - 170599734
+  - 6369672
+  - 230413074
+  - 101130274
+  - 476491114
+  - 9833174
+  - 129609852
+  - 4506626
+  - 1168181
+  TYPE: SemanticKITTI
+MODEL:
+  TYPE: DSC-AFC
+OPTIMIZER:
+  BASE_LR: 0.001
+  BETA1: 0.9
+  BETA2: 0.999
+  MOMENTUM: NA
+  TYPE: Adam
+  WEIGHT_DECAY: NA
+OUTPUT:
+  BEST_LOSS: 999999999999
+  BEST_METRIC: -999999999999
+  OUTPUT_PATH: ./outputs/DSC-AFC_SemanticKITTI_0611_000540
+  OUT_ROOT: ./outputs/
+SCHEDULER:
+  FREQUENCY: epoch
+  LR_POWER: 0.98
+  TYPE: power_iteration
+STATUS:
+  CONFIG: /home/jmwang/OccRWKV/cfgs/2024.6.11.yaml
+  LAST: ''
+  RESUME: false
+TRAIN:
+  BATCH_SIZE: 2
+  CHECKPOINT_PERIOD: 15
+  EPOCHS: 80
+  SUMMARY_PERIOD: 50
+VAL:
+  BATCH_SIZE: 2
+  SUMMARY_PERIOD: 20
diff --git a/cfgs/DSC-Base.yaml b/cfgs/DSC-Base.yaml
@@ -0,0 +1,88 @@
+DATALOADER:
+  NUM_WORKERS: 4
+DATASET:
+  CONFIG_FILE: data/semantic-kitti.yaml
+  DATA_ROOT: data/SemanticKITTI/dataset/sequences
+  GRID_METERS:
+  - 0.2
+  - 0.2
+  - 0.2
+  LIMS:
+  - - 0
+    - 51.2
+  - - -25.6
+    - 25.6
+  - - -2
+    - 4.4
+  NCLASS: 20
+  SC_CLASS_FREQ:
+  - 7632350044
+  - 15783539
+  - 125136
+  - 118809
+  - 646799
+  - 821951
+  - 262978
+  - 283696
+  - 204750
+  - 61688703
+  - 4502961
+  - 44883650
+  - 2269923
+  - 56840218
+  - 15719652
+  - 158442623
+  - 2061623
+  - 36970522
+  - 1151988
+  - 334146
+  SIZES:
+  - 256
+  - 256
+  - 32
+  SS_CLASS_FREQ:
+  - 55437630
+  - 320797
+  - 541736
+  - 2578735
+  - 3274484
+  - 552662
+  - 184064
+  - 78858
+  - 240942562
+  - 17294618
+  - 170599734
+  - 6369672
+  - 230413074
+  - 101130274
+  - 476491114
+  - 9833174
+  - 129609852
+  - 4506626
+  - 1168181
+  TYPE: SemanticKITTI
+MODEL:
+  TYPE: DSC-AFC
+OPTIMIZER:
+  BASE_LR: 0.001
+  BETA1: 0.9
+  BETA2: 0.999
+  MOMENTUM: NA
+  TYPE: Adam
+  WEIGHT_DECAY: NA
+OUTPUT:
+  OUT_ROOT: ./outputs/
+SCHEDULER:
+  FREQUENCY: epoch
+  LR_POWER: 0.98
+  TYPE: power_iteration
+STATUS:
+  RESUME: false
+TRAIN:
+  BATCH_SIZE: 2
+  CHECKPOINT_PERIOD: 15
+  EPOCHS: 80
+  SUMMARY_PERIOD: 50
+VAL:
+  BATCH_SIZE: 2
+  SUMMARY_PERIOD: 20