Add rerun and open3d multi visualizer (#3)

* feat: add o3d multi visualizer, allows parsing fwd and back

* feat: add rerun visualizer

* style(vis): open3d vis and add mode for people easily change in arg --mode.

* add example command for demo data

* fix(rerun): align rerun to fire also and enable mutiple input.
I agree it's too heavy since log times pc but only color difference. It looks like rerun rather than our problem.

* chore(vis): limit rerun frame id to avoid long time vis.
* update eval_metric to np.nan
* add rerun-sdk to sftool env.

* fix(o3d): align single MyVisualizer with Multi the same key.

---------

Co-authored-by: Kin <kinzhangglimmer@gmail.com>

Author: ajinkyakhoche

README.md

@@ -204,6 +204,14 @@ python tools/visualization.py --res_name 'seflow_best' --data_dir /home/kin/data
 
 https://github.com/user-attachments/assets/f031d1a2-2d2f-4947-a01f-834ed1c146e6
 
+Or another way to interact with [rerun](https://github.com/rerun-io/rerun) but please only vis scene by scene, not all at once.
+
+```bash
+python tools/visualization_rerun.py --data_dir /home/kin/data/av2/h5py/demo/train --res_name "['flow', 'deflow']"
+```
+
+https://github.com/user-attachments/assets/07e8d430-a867-42b7-900a-11755949de21
+
 
 ## Acknowledgement
 

dataprocess/README.md

@@ -16,7 +16,7 @@ If you want to use all datasets above, there is a specific process environment i
 
 ```bash
 conda env create -f envprocess.yaml
-conda activate dataprocess
+conda activate sftool
 # NOTE we need **manually reinstall numpy** (higher than 1.22)
 # * since waymo package force numpy==1.21.5, BUT!
 # * hdbscan w. numpy<1.22.0 will raise error: 'numpy.float64' object cannot be interpreted as an integer

envprocess.yaml

@@ -1,4 +1,4 @@
-name: dataprocess
+name: sftool
 channels:
   - conda-forge
   - pytorch
@@ -16,6 +16,7 @@ dependencies:
   - fire
   - hdbscan
   - s5cmd
+  - rerun-sdk
   - pip:
     - nuscenes-devkit
     - av2==0.2.1

src/utils/__init__.py

@@ -7,4 +7,15 @@ class bc:
     FAIL = '\033[91m'
     ENDC = '\033[0m'
     BOLD = '\033[1m'
-    UNDERLINE = '\033[4m'
+    UNDERLINE = '\033[4m'
+
+
+def hex_to_rgb(hex_color):
+    hex_color = hex_color.lstrip("#")
+    return tuple(int(hex_color[i:i + 2], 16) / 255.0 for i in (0, 2, 4))
+
+color_map_hex = ['#a6cee3', '#de2d26', '#1f78b4','#b2df8a','#33a02c','#fb9a99','#e31a1c','#fdbf6f','#ff7f00',\
+                 '#cab2d6','#6a3d9a','#ffff99','#b15928', '#8dd3c7','#ffffb3','#bebada','#fb8072','#80b1d3',\
+                 '#fdb462','#b3de69','#fccde5','#d9d9d9','#bc80bd','#ccebc5','#ffed6f']
+
+color_map = [hex_to_rgb(color) for color in color_map_hex]

src/utils/eval_metric.py

@@ -106,8 +106,8 @@ def __init__(self, class_names: List[str], speed_buckets: List[Tuple[float, floa
         ), f"speed_buckets must have at least one entry, got {len(self.speed_buckets)}"
 
         # By default, NaNs are not counted in np.nanmean
-        self.epe_storage_matrix = np.zeros((len(class_names), len(self.speed_buckets))) * np.NaN
-        self.speed_storage_matrix = np.zeros((len(class_names), len(self.speed_buckets))) * np.NaN
+        self.epe_storage_matrix = np.zeros((len(class_names), len(self.speed_buckets))) * np.nan
+        self.speed_storage_matrix = np.zeros((len(class_names), len(self.speed_buckets))) * np.nan
         self.count_storage_matrix = np.zeros(
             (len(class_names), len(self.speed_buckets)), dtype=np.int64
         )

src/utils/mics.py

@@ -234,13 +234,13 @@ def move_hue_on_S_axis(hues, factors):
 
 
 class HDF5Data:
-    def __init__(self, directory, flow_view=False, vis_name="flow"):
+    def __init__(self, directory, flow_view=False, vis_name=["flow"]):
         '''
         directory: the directory of the dataset
         t_x: how many past frames we want to extract
         '''
         self.flow_view = flow_view
-        self.vis_name = vis_name
+        self.vis_name = vis_name if isinstance(vis_name, list) else [vis_name]
         self.directory = directory
         with open(os.path.join(self.directory, 'index_total.pkl'), 'rb') as f:
             self.data_index = pickle.load(f)
@@ -285,7 +285,7 @@ def __getitem__(self, index):
             data_dict['pc0'] = f[key]['lidar'][:]
             data_dict['gm0'] = f[key]['ground_mask'][:]
             data_dict['pose0'] = f[key]['pose'][:]
-            for flow_key in [self.vis_name, 'dufo_label', 'label']:
+            for flow_key in self.vis_name + ['dufo_label', 'label']:
                 if flow_key in f[key]:
                     data_dict[flow_key] = f[key][flow_key][:]
 

src/utils/o3d_view.py

@@ -1,6 +1,6 @@
 '''
 # @date: 2023-1-26 16:38
-# @author: Qingwen Zhang  (https://kin-zhang.github.io/)
+# @author: Qingwen Zhang  (https://kin-zhang.github.io/), Ajinkya Khoche (https://ajinkyakhoche.github.io/)
 # Copyright (C) 2023-now, RPL, KTH Royal Institute of Technology
 # @detail:
 #  1. Play the data you want in open3d, and save the view control to json file.
@@ -13,12 +13,14 @@
 # 2024-08-23 21:41(Qingwen): remove totally on view setting from scratch but use open3d>=0.18.0 version for set_view from json text func.
 # 2024-04-15 12:06(Qingwen): show a example json text. add hex_to_rgb, color_map_hex, color_map (for color points if needed)
 # 2024-01-27 0:41(Qingwen): update MyVisualizer class, reference from kiss-icp: https://github.com/PRBonn/kiss-icp/blob/main/python/kiss_icp/tools/visualizer.py
+# 2024-09-10 (Ajinkya): Add MyMultiVisualizer class to view multiple windows at once, allow forward and backward playback, create bev square for giving a sense of metric scale.
 '''
 
 import open3d as o3d
 import os, time
 from typing import List, Callable
 from functools import partial
+import numpy as np
 
 def hex_to_rgb(hex_color):
     hex_color = hex_color.lstrip("#")
@@ -48,12 +50,12 @@ def __init__(self, view_file=None, window_title="Default", save_folder="logs/img
             "\t[SPACE] to pause/start\n"
             "\t[ESC/Q] to exit\n"
             "\t    [P] to save screen and viewpoint\n"
-            "\t    [N] to step\n"
+            "\t    [D] to step next\n"
         )
         self._register_key_callback(["Ā", "Q", "\x1b"], self._quit)
         self._register_key_callback(["P"], self._save_screen)
         self._register_key_callback([" "], self._start_stop)
-        self._register_key_callback(["N"], self._next_frame)
+        self._register_key_callback(["D"], self._next_frame)
 
     def show(self, assets: List):
         self.vis.clear_geometries()
@@ -113,6 +115,106 @@ def _save_screen(self, vis):
         vis.capture_screen_image(png_file)
         print(f"ScreenShot saved to: {png_file}, Please check it.")
 
+
+def create_bev_square(size=409.6, color=[68/255,114/255,196/255]):
+    # Create the vertices of the square
+    half_size = size / 2.0
+    vertices = np.array([
+        [-half_size, -half_size, 0],
+        [half_size, -half_size, 0],
+        [half_size, half_size, 0],
+        [-half_size, half_size, 0]
+    ])
+
+    # Define the square as a LineSet for visualization
+    lines = [[0, 1], [1, 2], [2, 3], [3, 0]]
+    colors = [color for _ in lines]  
+
+    line_set = o3d.geometry.LineSet(
+        points=o3d.utility.Vector3dVector(vertices),
+        lines=o3d.utility.Vector2iVector(lines)
+    )
+    line_set.colors = o3d.utility.Vector3dVector(colors)
+
+    return line_set
+
+class MyMultiVisualizer(MyVisualizer):
+    def __init__(self, view_file=None, flow_mode=['flow'], screen_width=2500, screen_height = 1375):
+        self.params = None
+        self.view_file = view_file
+        self.block_vis = True
+        self.play_crun = False
+        self.reset_bounding_box = True
+        self.playback_direction = 1 # 1:forward, -1:backward
+
+        self.vis = []
+        # self.o3d_vctrl = []
+
+        # Define width and height for each window
+        window_width = screen_width // 2
+        window_height = screen_height // 2
+        # Define positions for the four windows
+        epsilon = 150
+        positions = [
+            (0, 0),  # Top-left
+            (screen_width - window_width + epsilon, 0),  # Top-right
+            (0, screen_height - window_height + epsilon),  # Bottom-left
+            (screen_width - window_width + epsilon, screen_height - window_height + epsilon)  # Bottom-right
+        ]
+
+        for i, mode in enumerate(flow_mode):
+            window_title = f"view {'ground truth flow' if mode == 'flow' else f'{mode} flow'}, `SPACE` start/stop"
+            v = o3d.visualization.VisualizerWithKeyCallback()
+            v.create_window(window_name=window_title, width=window_width, height=window_height, left=positions[i%len(positions)][0], top=positions[i%len(positions)][1])
+            # self.o3d_vctrl.append(ViewControl(v.get_view_control(), view_file=view_file))
+            self.vis.append(v)
+
+        self._register_key_callback(["Ā", "Q", "\x1b"], self._quit)
+        self._register_key_callback([" "], self._start_stop)
+        self._register_key_callback(["D"], self._next_frame)
+        self._register_key_callback(["A"], self._prev_frame)
+        print(
+            f"\n{window_title.capitalize()} initialized. Press:\n"
+            "\t[SPACE] to pause/start\n"
+            "\t[ESC/Q] to exit\n"
+            "\t    [P] to save screen and viewpoint\n"
+            "\t    [D] to step next\n"
+            "\t    [A] to step previous\n"
+        )
+
+    def update(self, assets_list: List, clear: bool = True):
+        if clear:
+            [v.clear_geometries() for v in self.vis]
+
+        for i, assets in enumerate(assets_list):
+            [self.vis[i].add_geometry(asset, reset_bounding_box=False) for asset in assets]
+            self.vis[i].update_geometry(assets[-1])
+
+        if self.reset_bounding_box:
+            [v.reset_view_point(True) for v in self.vis]
+            if self.view_file is not None:
+                # [o.read_viewTfile(self.view_file) for o in self.o3d_vctrl]
+                [v.set_view_status(open(self.view_file).read()) for v in self.vis]
+            self.reset_bounding_box = False
+
+        [v.update_renderer() for v in self.vis]
+        while self.block_vis:
+            [v.poll_events() for v in self.vis]
+            if self.play_crun:
+                break
+        self.block_vis = not self.block_vis
+
+    def _register_key_callback(self, keys: List, callback: Callable):
+        for key in keys:
+            [v.register_key_callback(ord(str(key)), partial(callback)) for v in self.vis]
+    def _next_frame(self, vis):
+        self.block_vis = not self.block_vis
+        self.playback_direction = 1
+    def _prev_frame(self, vis):
+        self.block_vis = not self.block_vis
+        self.playback_direction = -1
+
+
 if __name__ == "__main__":
     json_content = """{
 	"class_name" : "ViewTrajectory",

tools/visualization.py

@@ -1,13 +1,21 @@
 """
 # Created: 2023-11-29 21:22
 # Copyright (C) 2023-now, RPL, KTH Royal Institute of Technology
-# Author: Qingwen Zhang  (https://kin-zhang.github.io/)
+# Author: Qingwen Zhang  (https://kin-zhang.github.io/), Ajinkya Khoche (https://ajinkyakhoche.github.io/)
 #
-# This file is part of DeFlow (https://github.com/KTH-RPL/DeFlow).
+# This file is part of OpenSceneFlow (https://github.com/KTH-RPL/OpenSceneFlow).
 # If you find this repo helpful, please cite the respective publication as 
 # listed on the above website.
 # 
 # Description: view scene flow dataset after preprocess.
+
+# CHANGELOG:
+# 2024-09-10 (Ajinkya): Add vis_multiple(), to visualize multiple flow modes at once.
+
+# Usage: (flow is ground truth flow, `other_name` is the estimated flow from the model)
+* python tools/visualization.py --data_dir /home/kin/data/av2/h5py/demo/train --res_name 'flow' --mode vis
+* python tools/visualization.py --data_dir /home/kin/data/av2/h5py/demo/train --res_name "['flow', 'deflow' , 'ssf']" --mode mul
+
 """
 
 import numpy as np
@@ -19,7 +27,7 @@
 BASE_DIR = os.path.abspath(os.path.join( os.path.dirname( __file__ ), '..' ))
 sys.path.append(BASE_DIR)
 from src.utils.mics import HDF5Data, flow_to_rgb
-from src.utils.o3d_view import MyVisualizer, color_map
+from src.utils.o3d_view import MyVisualizer, MyMultiVisualizer, color_map, create_bev_square
 
 
 VIEW_FILE = f"{BASE_DIR}/assets/view/av2.json"
@@ -64,7 +72,10 @@ def vis(
     res_name: str = "flow", # "flow", "flow_est"
     start_id: int = 0,
     point_size: float = 2.0,
+    mode: str = "vis",
 ):
+    if mode != "vis":
+        return
     dataset = HDF5Data(data_dir, vis_name=res_name, flow_view=True)
     o3d_vis = MyVisualizer(view_file=VIEW_FILE, window_title=f"view {'ground truth flow' if res_name == 'flow' else f'{res_name} flow'}, `SPACE` start/stop")
 
@@ -111,8 +122,83 @@ def vis(
             pcd.colors = o3d.utility.Vector3dVector(flow_color)
         o3d_vis.update([pcd, o3d.geometry.TriangleMesh.create_coordinate_frame(size=2)])
 
+
+def vis_multiple(
+    data_dir: str ="/home/kin/data/av2/preprocess/sensor/mini",
+    res_name: list = ["flow"],
+    start_id: int = 0,
+    point_size: float = 3.0,
+    tone: str = 'dark',
+    mode: str = "mul",
+):
+    if mode != "mul":
+        return
+    assert isinstance(res_name, list), "vis_multiple() needs a list as flow_mode"
+    dataset = HDF5Data(data_dir, vis_name=res_name, flow_view=True)
+    o3d_vis = MyMultiVisualizer(view_file=VIEW_FILE, flow_mode=res_name)
+
+    for v in o3d_vis.vis:
+        opt = v.get_render_option()
+        if tone == 'bright':
+            background_color = np.asarray([216, 216, 216]) / 255.0  # offwhite
+            # background_color = np.asarray([1, 1, 1])
+            pcd_color = [0.25, 0.25, 0.25]
+        elif tone == 'dark':
+            background_color = np.asarray([80/255, 90/255, 110/255])  # dark
+            pcd_color = [1., 1., 1.]
+
+        opt.background_color = background_color
+        opt.point_size = point_size
+
+    data_id = start_id
+    pbar = tqdm(range(0, len(dataset)))
+
+    while data_id >= 0 and data_id < len(dataset):
+        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']}")
+
+        pc0 = data['pc0']
+        gm0 = data['gm0']
+        pose0 = data['pose0']
+        pose1 = data['pose1']
+        ego_pose = np.linalg.inv(pose1) @ pose0
+
+        pose_flow = pc0[:, :3] @ ego_pose[:3, :3].T + ego_pose[:3, 3] - pc0[:, :3]
+
+        pcd_list = []
+        for mode in res_name:
+            pcd = o3d.geometry.PointCloud()
+            if mode in ['dufo_label', 'label']:
+                labels = data[mode]
+                pcd_i = o3d.geometry.PointCloud()
+                for label_i in np.unique(labels):
+                    pcd_i.points = o3d.utility.Vector3dVector(pc0[labels == label_i][:, :3])
+                    if label_i <= 0:
+                        pcd_i.paint_uniform_color([1.0, 1.0, 1.0])
+                    else:
+                        pcd_i.paint_uniform_color(color_map[label_i % len(color_map)])
+                    pcd += pcd_i
+            elif mode in data:
+                pcd.points = o3d.utility.Vector3dVector(pc0[:, :3])
+                flow = data[mode] - pose_flow # ego motion compensation here.
+                flow_color = flow_to_rgb(flow) / 255.0
+                is_dynamic = np.linalg.norm(flow, axis=1) > 0.1
+                flow_color[~is_dynamic] = pcd_color
+                flow_color[gm0] = pcd_color
+                pcd.colors = o3d.utility.Vector3dVector(flow_color)
+            pcd_list.append([pcd, create_bev_square(), 
+                            create_bev_square(size=204.8, color=[195/255,86/255,89/255]), 
+                            o3d.geometry.TriangleMesh.create_coordinate_frame(size=2)])
+        o3d_vis.update(pcd_list)
+
+        data_id += o3d_vis.playback_direction
+        pbar.update(o3d_vis.playback_direction)
+
+
 if __name__ == '__main__':
     start_time = time.time()
     # fire.Fire(check_flow)
     fire.Fire(vis)
+    fire.Fire(vis_multiple)
     print(f"Time used: {time.time() - start_time:.2f} s")