style: clean up the developed code previously.

* add readh5 and create_eval_pkl

Author: Kin-Zhang

README.md

@@ -198,7 +198,7 @@ wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/deflow_best.ckp
 ### Feed-Forward Self-Supervised Model Training
 
 Train Feed-forward SSL methods (e.g. SeFlow/SeFlow++/VoteFlow etc), we needed to:
-1) process auto-label process.
+1) process auto-label process for training. Check [dataprocess/README.md#self-supervised-process](dataprocess/README.md#self-supervised-process) for more details. We provide these inside the demo dataset already.
 2) specify the loss function, we set the config here for our best model in the leaderboard.
 
 #### SeFlow
@@ -257,7 +257,8 @@ python save.py model=fastnsf
 
 ## 3. Evaluation
 
-You can view Wandb dashboard for the training and evaluation results or upload result to online leaderboard.
+You can view Wandb dashboard for the training and evaluation results or upload result to online leaderboard. 
+<!-- Three-way EPE and Dynamic Bucket-normalized are evaluated within a 70x70m range (followed Argoverse 2 online leaderboard). No ground points are considered in the evaluation. -->
 
 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.
 

assets/slurm/dufolabel_sbatch.py

@@ -18,7 +18,7 @@ cd /proj/berzelius-2023-154/users/x_qinzh/OpenSceneFlow
 
 
 # data directory containing the extracted h5py files
-DATA_DIR="/proj/berzelius-2023-364/data/truckscenes/h5py/val"
+DATA_DIR="/proj/berzelius-2023-364/data/av2/h5py/sensor/train"
 
 TOTAL_SCENES=$(ls ${DATA_DIR}/*.h5 | wc -l)
 # Process every n-th frame into DUFOMap, no need to change at least for now.

assets/slurm/ssl-process.sh

@@ -247,3 +247,6 @@ Process train data for self-supervised learning. Only training data needs this s
 ```bash
 python process.py --data_dir /home/kin/data/av2/h5py/sensor/train --scene_range 0,701
 ```
+
+As some users must have multi-nodes for running, here I provide an example SLURM script to run the data process in parallel. 
+Check [assets/slurm/ssl-process.sh](../assets/slurm/ssl-process.sh) for more details.

dataprocess/README.md

@@ -186,7 +186,6 @@ def main(
     if not os.path.exists(gm_config_path) and run_gm:
         raise FileNotFoundError(f"Ground segmentation config file not found: {gm_config_path}. Please check folder")
 
-    
     data_path = Path(data_dir)
     dataset = HDF5Data(data_path) # single frame reading.
     all_scene_ids = list(dataset.scene_id_bounds.keys())

process.py

@@ -23,6 +23,7 @@
 import h5py, pickle, argparse
 from tqdm import tqdm
 import numpy as np
+from torchvision import transforms
 
 import os, sys
 BASE_DIR = os.path.abspath(os.path.join( os.path.dirname( __file__ ), '..' ))
@@ -185,8 +186,8 @@ def __call__(self, data_dict):
 class HDF5Dataset(Dataset):
     def __init__(self, directory, \
                 transform=None, n_frames=2, ssl_label=None, \
-                eval = False, eval_input_seq = False, leaderboard_version=1, \
-                vis_name='', flow_num=1):
+                eval = False, leaderboard_version=1, \
+                vis_name=''):
         '''
         Args:
             directory: the directory of the dataset, the folder should contain some .h5 file and index_total.pkl.
@@ -196,10 +197,8 @@ def __init__(self, directory, \
             * n_frames: the number of frames we use, default is 2: current (pc0), next (pc1); if it's more than 2, then it read the history from current.
             * ssl_label: if attr, it will read the dynamic cluster label. Otherwise, no dynamic cluster label in data dict.
             * eval: if True, use the eval index (only used it for leaderboard evaluation)
-            * eval_input_seq: I forgot what it is.... xox...
             * leaderboard_version: 1st or 2nd, default is 1. If '2', we will use the index_eval_v2.pkl from assets/docs.
             * vis_name: the data of the visualization, default is ''.
-            * flow_num: the number of future frames we read, default is 1. (pc0->pc1 flow)
         '''
         super(HDF5Dataset, self).__init__()
         self.directory = directory
@@ -209,12 +208,10 @@ def __init__(self, directory, \
             self.data_index = pickle.load(f)
 
         self.eval_index = False
-        self.eval_input_seq = eval_input_seq
         self.ssl_label = import_func(f"src.autolabel.{ssl_label}") if ssl_label is not None else None
         self.history_frames = n_frames - 2
         self.vis_name = vis_name if isinstance(vis_name, list) else [vis_name]
         self.transform = transform
-        self.flow_num = flow_num
 
         if eval:
             eval_index_file = os.path.join(self.directory, 'index_eval.pkl')
@@ -267,7 +264,7 @@ def __init__(self, directory, \
 
     def __len__(self):
         # return 100 # for testing
-        if self.eval_index and not self.eval_input_seq:
+        if self.eval_index:
             return len(self.eval_data_index)
         elif not self.eval_index and self.train_index is not None:
             return len(self.train_index)
@@ -278,25 +275,17 @@ def valid_index(self, index_):
         Check if the index is valid for the current mode and satisfy the constraints.
         """
         eval_flag = False
-        if self.eval_index and not self.eval_input_seq:
+        if self.eval_index:
             eval_index_ = index_
             scene_id, timestamp = self.eval_data_index[eval_index_]
             index_ = self.data_index.index([scene_id, timestamp])
             max_idx = self.scene_id_bounds[scene_id]["max_index"]
             if index_ >= max_idx:
                 _, index_ = self.valid_index(eval_index_ - 1)
             eval_flag = True
-        elif self.eval_index and self.eval_input_seq:
-            scene_id, timestamp = self.data_index[index_]
-            # to make sure we have continuous frames
-            if self.scene_id_bounds[scene_id]["max_index"] <= index_:
-                index_ = index_ - 1
-            scene_id, timestamp = self.data_index[index_]
-            eval_flag = True if [scene_id, timestamp] in self.eval_data_index else False
         elif self.train_index is not None:
             train_index_ = index_
             scene_id, timestamp = self.train_index[train_index_]
-            # FIXME: it works now, but self.flow_num is not possible in this case.
             max_idx = self.scene_id_bounds[scene_id]["max_index"]
             index_ = self.data_index.index([scene_id, timestamp])
             if index_ >= max_idx:
@@ -306,7 +295,7 @@ def valid_index(self, index_):
             max_idx = self.scene_id_bounds[scene_id]["max_index"]
             min_idx = self.scene_id_bounds[scene_id]["min_index"]
 
-            max_valid_index_for_flow = max_idx - self.flow_num
+            max_valid_index_for_flow = max_idx - 1
             min_valid_index_for_flow = min_idx + self.history_frames
             index_ = max(min_valid_index_for_flow, min(max_valid_index_for_flow, index_))
         return eval_flag, index_

src/dataset.py

@@ -0,0 +1,60 @@
+"""
+# Created: 2025-11-21 15:13
+# Copyright (C) 2023-now, RPL, KTH Royal Institute of Technology
+# Author: Qingwen Zhang  (https://kin-zhang.github.io/)
+
+# Description:
+    Create evaluation index pickle file from the total index pickle file.
+    - need have enough non-ground points (as some of waymo frames have data quality issues)
+    - sample every 5 frames for evaluation (followed the leaderboard setting) it can also save 5x validation time for optimization-based methods also.
+"""
+
+import os, fire, pickle, time
+import h5py, torch
+from tqdm import tqdm
+
+def create_evalpkl(
+    data_dir: str = "/home/kin/data/waymo/valid",
+    interval: int = 5,
+):
+    with open(os.path.join(data_dir, "index_total.pkl"), 'rb') as f:
+        total_index = pickle.load(f)
+
+    scene_id_bounds = {}
+    for idx, (scene_id, timestamp) in enumerate(total_index):
+        if scene_id not in scene_id_bounds:
+            scene_id_bounds[scene_id] = {
+                "min_timestamp": timestamp, "max_timestamp": timestamp,
+                "min_index": idx, "max_index": idx
+            }
+        else:
+            bounds = scene_id_bounds[scene_id]
+            if timestamp < bounds["min_timestamp"]:
+                bounds["min_timestamp"] = timestamp
+                bounds["min_index"] = idx
+            if timestamp > bounds["max_timestamp"]:
+                bounds["max_timestamp"] = timestamp
+                bounds["max_index"] = idx
+
+    # split the index by 5 - 5 frame, start with the fifth frame
+    eval_data_index = []
+    for scene_id, bounds in tqdm(scene_id_bounds.items(), desc="Creating eval index", total=len(scene_id_bounds), dynamic_ncols=True):
+        with h5py.File(os.path.join(data_dir, f'{scene_id}.h5'), 'r') as f:
+            for idx in range(bounds["min_index"] + interval*2, bounds["max_index"] - interval*2, interval):
+                scene_id, timestamp = total_index[idx]
+                key = str(timestamp)
+                pc = torch.tensor(f[key]['lidar'][:][:,:3])
+                gm = torch.tensor(f[key]['ground_mask'][:])
+                if pc[~gm].shape[0] < 10000:
+                    continue
+                eval_data_index.append(total_index[idx])
+    
+    # print(f"Demo: {eval_data_index[:10]}")
+    print(f"Total {len(eval_data_index)} frames for evaluation in {data_dir}.")
+    with open(os.path.join(data_dir, "index_eval.pkl"), 'wb') as f:
+        pickle.dump(eval_data_index, f)
+
+if __name__ == '__main__':
+    start_time = time.time()
+    fire.Fire(create_evalpkl)
+    print(f"Create reading index Successfully, cost: {time.time() - start_time:.2f} s")

tools/create_evalpkl.py

@@ -0,0 +1,35 @@
+"""
+# Created: 2023-12-31 22:19
+# LastEdit: 2024-01-12 18:46
+# Copyright (C) 2023-now, RPL, KTH Royal Institute of Technology
+# Author: Qingwen Zhang  (https://kin-zhang.github.io/)
+
+# Description:
+#   Quick Read the keys in an h5 file, print out their shapes and data types etc.
+
+# Example Running:
+    python tools/readh5.py --scene_path /home/kin/data/av2/h5py/sensor/test/0c6e62d7-bdfa-3061-8d3d-03b13aa21f68.h5
+"""
+
+import os
+os.environ["OMP_NUM_THREADS"] = "1"
+import fire, time, h5py
+
+def readh5key(
+    scene_path: str = "/home/kin/data/av2/h5py/sensor/test/0c6e62d7-bdfa-3061-8d3d-03b13aa21f68.h5"
+):
+    with h5py.File(scene_path, 'r') as f:
+        for cnt, k in enumerate(f.keys()):
+            if cnt % 2 == 1:
+                continue
+            print(f"id: {cnt}; Key (TimeStamp): {k}")
+            for sub_k in f[k].keys():
+                print(f"  Sub-Key: {sub_k}, Shape: {f[k][sub_k].shape}, Dtype: {f[k][sub_k].dtype}")
+            if cnt >= 10:
+                break
+        print(f"\nTotal {len(f.keys())} timestamps in the file.")
+
+if __name__ == '__main__':
+    start_time = time.time()
+    fire.Fire(readh5key)
+    print(f"\nTime used: {(time.time() - start_time)/60:.2f} mins")

tools/readh5.py

@@ -102,7 +102,7 @@ def vis(
         if res_name == 'raw': # no result, only show **raw point cloud**
             pcd.points = o3d.utility.Vector3dVector(pc0[:, :3])
             pcd.paint_uniform_color([1.0, 1.0, 1.0])
-        elif res_name in ['dufo_label', 'label']:
+        elif res_name in ['dufo', 'label']:
             labels = data[res_name]
             pcd_i = o3d.geometry.PointCloud()
             for label_i in np.unique(labels):
@@ -169,7 +169,7 @@ def vis_multiple(
         pcd_list = []
         for mode in res_name:
             pcd = o3d.geometry.PointCloud()
-            if mode in ['dufo_label', 'label']:
+            if mode in ['dufo', 'label']:
                 labels = data[mode]
                 pcd_i = o3d.geometry.PointCloud()
                 for label_i in np.unique(labels):

tools/visualization.py

@@ -104,7 +104,7 @@ def main(
             flow_color = np.tile(pcd_color, (pc0.shape[0], 1))
             flow_color[gm0] = ground_color
 
-            if mode in ['dufo_label', 'label']:
+            if mode in ['dufo', 'label']:
                 if mode in data:
                     labels = data[mode]
                     for label_i in np.unique(labels):