Vis cleanups (#18)

* Updated vis to use non-causal loader

* Fix mislabeled trackflow as supervised

* Added cleanups to vis support

docker/Dockerfile_cuda

@@ -41,7 +41,7 @@ RUN pip install -e setup
 
 RUN pip install FastGeodis==1.0.5 --no-build-isolation
 
-RUN pip install bucketed-scene-flow-eval==2.0.9
+RUN pip install bucketed-scene-flow-eval==2.0.10
 
 
 ENV PYTHONPATH=/project:/bucketed_scene_flow_eval:/:${PYTHONPATH}

paper_figures/perf_data/av2_test_volume_split/bucketed_epe/raw_to_leaderboard.py

@@ -52,7 +52,7 @@ def process(raw_file: Path, save_file: Path, is_supervised: bool):
 process(root_folder / "deflow_raw.json", root_folder / "deflow.json", is_supervised=True)
 process(root_folder / "fastflow3d_raw.json", root_folder / "fastflow3d.json", is_supervised=True)
 process(root_folder / "nsfp_raw.json", root_folder / "nsfp.json", is_supervised=False)
-process(root_folder / "trackflow_raw.json", root_folder / "trackflow.json", is_supervised=False)
+process(root_folder / "trackflow_raw.json", root_folder / "trackflow.json", is_supervised=True)
 process(root_folder / "zeroflow_1x_raw.json", root_folder / "zeroflow_1x.json", is_supervised=False)
 process(root_folder / "zeroflow_3x_raw.json", root_folder / "zeroflow_3x.json", is_supervised=False)
 process(root_folder / "zeroflow_5x_raw.json", root_folder / "zeroflow_5x.json", is_supervised=False)

paper_figures/perf_data/av2_test_volume_split/bucketed_epe/trackflow.json

@@ -2,7 +2,7 @@
     {
         "mean Dynamic": 0.21888200000000002,
         "mean Static": 0.06927174999999999,
-        "Is Supervised?": 0,
+        "Is Supervised?": 1,
         "BACKGROUND Static": 0.003346,
         "LARGE Static": 0.180365,
         "LARGE Dynamic": 0.35412,

paper_figures/plot_lib.py

@@ -157,7 +157,7 @@ def color(count, total_elements, intensity=1.3):
     start = 0.2
     stop = 0.7
 
-    colormap = matplotlib.cm.get_cmap(color_map())
+    colormap = plt.get_cmap(color_map())
     cm_subsection = np.linspace(start, stop, total_elements)
     # color = [matplotlib.cm.gist_earth(x) for x in cm_subsection][count]
     color = [colormap(x) for x in cm_subsection][count]

visualization/visualize_tto_itermediary_results.py

@@ -17,20 +17,49 @@
 import open3d as o3d
 
 
+@dataclass(kw_only=True)
+class SubSubsequenceInfo:
+    start_idx: int
+    size: int
+
+    def __repr__(self) -> str:
+        return f"Start: {self.start_idx} Size: {self.size}"
+
+
 @dataclass(kw_only=True)
 class VisState:
     flow_color: ColorEnum = ColorEnum.RED
     intermedary_result_idx: int = 0
+    subsubsequence: SubSubsequenceInfo | None = None
+
+    def __repr__(self) -> str:
+        ret_str = f"Vis State: Opt Step {self.intermedary_result_idx}"
+        if self.subsubsequence is not None:
+            ret_str += f" SubSubsequence: ({self.subsubsequence})"
+        return ret_str
 
 
 class ResultsVisualizer(BaseCallbackVisualizer):
 
-    def __init__(self, dataset: AbstractDataset, intermediary_results_folder: Path):
+    def __init__(
+        self,
+        dataset: AbstractDataset,
+        intermediary_results_folder: Path,
+        subsequence_length: int,
+        subsubsequence_length: int | None,
+    ):
 
-        self.vis_state = VisState()
         self.dataset = dataset
         self.intermediary_results_folder = intermediary_results_folder
         self._load_intermediary_flows()
+        self.subsequence_length = subsequence_length
+
+        self.vis_state = VisState(intermedary_result_idx=len(self.intermediary_files) - 1)
+        if subsubsequence_length is not None:
+            self.vis_state.subsubsequence = SubSubsequenceInfo(
+                start_idx=0,
+                size=subsubsequence_length,
+            )
 
         super().__init__(
             screenshot_path=Path()
@@ -39,6 +68,12 @@ def __init__(self, dataset: AbstractDataset, intermediary_results_folder: Path):
         )
 
     def _load_intermediary_flows(self):
+        """
+        Load / update a cache of all existing intermediary results.
+
+        This is called whenever the user changes the intermediary result index, in order to support
+        live updates of the visualizer from a running method.
+        """
         intermediary_results_folder = self.intermediary_results_folder
         dataset_idx = int(intermediary_results_folder.name.split("_")[-1])
         assert intermediary_results_folder.exists(), f"{intermediary_results_folder} does not exist"
@@ -57,6 +92,20 @@ def _load_ego_lidar(self) -> list[EgoLidarFlow]:
         return [EgoLidarFlow(*raw_flow) for raw_flow in raw_intermediary_flows]
 
     def _get_result_data(self) -> list[TimeSyncedSceneFlowFrame]:
+        """
+        Loads the dataset frame list and adds the intermediary ego flow to each frame.
+
+        Caches the dataset frame list and intermediary ego flows if the dataset index has not changed.
+        """
+
+        # Check to see if self._cached_dataset_idx is the same as self.dataset_idx.
+        # The field will not exist if it has not been set yet.
+        if hasattr(self, "_cached_dataset_idx") and self._cached_dataset_idx == self.dataset_idx:
+            assert hasattr(
+                self, "_cached_dataset_frame_list"
+            ), "Cached dataset frame list not found"
+            return self._cached_dataset_frame_list
+
         dataset_frame_list = self.dataset[self.dataset_idx]
         intermediary_ego_flows = self._load_ego_lidar()
 
@@ -68,10 +117,16 @@ def _get_result_data(self) -> list[TimeSyncedSceneFlowFrame]:
             # Add the intermediary ego flow to the frame info
             frame_info.flow = ego_flow
 
+        self._cached_dataset_idx = self.dataset_idx
+        self._cached_dataset_frame_list = dataset_frame_list
+
         return dataset_frame_list
 
     def _get_screenshot_path(self) -> Path:
-        return self.screenshot_path / f"{self.vis_state.intermedary_result_idx:08d}.png"
+        file_stem = f"{self.vis_state.intermedary_result_idx:08d}"
+        if self.vis_state.subsubsequence is not None:
+            file_stem += f"_subsubsequence_{self.vis_state.subsubsequence.start_idx:08d}"
+        return self.screenshot_path / f"{file_stem}.png"
 
     def _print_instructions(self):
         print("#############################################################")
@@ -83,41 +138,70 @@ def _print_instructions(self):
 
     def _register_callbacks(self, vis: o3d.visualization.VisualizerWithKeyCallback):
         super()._register_callbacks(vis)
-        # up arrow increase sequence_idx
-        vis.register_key_callback(265, self.increase_sequence_idx)
-        # down arrow decrease sequence_idx
-        vis.register_key_callback(264, self.decrease_sequence_idx)
+        # up arrow increase result_idx
+        vis.register_key_callback(265, self.increase_result_idx)
+        # down arrow decrease result_idx
+        vis.register_key_callback(264, self.decrease_result_idx)
+        # right arrow increase subsubsequence_idx
+        vis.register_key_callback(262, self.increase_subsubsequence_idx)
+        # left arrow decrease subsubsequence_idx
+        vis.register_key_callback(263, self.decrease_subsubsequence_idx)
         # E to jump to end of sequence
         vis.register_key_callback(ord("E"), self.jump_to_end_of_sequence)
 
-    def increase_sequence_idx(self, vis):
-        self._load_intermediary_flows()
+    def increase_result_idx(self, vis):
+        self._load_intermediary_flows()  # Refreshj cache to support live updates
         self.vis_state.intermedary_result_idx += 1
         if self.vis_state.intermedary_result_idx >= len(self.intermediary_files):
             self.vis_state.intermedary_result_idx = 0
         self.draw_everything(vis, reset_view=False)
 
-    def decrease_sequence_idx(self, vis):
-        self._load_intermediary_flows()
+    def decrease_result_idx(self, vis):
+        self._load_intermediary_flows()  # Refreshj cache to support live updates
         self.vis_state.intermedary_result_idx -= 1
         if self.vis_state.intermedary_result_idx < 0:
             self.vis_state.intermedary_result_idx = len(self.intermediary_files) - 1
         self.draw_everything(vis, reset_view=False)
 
+    def increase_subsubsequence_idx(self, vis):
+        if self.vis_state.subsubsequence is None:
+            return
+        self.vis_state.subsubsequence.start_idx += 1
+        end_idx = self.vis_state.subsubsequence.start_idx + self.vis_state.subsubsequence.size
+        if end_idx >= self.subsequence_length:
+            self.vis_state.subsubsequence.start_idx = 0
+        self.draw_everything(vis, reset_view=False)
+
+    def decrease_subsubsequence_idx(self, vis):
+        if self.vis_state.subsubsequence is None:
+            return
+        self.vis_state.subsubsequence.start_idx -= 1
+        if self.vis_state.subsubsequence.start_idx < 0:
+            self.vis_state.subsubsequence.start_idx = (
+                self.subsequence_length - self.vis_state.subsubsequence.size
+            )
+        self.draw_everything(vis, reset_view=False)
+
     def jump_to_end_of_sequence(self, vis):
         self._load_intermediary_flows()
         self.vis_state.intermedary_result_idx = len(self.intermediary_files) - 1
         self.draw_everything(vis, reset_view=False)
 
     def draw_everything(self, vis, reset_view=False):
         self.geometry_list.clear()
-        print(f"Vis State: {self.intermediary_files[self.vis_state.intermedary_result_idx].stem}")
+        print(f"Vis State: {self.vis_state}")
         frame_list = self._get_result_data()
         color_list = self._frame_list_to_color_list(len(frame_list))
 
-        for idx, flow_frame in enumerate(frame_list):
+        if self.vis_state.subsubsequence is not None:
+            start_idx = self.vis_state.subsubsequence.start_idx
+            size = self.vis_state.subsubsequence.size
+            frame_list = frame_list[start_idx : start_idx + size]
+            color_list = color_list[start_idx : start_idx + size]
+
+        for idx, (flow_frame, color) in enumerate(zip(frame_list, color_list)):
             pc = flow_frame.pc.global_pc
-            self.add_pointcloud(pc, color=color_list[idx])
+            self.add_pointcloud(pc, color=color)
             flowed_pc = flow_frame.pc.flow(flow_frame.flow).global_pc
 
             draw_color = self.vis_state.flow_color.rgb
@@ -131,23 +215,33 @@ def draw_everything(self, vis, reset_view=False):
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--dataset_name", type=str, default="Argoverse2CausalSceneFlow")
+    parser.add_argument("--dataset_name", type=str, default="Argoverse2NonCausalSceneFlow")
     parser.add_argument("root_dir", type=Path)
     parser.add_argument("subsequence_length", type=int)
+    parser.add_argument("--subsubsequence_length", type=int, default=-1)
     parser.add_argument("intermediary_results_folder", type=Path)
     args = parser.parse_args()
 
+    subsequence_length = args.subsequence_length
+
     dataset = construct_dataset(
         name=args.dataset_name,
         args=dict(
-            root_dir=args.root_dir, subsequence_length=args.subsequence_length, with_ground=False
+            root_dir=args.root_dir,
+            subsequence_length=subsequence_length,
+            with_ground=False,
         ),
     )
+    subsubsequence_length = args.subsubsequence_length
+    if subsubsequence_length < 0:
+        subsubsequence_length = None
     intermediary_results_folder = args.intermediary_results_folder
     # If the results folder is a pkl file, then grab the parent directory
     if intermediary_results_folder.is_file():
         intermediary_results_folder = intermediary_results_folder.parent
-    visualizer = ResultsVisualizer(dataset, intermediary_results_folder)
+    visualizer = ResultsVisualizer(
+        dataset, intermediary_results_folder, subsequence_length, subsubsequence_length
+    )
 
     visualizer.run()