feat(data): update waymo process script.

scripts from running exp on our new SSL method SeFlow

Author: Kin-Zhang

README.md

@@ -5,18 +5,19 @@ DeFlow: Decoder of Scene Flow Network in Autonomous Driving
 [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/deflow-decoder-of-scene-flow-network-in/scene-flow-estimation-on-argoverse-2)](https://paperswithcode.com/sota/scene-flow-estimation-on-argoverse-2?p=deflow-decoder-of-scene-flow-network-in) 
 [![poster](https://img.shields.io/badge/ICRA24|Poster-6495ed?style=flat&logo=Shotcut&logoColor=wihte)](https://hkustconnect-my.sharepoint.com/:b:/g/personal/qzhangcb_connect_ust_hk/EXP_uXYmm_tItTWc8MafXHoB-1dVrMnvF1-lCzU1PXAvqQ?e=2FPfBS) 
 [![video](https://img.shields.io/badge/video-YouTube-FF0000?logo=youtube&logoColor=white)](https://youtu.be/bZ4uUv0nDa0)
+[![blog](https://img.shields.io/badge/Blog%7C%E7%9F%A5%E4%B9%8E%E4%B8%AD%E6%96%87-1772f6?style=flat&logo=Shotcut)](https://zhuanlan.zhihu.com/p/706514747) 
 
 Task: Scene Flow Estimation in Autonomous Driving. 
 
 🔥 2024/07/02: Check the self-supervised version in our new ECCV'24 [SeFlow](https://github.com/KTH-RPL/SeFlow). The 1st ranking in new leaderboard among self-supervise methods.
 
-Pre-trained weights for models are available in [Zenodo](https://zenodo.org/records/12173874) or [Onedrive link](https://hkustconnect-my.sharepoint.com/:f:/g/personal/qzhangcb_connect_ust_hk/Et85xv7IGMRKgqrVeJEVkMoB_vxlcXk6OZUyiPjd4AArIg?e=lqRGhx). 
+Pre-trained weights for models are available in [Zenodo](https://zenodo.org/records/12632962) or [Onedrive link](https://hkustconnect-my.sharepoint.com/:f:/g/personal/qzhangcb_connect_ust_hk/Et85xv7IGMRKgqrVeJEVkMoB_vxlcXk6OZUyiPjd4AArIg?e=lqRGhx). 
 Check usage in [2. Evaluation](#2-evaluation) or [3. Visualization](#3-visualization). 
 
 **Scripts** quick view in our scripts:
 
 - `dataprocess/extract_*.py` : pre-process data before training to speed up the whole training time. 
-  [Dataset we included now: Argoverse 2, more on the way: Waymo and Nuscenes, custom data.]
+  [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.
 
@@ -77,7 +78,9 @@ python 1_train.py model=fastflow3d lr=2e-6 epochs=50 batch_size=16
 python 1_train.py model=deflow lr=2e-6 epochs=50 batch_size=16
 ```
 
-To help community benchmarking, we provide our weights including fastflow3d, deflow [Onedrive link](https://hkustconnect-my.sharepoint.com/:f:/g/personal/qzhangcb_connect_ust_hk/Et85xv7IGMRKgqrVeJEVkMoB_vxlcXk6OZUyiPjd4AArIg?e=lqRGhx). These checkpoints also include parameters and status of that epoch inside it. If you are interested in weights of ablation studies, please contact us.
+To help community benchmarking, we provide our weights including fastflow3d, deflow [Zendo](https://zenodo.org/records/12632962). 
+These checkpoints also include parameters and status of that epoch inside it. If you are interested in weights of ablation studies, please contact us.
+Note: Please use these weights by following the term of use of the trained dataset (since weights are trained on these datasets) as [Argoverse 2 Term of Use](https://www.argoverse.org/about.html) mentioned: Using it under Non-Commercially (CC BY-NC-SA 4.0).
 
 ## 2. Evaluation
 
@@ -88,7 +91,7 @@ You can view Wandb dashboard for the training and evaluation results or [run/sub
 Since in training, we save all hyper-parameters and model checkpoints, the only thing you need to do is to specify the checkpoint path. Remember to set the data path correctly also.
 ```bash
 # downloaded pre-trained weight, or train by yourself
-wget https://zenodo.org/records/12173874/files/deflow_best.ckpt
+wget https://zenodo.org/records/12632962/files/deflow_best.ckpt
 
 python 2_eval.py checkpoint=/home/kin/deflow_best.ckpt av2_mode=val # it will directly prints all metric
 # it will output the av2_submit.zip or av2_submit_v2.zip for you to submit to leaderboard
@@ -121,7 +124,7 @@ We provide a script to visualize the results of the model. You can specify the c
 
 ```bash
 # downloaded pre-trained weight, or train by yourself
-wget https://zenodo.org/records/12173874/files/deflow_best.ckpt
+wget https://zenodo.org/records/12632962/files/deflow_best.ckpt
 
 python 3_vis.py checkpoint=/home/kin/deflow_best.ckpt dataset_path=/home/kin/data/av2/preprocess/sensor/vis
 
@@ -153,6 +156,12 @@ https://github.com/KTH-RPL/DeFlow/assets/35365764/9b265d56-06a9-4300-899c-96047a
   journal={arXiv preprint arXiv:2401.16122},
   year={2024}
 }
+@article{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},
+  journal={arXiv preprint arXiv:2407.01702},
+  year={2024}
+}
 ```
 
 This implementation is based on codes from several repositories. Thanks to these authors who kindly open-sourcing their work to the community. Please see our paper reference part to get more information. 

dataprocess/README.md

@@ -8,8 +8,8 @@ README for downloading and preprocessing the dataset. We includes waymo, argover
 
 We've updated the process dataset for:
 
-- [x] Argoverse 2.0: check [here](#argoverse-20). Involved from DeFlow
-- [ ] Waymo: done coding, public after review. Will be involved later by another paper.
+- [x] Argoverse 2.0: check [here](#argoverse-20). The process script Involved from [DeFlow](https://github.com/KTH-RPL/DeFlow).
+- [x] Waymo: check [here](#waymo-dataset). The process script was involved from [SeFlow](https://github.com/KTH-RPL/SeFlow).
 - [ ] nuScenes: done coding, public after review. Will be involved later by another paper.
 
 ## Download
@@ -83,13 +83,39 @@ gsutil -m cp -r "gs://waymo_open_dataset_scene_flow/valid" .
 gsutil -m cp -r "gs://waymo_open_dataset_scene_flow/train" .
 ```
 
+And flowlabel data can be downloaded here with ground segmentation by HDMap follow the same style of [ZeroFlow](https://github.com/kylevedder/zeroflow/blob/master/data_prep_scripts/waymo/extract_flow_and_remove_ground.py).
+
+You can download the processed map folder here to free yourself downloaded another type of data again:
+
+```bash
+wget https://zenodo.org/records/12632962/files/waymo_map.tar.gz
+tar -xvf waymo_map.tar.gz -C /home/kin/data/waymo/flowlabel
+# you will see there is a `map` folder in the `flowlabel` folder now.
+```
+
+#### Dataset frames
+
+| Dataset | # Total Scene | # Total Frames |
+| ------- | ------------- | -------------- |
+| train   | 799           | 155687         |
+| val     | 203           | 39381          |
+
 ## Process
 This directory contains the scripts to preprocess the datasets. 
 
 - `extract_av2.py`: Process the datasets in Argoverse 2.0.
 - `extract_nus.py`: Process the datasets in nuScenes.
 - `extract_waymo.py`: Process the datasets in Waymo.
 
+Example Running command:`
+```bash
+# av2:
+python dataprocess/extract_av2.py --av2_type sensor --data_mode train --argo_dir /home/kin/data/av2 --output_dir /home/kin/data/av2/preprocess
+
+# waymo:
+python dataprocess/extract_waymo.py --mode train --flow_data_dir /home/kin/data/waymo/flowlabel --map_dir /home/kin/data/waymo/flowlabel/map --output_dir /home/kin/data/waymo/preprocess  --nproc 48
+```
+
 All these preprocess scripts will generate the same format `.h5` file. The file contains the following in codes:
 
 File: `[*:logid].h5` file named in logid. Every timestamp is the key of group (f[key]).
@@ -120,6 +146,8 @@ After preprocessing, all data can use the same dataloader to load the data. As a
 Or you can run testing file to visualize the data. 
 
 ```bash
-python3 tests/preview_data.py --data_dir /home/kin/data/av2/preprocess/sensor/mini
-python3 tests/preview_data.py --data_dir /home/kin/data/nuScenes/preprocess/v1.0-mini
+# view gt flow
+python3 tests/scene_flow.py --data_dir /home/kin/data/av2/preprocess/sensor/mini --flow_mode flow
+
+python3 tests/scene_flow.py --data_dir /home/kin/data/waymo/preprocess/val --flow_mode flow
 ```

dataprocess/extract_av2.py

@@ -35,30 +35,17 @@
 from zipfile import ZipFile
 import pandas as pd
 
+import os, sys
+BASE_DIR = os.path.abspath(os.path.join( os.path.dirname( __file__ ), '..' ))
+sys.path.append(BASE_DIR)
+from scripts.utils.mics import create_reading_index
+
 BOUNDING_BOX_EXPANSION: Final = 0.2
 CATEGORY_TO_INDEX: Final = {
     **{"NONE": 0},
     **{k.value: i + 1 for i, k in enumerate(AnnotationCategories)},
 }
 
-def create_reading_index(data_dir: Path):
-    start_time = time.time()
-    data_index = []
-    for file_name in tqdm(os.listdir(data_dir), ncols=100, desc='Create reading index'):
-        if not file_name.endswith(".h5"):
-            continue
-        scene_id = file_name.split(".")[0]
-        timestamps = []
-        with h5py.File(data_dir/file_name, 'r') as f:
-            timestamps.extend(f.keys())
-        timestamps.sort(key=lambda x: int(x)) # make sure the timestamps are in order
-        for timestamp in timestamps:
-            data_index.append([scene_id, timestamp])
-
-    with open(data_dir/'index_total.pkl', 'wb') as f:
-        pickle.dump(data_index, f)
-        print(f"Create reading index Successfully, cost: {time.time() - start_time:.2f} s")
-
 def create_eval_mask(data_mode: str, output_dir_: Path, mask_dir: str):
     """
     Need download the official mask file run: `s5cmd --no-sign-request cp "s3://argoverse/tasks/3d_scene_flow/zips/*" .`