diff --git a/configs/fast_nsf/argo/val_debug.py b/configs/fast_nsf/argo/val_debug.py index 61c1910..f1eb3b9 100644 --- a/configs/fast_nsf/argo/val_debug.py +++ b/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)) +) diff --git a/models/optimization/test_time_optimizer_loop.py b/models/optimization/test_time_optimizer_loop.py index 4f487ea..67f3a07 100644 --- a/models/optimization/test_time_optimizer_loop.py +++ b/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 diff --git a/util_scripts/logging/tb_logging.py b/util_scripts/logging/tb_logging.py index 9bf2b54..2af5faf 100644 --- a/util_scripts/logging/tb_logging.py +++ b/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 diff --git a/visualization/vis_lib/__init__.py b/visualization/vis_lib/__init__.py index cb39ab5..2c140b8 100644 --- a/visualization/vis_lib/__init__.py +++ b/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"] diff --git a/visualization/vis_lib/base_callback_visualizer.py b/visualization/vis_lib/base_callback_visualizer.py new file mode 100644 index 0000000..f2e01df --- /dev/null +++ b/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() diff --git a/visualization/vis_lib/sequence_visualizer.py b/visualization/vis_lib/sequence_visualizer.py index 3b3e3bc..d8f3ae3 100644 --- a/visualization/vis_lib/sequence_visualizer.py +++ b/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() diff --git a/visualization/visualize_tto_itermediary_results.py b/visualization/visualize_tto_itermediary_results.py new file mode 100644 index 0000000..defc9f6 --- /dev/null +++ b/visualization/visualize_tto_itermediary_results.py @@ -0,0 +1,158 @@ +import argparse +from pathlib import Path +from bucketed_scene_flow_eval.utils import load_pickle +from bucketed_scene_flow_eval.datastructures import ( + O3DVisualizer, + EgoLidarFlow, + PointCloud, + TimeSyncedSceneFlowFrame, +) +from bucketed_scene_flow_eval.datasets import ( + construct_dataset, + AbstractDataset, +) +from visualization.vis_lib import ColorEnum, BaseCallbackVisualizer +from dataclasses import dataclass +import numpy as np +import open3d as o3d + + +@dataclass(kw_only=True) +class VisState: + flow_color: ColorEnum = ColorEnum.RED + intermedary_result_idx: int = 0 + + +class ResultsVisualizer(BaseCallbackVisualizer): + + def __init__(self, dataset: AbstractDataset, intermediary_results_folder: Path): + + self.vis_state = VisState() + self.dataset = dataset + self.dataset_idx, self.intermediary_files = self._load_intermediary_flows( + intermediary_results_folder + ) + + super().__init__( + screenshot_path=Path() + / "intermediary_screenshots" + / f"dataset_idx_{self.dataset_idx:010d}" + ) + + def _load_intermediary_flows(self, intermediary_results_folder: Path) -> tuple[int, list[Path]]: + dataset_idx = int(intermediary_results_folder.name.split("_")[-1]) + assert intermediary_results_folder.exists(), f"{intermediary_results_folder} does not exist" + intermediary_files = sorted(intermediary_results_folder.glob("*.pkl")) + assert ( + len(intermediary_files) > 0 + ), f"No intermediary files found in {intermediary_results_folder}" + + print(f"Found {len(intermediary_files)} intermediary files for dataset idx {dataset_idx}") + return dataset_idx, intermediary_files + + def _load_ego_lidar(self) -> list[EgoLidarFlow]: + raw_intermediary_flows: list[tuple[np.ndarray, np.ndarray]] = load_pickle( + self.intermediary_files[self.vis_state.intermedary_result_idx], verbose=False + ) + return [EgoLidarFlow(*raw_flow) for raw_flow in raw_intermediary_flows] + + def _get_result_data(self) -> list[TimeSyncedSceneFlowFrame]: + dataset_frame_list = self.dataset[self.dataset_idx] + intermediary_ego_flows = self._load_ego_lidar() + + assert ( + len(dataset_frame_list) == len(intermediary_ego_flows) + 1 + ), f"Expected one more frame in dataset than intermediary flows; instead found {len(dataset_frame_list)} and {len(intermediary_ego_flows)}" + + for frame_info, ego_flow in zip(dataset_frame_list, intermediary_ego_flows): + # Add the intermediary ego flow to the frame info + frame_info.flow = ego_flow + + return dataset_frame_list + + def _get_screenshot_path(self) -> Path: + return self.screenshot_path / f"{self.vis_state.intermedary_result_idx:08d}.png" + + def _print_instructions(self): + print("#############################################################") + print("Flow moves from the gray point cloud to the white point cloud\n") + print("Press up or down arrow to change intermediary result") + print(f"Press S to save screenshot (saved to {self.screenshot_path.absolute()})") + print("#############################################################") + + 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) + + def increase_sequence_idx(self, vis): + 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.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 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}") + 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): + pc = flow_frame.pc.global_pc + self.add_pointcloud(pc, color=color_list[idx]) + flowed_pc = flow_frame.pc.flow(flow_frame.flow).global_pc + + draw_color = self.vis_state.flow_color.rgb + + # Add flowed point cloud + if (draw_color is not None) and (flowed_pc is not None) and (idx < len(frame_list) - 1): + self.add_lineset(pc, flowed_pc, color=draw_color) + vis.clear_geometries() + self.render(vis, reset_view=reset_view) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--dataset_name", type=str, default="Argoverse2CausalSceneFlow") + parser.add_argument("root_dir", type=Path) + parser.add_argument("subsequence_length", type=int) + parser.add_argument("intermediary_results_folder", type=Path) + args = parser.parse_args() + + # def __init__( + # self, + # root_dir: Union[Path, list[Path]], + # subsequence_length: int = 2, + # with_ground: bool = True, + # with_rgb: bool = False, + # cache_root: Path = Path("/tmp/"), + # use_gt_flow: bool = True, + # flow_data_path: Optional[Union[Path, list[Path]]] = None, + # eval_type: str = "bucketed_epe", + # eval_args=dict(), + # expected_camera_shape: tuple[int, int, int] = (1550, 2048, 3), + # point_cloud_range: Optional[PointCloudRange] = DEFAULT_POINT_CLOUD_RANGE, + # use_cache=True, + # load_flow: bool = True, + # ) -> None: + + dataset = construct_dataset( + name=args.dataset_name, + args=dict(root_dir=args.root_dir, subsequence_length=args.subsequence_length), + ) + + visualizer = ResultsVisualizer(dataset, args.intermediary_results_folder) + + visualizer.run() + + +if __name__ == "__main__": + main()