chore: refactor codes but it didn't compatible with previous pretrained weight

because of the model path is changed.

Author: Kin-Zhang

README.md

@@ -20,13 +20,13 @@ We directly follow our previous work [code structure](https://github.com/KTH-RPL
 - `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.]
 
-- `0_process.py`: process data with save dufomap, cluster labels inside file. Only needed for training.
+- `process.py`: process data with save dufomap, cluster labels inside file. Only needed once for training.
 
-- `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 upload to online leaderboard.
 
-- `3_vis.py` : For visualization of the results with a video.
+- `save.py` : For visualization of the results with a video.
 
 <details> <summary>🎁 <b>One repository, All methods!</b> </summary>
 <!-- <br> -->
@@ -101,15 +101,15 @@ python dataprocess/extract_av2.py --av2_type sensor --data_mode test --mask_dir
 Process train data for self-supervised learning. Only training data needs this step. [Runtime: Normally need 15 hours for my desktop, 3 hours for the cluster with five available nodes parallel running.]
 
 ```bash
-python 0_process.py --data_dir /home/kin/data/av2/preprocess_v2/sensor/train --scene_range 0,701
+python process.py --data_dir /home/kin/data/av2/preprocess_v2/sensor/train --scene_range 0,701
 ```
 
 ### Train the model
 
 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 1_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" "model.val_monitor=val/Dynamic/Mean"
 ```
 
 Or you can directly download the pre-trained weight from [Zenodo](https://zenodo.org/records/12751363/files/seflow_best.ckpt) and skip the training step. 
@@ -118,8 +118,8 @@ Or you can directly download the pre-trained weight from [Zenodo](https://zenodo
 
 You can also train the supervised baseline model in our paper with the following command. [Runtime: Around 10 hours in 4x A100 GPUs.] 
 ```bash
-python 1_train.py model=fastflow3d lr=2e-4 epochs=20 batch_size=16 loss_fn=ff3dLoss
-python 1_train.py model=deflow lr=2e-4 epochs=20 batch_size=16 loss_fn=deflowLoss
+python 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
 ```
 
 > [!NOTE]  
@@ -138,11 +138,11 @@ Since in training, we save all hyper-parameters and model checkpoints, the only
 wget https://zenodo.org/records/12751363/files/seflow_best.ckpt
 
 # it will directly prints all metric
-python 2_eval.py checkpoint=/home/kin/seflow_best.ckpt av2_mode=val
+python eval.py checkpoint=/home/kin/seflow_best.ckpt av2_mode=val
 
 # it will output the av2_submit.zip or av2_submit_v2.zip for you to submit to leaderboard
-python 2_eval.py checkpoint=/home/kin/seflow_best.ckpt av2_mode=test leaderboard_version=1
-python 2_eval.py checkpoint=/home/kin/seflow_best.ckpt av2_mode=test leaderboard_version=2
+python eval.py checkpoint=/home/kin/seflow_best.ckpt av2_mode=test leaderboard_version=1
+python eval.py checkpoint=/home/kin/seflow_best.ckpt av2_mode=test leaderboard_version=2
 ```
 
 And the terminal will output the command for you to submit the result to the online leaderboard. You can follow [this section for evalai](https://github.com/KTH-RPL/DeFlow?tab=readme-ov-file#2-evaluation).
@@ -154,7 +154,7 @@ Check all detailed result files (presented in our paper Table 1) in [this discus
 We provide a script to visualize the results of the model also. You can specify the checkpoint path and the data path to visualize the results. The step is quickly similar to evaluation.
 
 ```bash
-python 3_vis.py checkpoint=/home/kin/seflow_best.ckpt dataset_path=/home/kin/data/av2/preprocess_v2/sensor/vis
+python save.py checkpoint=/home/kin/seflow_best.ckpt dataset_path=/home/kin/data/av2/preprocess_v2/sensor/vis
 
 # The output of above command will be like:
 Model: DeFlow, Checkpoint from: /home/kin/model_zoo/v2/seflow_best.ckpt

assets/README.md

@@ -80,6 +80,6 @@ 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`.
+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,13 +25,13 @@ echo "Copy ${SOURCE} to ${DEST} Total time: ${elapsed} seconds"
 echo "Start training..."
 
 # ====> paper model = seflow_official
-# /proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/seflow/bin/python 1_train.py \
+# /proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/seflow/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=20 "model.target.num_iters=2" "model.val_monitor=val/Dynamic/Mean" \
 #     loss_fn=seflowLoss "add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}"
 
 # ====> leaderboard model = seflow_best
-/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/seflow/bin/python 1_train.py \
+/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/seflow/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-4 epochs=9 batch_size=16 "model.target.num_iters=2" "model.val_monitor=val/Dynamic/Mean" \
     loss_fn=seflowLoss "add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}"

assets/slurm/2_eval.sh

@@ -23,11 +23,11 @@ echo "Copy ${SOURCE} to ${DEST} Total time: ${elapsed} seconds"
 echo "Start training..."
 
 # ====> leaderboard model
-# /proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/seflow/bin/python 2_eval.py \
+# /proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/seflow/bin/python eval.py \
 #     wandb_mode=online dataset_path=/scratch/local/av2/sensor \
 #     checkpoint=/proj/berzelius-2023-154/users/x_qinzh/seflow/logs/wandb/seflow-10086990/checkpoints/epoch_19_seflow.ckpt \
 #     av2_mode=test save_res=True
 
-/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/seflow/bin/python 2_eval.py \
+/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/seflow/bin/python eval.py \
     wandb_mode=online dataset_path=/scratch/local/av2/sensor av2_mode=val \
     checkpoint=/proj/berzelius-2023-154/users/x_qinzh/seflow/logs/wandb/seflow-10086990/checkpoints/epoch_19_seflow.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.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.FastFlow3D
   voxel_size: ${voxel_size}
   point_cloud_range: ${point_cloud_range}