Added optimization step visualization

configs/fast_nsf/argo/val_debug.py

@@ -3,4 +3,6 @@
 test_dataset_root = "/efs/argoverse2_small/val/"
 save_output_folder = "/efs/argoverse2_small/val_fast_nsf_flow_replicate/"
 
-test_dataset = dict(args=dict(root_dir=test_dataset_root, split=dict(split_idx=56, num_splits=314)))
+test_dataset = dict(
+    args=dict(root_dir=test_dataset_root, split=dict(split_idx=56 // 2, num_splits=314 // 2))
+)

models/optimization/test_time_optimizer_loop.py

@@ -8,6 +8,7 @@
 import numpy as np
 from pytorch_lightning.loggers import Logger
 from pathlib import Path
+from bucketed_scene_flow_eval.utils import save_pickle
 
 
 class OptimizationLoop:
@@ -28,8 +29,29 @@ def __init__(
         self.min_delta = min_delta
         self.compile = compile
 
-    def _save_intermediary_results(self) -> None:
-        pass
+    def _save_intermediary_results(
+        self,
+        model: BaseNeuralRep,
+        problem: BucketedSceneFlowInputSequence,
+        logger: Logger,
+        optimization_step: int,
+    ) -> None:
+
+        output = model.forward_single(problem, logger)
+        ego_flows = output.to_ego_lidar_flow_list()
+        raw_ego_flows = [
+            (
+                ego_flow.full_flow,
+                ego_flow.mask,
+            )
+            for ego_flow in ego_flows
+        ]
+        save_path = (
+            Path(logger.log_dir)
+            / f"dataset_idx_{problem.dataset_idx:010d}"
+            / f"opt_step_{optimization_step:08d}.pkl"
+        )
+        save_pickle(save_path, raw_ego_flows, verbose=True)
 
     def optimize(
         self,
@@ -40,7 +62,7 @@ def optimize(
         min_delta: Optional[float] = None,
         title: Optional[str] = "Optimizing Neur Rep",
         leave: bool = False,
-        intermediary_results_folder: Optional[Path] = None,
+        save_flow_every: Optional[int] = None,
     ) -> BucketedSceneFlowOutputSequence:
         model = model.train()
         if self.compile:
@@ -70,6 +92,9 @@ def optimize(
                 {f"log/{problem.sequence_log_id}/{problem.dataset_idx:06d}": cost.item()}, step=step
             )
 
+            if save_flow_every is not None and step % save_flow_every == 0:
+                self._save_intermediary_results(model, problem, logger, step)
+
             if cost.item() < lowest_cost:
                 lowest_cost = cost.item()
                 # Run in eval mode to avoid unnecessary computation

util_scripts/logging/tb_logging.py

@@ -17,7 +17,7 @@ def _extract_experiment_name(cfg: Config) -> str:
 
     # If configs subdir is a parent of the config file, we can just use the path below the configs subdir
     if configs_subdir in cfg_filename.parents:
-        return cfg_filename.relative_to(configs_subdir).with_suffix("").as_posix()
+        return cfg_filename.relative_to(configs_subdir).with_suffix("")
 
     # If "configs" is in the config filename, we can use the path below the "configs" subdir
     if "configs" in cfg_filename.parts:
@@ -29,10 +29,11 @@ def _extract_experiment_name(cfg: Config) -> str:
 
     # If "launch_files" is in the config filename, we use the name of the folders one and two levels below
     if "launch_files" in cfg_filename.parts:
+        launch_files_idx = cfg_filename.parts.index("launch_files")
         return (
-            cfg_filename.parts[cfg_filename.parts.index("launch_files") + 1]
+            cfg_filename.parts[launch_files_idx + 1]
             + "/"
-            + cfg_filename.parts[cfg_filename.parts.index("launch_files") + 2]
+            + cfg_filename.parts[launch_files_idx + 2]
         )
 
     return cfg_filename.absolute().with_suffix("").as_posix()
@@ -41,13 +42,11 @@ def _extract_experiment_name(cfg: Config) -> str:
 def setup_tb_logger(cfg: Config, script_name: str) -> TensorBoardLogger:
 
     # Save Dir
-    logger_save_dir = (Path() / "tb_logs" / script_name).absolute()
-
-    # Experiment name
+    base_logging_dir = (Path() / "tb_logs" / script_name).absolute()
     experiment_name = _extract_experiment_name(cfg)
 
-    # Version
-    checkpoint_dir_name = datetime.datetime.now().strftime("%Y_%m_%d-%I_%M_%S_%p")
-    tbl = TensorBoardLogger(logger_save_dir, name=experiment_name, version=checkpoint_dir_name)
+    # Name
+    version = datetime.datetime.now().strftime("%Y_%m_%d-%I_%M_%S_%p")
+    tbl = TensorBoardLogger(base_logging_dir, name=experiment_name, version=version)
     print("Tensorboard logs will be saved to:", tbl.log_dir, flush=True)
     return tbl

visualization/vis_lib/__init__.py

@@ -1,3 +1,4 @@
-from .sequence_visualizer import SequenceVisualizer
+from .base_callback_visualizer import BaseCallbackVisualizer
+from .sequence_visualizer import SequenceVisualizer, ColorEnum
 
-__all__ = ["SequenceVisualizer"]
+__all__ = ["SequenceVisualizer", "ColorEnum", "BaseCallbackVisualizer"]

visualization/vis_lib/base_callback_visualizer.py

@@ -0,0 +1,58 @@
+from bucketed_scene_flow_eval.datastructures import O3DVisualizer
+import open3d as o3d
+from pathlib import Path
+import datetime
+
+
+class BaseCallbackVisualizer(O3DVisualizer):
+
+    def __init__(self, screenshot_path: Path = Path() / "screenshots", point_size: float = 0.1):
+        super().__init__(point_size=point_size)
+        self.screenshot_path = screenshot_path
+
+    def _get_screenshot_path(self) -> Path:
+        return self.screenshot_path / f"{datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.png"
+
+    def save_screenshot(self, vis: o3d.visualization.VisualizerWithKeyCallback):
+        save_name = self._get_screenshot_path()
+        save_name.parent.mkdir(exist_ok=True, parents=True)
+        vis.capture_screen_image(str(save_name))
+
+    def _frame_list_to_color_list(
+        self,
+        num_frames: int,
+        color1: tuple[float, float, float] = (0, 1, 0),
+        color2: tuple[float, float, float] = (0, 0, 1),
+    ) -> list[tuple[float, float, float]]:
+        """
+        Interpolate between the two colors based on the number of frames in the list
+        """
+
+        def interpolate_color(color1, color2, fraction):
+            return tuple(c1 * (1 - fraction) + c2 * fraction for c1, c2 in zip(color1, color2))
+
+        return [
+            interpolate_color(color1, color2, idx / (num_frames - 1)) for idx in range(num_frames)
+        ]
+
+    def _register_callbacks(self, vis: o3d.visualization.VisualizerWithKeyCallback):
+        vis.register_key_callback(ord("S"), self.save_screenshot)
+
+    def draw_everything(self, vis, reset_view=False):
+        self.render(vis, reset_view=reset_view)
+
+    def _print_instructions(self):
+        print("#############################################################")
+        print("Flow moves from the gray point cloud to the white point cloud\n")
+        print(f"Press S to save screenshot (saved to {self.screenshot_path.absolute()})")
+        print("#############################################################")
+
+    def run(self):
+        self._print_instructions()
+        vis = o3d.visualization.VisualizerWithKeyCallback()
+        vis.create_window()
+        vis.get_render_option().background_color = (1, 1, 1)
+        vis.get_view_control().set_up([0, 0, 1])
+        self._register_callbacks(vis)
+        self.draw_everything(vis, reset_view=True)
+        vis.run()

visualization/vis_lib/sequence_visualizer.py

@@ -4,15 +4,14 @@
     AbstractSequenceLoader,
     AbstractAVLidarSequence,
 )
-from .lazy_frame_matrix import CausalLazyFrameMatrix, NonCausalLazyFrameMatrix
+from .lazy_frame_matrix import NonCausalLazyFrameMatrix
 
 from dataclasses import dataclass
 from .lazy_frame_matrix import AbstractFrameMatrix
 from pathlib import Path
 import enum
-import numpy as np
 from typing import Optional
-from bucketed_scene_flow_eval.interfaces import LoaderType
+from visualization.vis_lib import BaseCallbackVisualizer
 
 
 class ColorEnum(enum.Enum):
@@ -54,16 +53,12 @@ class VisState:
     frame_idx: int
     flow_color: ColorEnum = ColorEnum.RED
     frame_step_size: int = 1
-    screenshot_path: Path = Path() / "screenshots"
-
-    def __post_init__(self):
-        self.screenshot_path.mkdir(exist_ok=True)
 
     def __str__(self):
         return f"sequence_idx: {self.sequence_idx}, frame__idx: {self.frame_idx}, "
 
 
-class SequenceVisualizer(O3DVisualizer):
+class SequenceVisualizer(BaseCallbackVisualizer):
 
     def __init__(
         self,
@@ -99,31 +94,15 @@ def _register_callbacks(self, vis: o3d.visualization.VisualizerWithKeyCallback):
         # down arrow decrease sequence_idx
         vis.register_key_callback(264, self.decrease_sequence_idx)
 
-    def _frame_list_to_color_list(self, num_frames: int) -> list[tuple[float, float, float]]:
-        """
-        Interpolate between red and blue based on the number of frames in the list
-        """
-
-        def interpolate_color(color1, color2, fraction):
-            return tuple(c1 * (1 - fraction) + c2 * fraction for c1, c2 in zip(color1, color2))
-
-        color1 = (0, 1, 0)
-        color2 = (0, 0, 1)
-        return [
-            interpolate_color(color1, color2, idx / (num_frames - 1)) for idx in range(num_frames)
-        ]
-
     def get_current_method_name(self) -> str:
         return self.name_lst[self.vis_state.sequence_idx]
 
-    def save_screenshot(self, vis: o3d.visualization.VisualizerWithKeyCallback):
-        save_name = (
-            self.vis_state.screenshot_path
+    def _get_screenshot_path(self) -> Path:
+        return (
+            self.screenshot_path
             / self.sequence_id
             / f"{self.vis_state.frame_idx:06d}_{self.get_current_method_name()}.png"
         )
-        save_name.parent.mkdir(exist_ok=True, parents=True)
-        vis.capture_screen_image(str(save_name))
 
     def increase_sequence_idx(self, vis):
         self.vis_state.sequence_idx += 1
@@ -185,13 +164,3 @@ def _print_instructions(self):
         print("Press left or right arrow to change starter_idx")
         print(f"Press S to save screenshot (saved to {self.vis_state.screenshot_path.absolute()})")
         print("#############################################################")
-
-    def run(self):
-        self._print_instructions()
-        vis = o3d.visualization.VisualizerWithKeyCallback()
-        vis.create_window()
-        vis.get_render_option().background_color = (1, 1, 1)
-        vis.get_view_control().set_up([0, 0, 1])
-        self._register_callbacks(vis)
-        self.draw_everything(vis, reset_view=True)
-        vis.run()