feat(check_flow): visualize it easily with a inside func.

Author: Kin-Zhang

scripts/utils/mics.py

@@ -276,34 +276,23 @@ def __getitem__(self, index):
             scene_id, timestamp = self.data_index[index]
 
         key = str(timestamp)
+        data_dict = {
+            'scene_id': scene_id,
+            'timestamp': timestamp,
+        }
         with h5py.File(os.path.join(self.directory, f'{scene_id}.h5'), 'r') as f:
             # original data
-            pc0 = f[key]['lidar'][:]
-            gm0 = f[key]['ground_mask'][:]
-            pose0 = f[key]['pose'][:]
-
-            label = None
-            if 'label' in f[key]:
-                label = f[key]['label'][:]
+            data_dict['pc0'] = f[key]['lidar'][:]
+            data_dict['gm0'] = f[key]['ground_mask'][:]
+            data_dict['pose0'] = f[key]['pose'][:]
 
             if self.flow_view:
-                flow = f[key][self.vis_name][:]
+                data_dict[self.vis_name] = f[key][self.vis_name][:]
                 next_timestamp = str(self.data_index[index+1][1])
-                pose1 = f[next_timestamp]['pose'][:]
-        
-        data_dict = {
-            'scene_id': scene_id,
-            'timestamp': timestamp,
-            'pc0': pc0,
-            'gm0': gm0,
-            'pose0': pose0,
-            'label': label,
-        }
-        
-        if self.flow_view:
-            data_dict[self.vis_name] = flow
-            data_dict['pose1'] = pose1
-            
+                data_dict['pose1'] = f[next_timestamp]['pose'][:]
+                data_dict['pc1'] = f[next_timestamp]['lidar'][:]
+                data_dict['gm1'] = f[next_timestamp]['ground_mask'][:]
+
         return data_dict
 
 from av2.geometry.se3 import SE3

tests/scene_flow.py

@@ -1,7 +1,7 @@
 """
 # Created: 2023-11-29 21:22
 # Copyright (C) 2023-now, RPL, KTH Royal Institute of Technology
-# Author: Kin ZHANG  (https://kin-zhang.github.io/)
+# Author: Qingwen Zhang  (https://kin-zhang.github.io/)
 #
 # This work is licensed under the terms of the MIT license.
 # For a copy, see <https://opensource.org/licenses/MIT>.
@@ -22,6 +22,44 @@
 
 VIEW_FILE = f"{BASE_DIR}/assets/view/av2.json"
 
+def check_flow(
+    data_dir: str ="/home/kin/data/av2/preprocess/sensor/mini",
+    flow_mode: str = "flow", # "flow", "flow_est"
+    start_id: int = -1,
+    point_size: float = 3.0,
+):
+    dataset = HDF5Data(data_dir, vis_name=flow_mode, flow_view=True)
+    o3d_vis = MyVisualizer(view_file=VIEW_FILE, window_title=f"view {'ground truth flow' if flow_mode == 'flow' else f'{flow_mode} flow'}, `SPACE` start/stop")
+
+    opt = o3d_vis.vis.get_render_option()
+    opt.background_color = np.asarray([80/255, 90/255, 110/255])
+    opt.point_size = point_size
+
+    for data_id in (pbar := tqdm(range(0, len(dataset)))):
+        # for easy stop and jump to any id, and save same id always from 0.
+        if data_id < start_id and start_id != -1:
+            continue
+        data = dataset[data_id]
+        now_scene_id = data['scene_id']
+        pbar.set_description(f"id: {data_id}, scene_id: {now_scene_id}, timestamp: {data['timestamp']}")
+        gm0 = data['gm0']
+        pc0 = data['pc0'][~gm0]
+        
+        pcd = o3d.geometry.PointCloud()
+        pcd.points = o3d.utility.Vector3dVector(pc0[:, :3])
+        pcd.paint_uniform_color([1.0, 0.0, 0.0]) # red: pc0
+
+        pc1 = data['pc1']
+        pcd1 = o3d.geometry.PointCloud()
+        pcd1.points = o3d.utility.Vector3dVector(pc1[:, :3][~data['gm1']])
+        pcd1.paint_uniform_color([0.0, 1.0, 0.0]) # green: pc1
+
+        pcd2 = o3d.geometry.PointCloud()
+        # pcd2.points = o3d.utility.Vector3dVector(pc0[:, :3] + pose_flow) # if you want to check pose_flow
+        pcd2.points = o3d.utility.Vector3dVector(pc0[:, :3] + data[flow_mode][~gm0])
+        pcd2.paint_uniform_color([0.0, 0.0, 1.0]) # blue: pc0 + flow
+        o3d_vis.update([pcd, pcd1, pcd2, o3d.geometry.TriangleMesh.create_coordinate_frame(size=2)])
+
 def vis(
     data_dir: str ="/home/kin/data/av2/preprocess/sensor/mini",
     flow_mode: str = "flow", # "flow", "flow_est"
@@ -38,14 +76,13 @@ def vis(
     opt.point_size = point_size
 
     for data_id in (pbar := tqdm(range(0, len(dataset)))):
+        # for easy stop and jump to any id, and save same id always from 0.
+        if data_id < start_id and start_id != -1:
+            continue
         data = dataset[data_id]
         now_scene_id = data['scene_id']
         pbar.set_description(f"id: {data_id}, scene_id: {now_scene_id}, timestamp: {data['timestamp']}")
 
-        # for easy stop and jump to any id
-        if data_id < start_id and start_id != -1:
-            continue
-
         pc0 = data['pc0']
         gm0 = data['gm0']
         pose0 = data['pose0']
@@ -69,5 +106,6 @@ def vis(
 
 if __name__ == '__main__':
     start_time = time.time()
+    # fire.Fire(check_flow)
     fire.Fire(vis)
     print(f"Time used: {time.time() - start_time:.2f} s")