chore: refactor structure and updates latest best model for all.

Merge branch 'SeFlow' into DeFlow

Author: Kin-Zhang

README.md

@@ -21,11 +21,11 @@ Check usage in [2. Evaluation](#2-evaluation) or [3. Visualization](#3-visualiza
 - `dataprocess/extract_*.py` : pre-process data before training to speed up the whole training time. 
   [Dataset we included now: Argoverse 2 and Waymo, more on the way: Nuscenes, custom data.]
 
-- `1_train.py`: Train the model and get model checkpoints. Pls remember to check the config.
+- `train.py`: Train the model and get model checkpoints. Pls remember to check the config.
 
-- `2_eval.py` : Evaluate the model on the validation/test set. And also upload to online leaderboard.
+- `eval.py` : Evaluate the model on the validation/test set. And also output the zip file to upload to online leaderboard.
 
-- `3_vis.py` : For visualization of the results with a video.
+- `save.py` : Will save result into h5py file, using [tool/visualization.py] to show results with interactive window.
 
 
 <details> <summary>🎁 <b>One repository, All methods!</b> </summary>
@@ -105,8 +105,8 @@ All local benchmarking methods and ablation studies can be done through command
 Best fine-tuned model train with following command by other default config in [conf/config.yaml](conf/config.yaml) and [conf/model/deflow.yaml](conf/model/deflow.yaml), if you will set wandb_mode=online, maybe change all `entity="kth-rpl"` to your own account name.
 
 ```bash
-python 1_train.py model=deflow lr=2e-4 epochs=20 batch_size=16 loss_fn=deflowLoss
-python 1_train.py model=fastflow3d lr=2e-4 epochs=20 batch_size=16 loss_fn=ff3dLoss
+python train.py model=deflow lr=2e-4 epochs=20 batch_size=16 loss_fn=deflowLoss
+python train.py model=fastflow3d lr=4e-5 epochs=20 batch_size=16 loss_fn=ff3dLoss
 ```
 
 > [!NOTE]  
@@ -126,10 +126,10 @@ Since in training, we save all hyper-parameters and model checkpoints, the only
 # downloaded pre-trained weight, or train by yourself
 wget https://zenodo.org/records/12751363/files/deflow_best.ckpt
 
-python 2_eval.py checkpoint=/home/kin/deflow_best.ckpt av2_mode=val # it will directly prints all metric
+python eval.py checkpoint=/home/kin/deflow_best.ckpt av2_mode=val # it will directly prints all metric
 # it will output a command with absolute path of a zip file for you to submit to leaderboard
-python 2_eval.py checkpoint=/home/kin/deflow_best.ckpt av2_mode=test leaderboard_version=1
-python 2_eval.py checkpoint=/home/kin/deflow_best.ckpt av2_mode=test leaderboard_version=2
+python eval.py checkpoint=/home/kin/deflow_best.ckpt av2_mode=test leaderboard_version=1
+python eval.py checkpoint=/home/kin/deflow_best.ckpt av2_mode=test leaderboard_version=2
 ```
 
 Check all detailed result files (presented in our paper Table 1) in [this discussion](https://github.com/KTH-RPL/DeFlow/discussions/2).
@@ -159,16 +159,16 @@ We provide a script to visualize the results of the model. You can specify the c
 # downloaded pre-trained weight, or train by yourself
 wget https://zenodo.org/records/12751363/files/deflow_best.ckpt
 
-python 3_vis.py checkpoint=/home/kin/deflow_best.ckpt dataset_path=/home/kin/data/av2/preprocess/sensor/vis
+python save.py checkpoint=/home/kin/deflow_best.ckpt dataset_path=/home/kin/data/av2/preprocess/sensor/vis
 
 # Then terminal will tell you the command you need run. For example here is the output of the above:
 Model: DeFlow, Checkpoint from: /home/kin/deflow_best.ckpt
 We already write the estimate flow: deflow_best into the dataset, please run following commend to visualize the flow. Copy and paste it to your terminal:
-python tools/scene_flow.py --flow_mode 'deflow_best' --data_dir /home/kin/data/av2/preprocess/sensor/mini
+python tools/visualization.py --flow_mode 'deflow_best' --data_dir /home/kin/data/av2/preprocess/sensor/mini
 Enjoy! ^v^ ------ 
 
 # Then run the command in the terminal:
-python tools/scene_flow.py --flow_mode 'deflow_best' --data_dir /home/kin/data/av2/preprocess/sensor/mini
+python tools/visualization.py --res_name 'deflow_best' --data_dir /home/kin/data/av2/preprocess_v2/sensor/mini
 ```
 
 https://github.com/KTH-RPL/DeFlow/assets/35365764/9b265d56-06a9-4300-899c-96047a0da505

assets/README.md

@@ -76,20 +76,16 @@ cd DeFlow && git submodule update --init --recursive
 mamba env create -f assets/environment.yml
 ```
 
-Install mmcv lite version from us: [~5 mins] needed CUDA inside the env, echo ${PATH}
-
+CUDA package (nvcc compiler already installed through conda), the compile time is around 1-5 minutes:
 ```bash
-# change it if you use different cuda version (I tested 11.3, 11.4, 11.7, 11.8 all works)
-export PATH=/usr/local/cuda-11.7/bin:$PATH
-export LD_LIBRARY_PATH=/usr/local/cuda-11.7/lib64:$LD_LIBRARY_PATH
-
+mamba activate seflow
 cd assets/cuda/mmcv && python ./setup.py install && cd ../../..
 ```
 
 
 Checking the environment:
 ```bash
-mamba activate deflow
+mamba activate seflow
 python -c "import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.version.cuda)"
 python -c "import lightning.pytorch as pl"
 python -c "from assets.cuda.mmcv import Voxelization, DynamicScatter;print('success test on mmcv package')"
@@ -115,8 +111,8 @@ python -c "from assets.cuda.mmcv import Voxelization, DynamicScatter;print('succ
 
 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: scripts/network/models](../scripts/network/models). You can view deflow and fastflow3d to know how to implement a new model.
-3. Loss: All loss files can be found [here: scripts/network/loss_func.py](../scripts/network/loss_func.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 1_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`.
+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.

assets/slurm/1_train.sh

@@ -25,11 +25,11 @@ echo "Copy ${SOURCE} to ${DEST} Total time: ${elapsed} seconds"
 echo "Start training..."
 
 # ====> leaderboard model = [fastflow3d, deflow]
-# /proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/deflow/bin/python 1_train.py \
+# /proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/deflow/bin/python train.py \
 #     slurm_id=$SLURM_JOB_ID wandb_mode=online train_data=/scratch/local/av2/sensor/train val_data=/scratch/local/av2/sensor/val \
 #     num_workers=16 model=deflow lr=2e-6 epochs=50 batch_size=10 loss_fn=deflowLoss
 
-# /proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/deflow/bin/python 1_train.py \
+# /proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/deflow/bin/python train.py \
 #     slurm_id=$SLURM_JOB_ID wandb_mode=online train_data=/scratch/local/av2/sensor/train val_data=/scratch/local/av2/sensor/val \
 #     num_workers=16 model=fastflow3d lr=2e-6 epochs=50 batch_size=16 loss_fn=ff3dLoss
 

assets/slurm/2_eval.sh

@@ -30,14 +30,14 @@ echo "Start training..."
 #     av2_mode=test save_res=True
 
 
-/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/deflow/bin/python 2_eval.py \
+/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/deflow/bin/python eval.py \
     wandb_mode=online dataset_path=/scratch/local/av2/sensor av2_mode=val \
     checkpoint=/proj/berzelius-2023-154/users/x_qinzh/deflow/logs/wandb/fastflow3d-10086990/checkpoints/epoch_49_fastflow3d.ckpt
 
-/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/deflow/bin/python 2_eval.py \
+/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/deflow/bin/python eval.py \
     wandb_mode=online dataset_path=/scratch/local/av2/sensor av2_mode=val \
     checkpoint=/proj/berzelius-2023-154/users/x_qinzh/deflow/logs/wandb/fastflow3d-10088873/checkpoints/epoch_49_fastflow3d.ckpt
 
-/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/deflow/bin/python 2_eval.py \
+/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/deflow/bin/python eval.py \
     wandb_mode=online dataset_path=/scratch/local/av2/sensor av2_mode=val \
     checkpoint=/proj/berzelius-2023-154/users/x_qinzh/deflow/logs/wandb/fastflow3d-10088874/checkpoints/epoch_49_fastflow3d.ckpt

assets/slurm/dufolabel_sbatch.py

@@ -40,7 +40,7 @@ def main(
 cd /proj/berzelius-2023-154/users/x_qinzh/seflow
 export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/proj/berzelius-2023-154/users/x_qinzh/mambaforge/lib
 
-/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/seflow/bin/python 0_process.py \\
+/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/seflow/bin/python process.py \\
     --data_dir {data_dir} \\
     --interval {interval} \\
     --scene_range {scene_range[0]},{scene_range[1]}

assets/tests/chamferdis_speed_test.py

@@ -72,7 +72,7 @@ def faiss_knn(pc1, pc2):
 
     if PYTORCH3D_TEST:
         print("------ START Pytorch3d Chamfer Distance Cal ------")
-        from scripts.network.models.unsfp.model import my_chamfer_fn
+        from src.models.basic.nsfp_module import my_chamfer_fn
         start_time = time.time()
         loss0, _ = my_chamfer_fn(pc0.unsqueeze(0), pc1.unsqueeze(0), truncate_dist=False)
 

assets/tests/hdbscan_speed.py

@@ -18,7 +18,7 @@
 from sklearn.cluster import HDBSCAN
 import hdbscan as cpu_hdbscan
 
-from scripts.utils.o3d_view import MyVisualizer, color_map
+from src.utils.o3d_view import MyVisualizer, color_map
 import open3d as o3d
 
 MAX_AXIS_RANGE = 60

conf/model/deflow.yaml

@@ -1,7 +1,7 @@
 name: deflow
 
 target:
-  _target_: scripts.network.models.deflow.DeFlow
+  _target_: src.models.DeFlow
   decoder_option: gru # choices: [linear, gru]
   num_iters: 4
   voxel_size: ${voxel_size}

conf/model/fastflow3d.yaml

@@ -1,8 +1,7 @@
 name: fastflow3d
 
 target:
-  _target_: scripts.network.models.fastflow3d.FastFlow3D
-  bottleneck_head: False
+  _target_: src.models.FastFlow3D
   voxel_size: ${voxel_size}
   point_cloud_range: ${point_cloud_range}