diff --git a/.gitignore b/.gitignore index 857e741..4a5c292 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,8 @@ log/ results/ _skbuild/ .gitlab-ci-local/ +.polyscope.ini +imgui.ini # Created by https://www.toptal.com/developers/gitignore/api/python,c++ # Edit at https://www.toptal.com/developers/gitignore?templates=python,c++ diff --git a/src/mapmos/mapping.py b/src/mapmos/mapping.py index 6f81df4..f72b1bc 100644 --- a/src/mapmos/mapping.py +++ b/src/mapmos/mapping.py @@ -50,6 +50,10 @@ def point_cloud_with_timestamps(self): map_points, map_timestamps = self._internal_map._point_cloud_with_timestamps() return np.asarray(map_points), np.asarray(map_timestamps) + def voxels_with_belief(self): + voxels, belief = self._internal_map._voxels_with_belief() + return np.asarray(voxels), np.asarray(belief) + def remove_voxels_far_from_location(self, location: np.ndarray): self._internal_map._remove_far_away_points(location) diff --git a/src/mapmos/paper_pipeline.py b/src/mapmos/paper_pipeline.py index 29f1000..8536ef3 100644 --- a/src/mapmos/paper_pipeline.py +++ b/src/mapmos/paper_pipeline.py @@ -158,19 +158,14 @@ def _run_pipeline(self): torch.tensor(gt_labels, dtype=torch.int32), ) - belief_labels_scan = belief_labels_with_map - if self.visualize: - belief_map = self.belief.get_belief(map_points) - belief_labels_map = self.model.to_label(belief_map) - self.visualizer.update( - scan_points, - map_points, - pred_labels_scan, - pred_labels_map, - belief_labels_scan, - belief_labels_map, - self.odometry.last_pose, - ) + self.visualizer.update( + scan_points, + map_points, + pred_labels_scan, + pred_labels_map, + self.belief, + self.odometry.last_pose, + ) # Probabilistic volumetric fusion with scan and moving map predictions and delay self.buffer.append([scan_index, scan_points, gt_labels]) diff --git a/src/mapmos/pipeline.py b/src/mapmos/pipeline.py index f704d32..7a1d4e1 100644 --- a/src/mapmos/pipeline.py +++ b/src/mapmos/pipeline.py @@ -98,7 +98,7 @@ def __init__( # Visualizer self.visualize = visualize self.visualizer = MapMOSVisualizer() if visualize else StubVisualizer() - + self.visualizer.set_voxel_size(self.config.mos.voxel_size_belief) self.ply_writer = PlyWriter() if save_ply else StubWriter() self.kitti_writer = KITTIWriter() if save_kitti else StubWriter() @@ -174,22 +174,17 @@ def _run_pipeline(self): start_time = time.perf_counter_ns() self.belief.update_belief(points_stacked, logits_stacked) - belief_scan = self.belief.get_belief(scan_points) + self.belief.get_belief(scan_points) self.times_belief.append(time.perf_counter_ns() - start_time) - belief_labels_scan = self.model.to_label(belief_scan) - - if self.visualize: - belief_map = self.belief.get_belief(map_points) - belief_labels_map = self.model.to_label(belief_map) - self.visualizer.update( - scan_points, - map_points, - pred_labels_scan, - pred_labels_map, - belief_labels_scan, - belief_labels_map, - self.odometry.last_pose, - ) + + self.visualizer.update( + scan_points, + map_points, + pred_labels_scan, + pred_labels_map, + self.belief, + self.odometry.last_pose, + ) # Evaluate and save with delay self.buffer.append([scan_index, scan_points, gt_labels]) diff --git a/src/mapmos/pybind/VoxelHashMap.cpp b/src/mapmos/pybind/VoxelHashMap.cpp index 3c68311..863bea1 100644 --- a/src/mapmos/pybind/VoxelHashMap.cpp +++ b/src/mapmos/pybind/VoxelHashMap.cpp @@ -77,6 +77,17 @@ std::tuple, std::vector> VoxelHashMap::Pointcl return std::make_tuple(points, timestamps); } +std::tuple, std::vector> VoxelHashMap::VoxelsWithBelief() + const { + std::vector voxels; + std::vector belief; + for (auto map_element : map_) { + voxels.push_back(map_element.first); + belief.push_back(map_element.second.belief.value_); + } + return make_tuple(voxels, belief); +} + void VoxelHashMap::Update(const std::vector &points, const Eigen::Vector3d &origin, const int timestamp) { diff --git a/src/mapmos/pybind/VoxelHashMap.hpp b/src/mapmos/pybind/VoxelHashMap.hpp index 5f5b4d2..a3dff73 100644 --- a/src/mapmos/pybind/VoxelHashMap.hpp +++ b/src/mapmos/pybind/VoxelHashMap.hpp @@ -97,6 +97,7 @@ struct VoxelHashMap { void RemovePointsFarFromLocation(const Eigen::Vector3d &origin); std::vector Pointcloud() const; std::tuple, std::vector> PointcloudWithTimestamps() const; + std::tuple, std::vector> VoxelsWithBelief() const; std::vector GetPoints(const std::vector &query_voxels) const; double voxel_size_; diff --git a/src/mapmos/pybind/mapmos_pybind.cpp b/src/mapmos/pybind/mapmos_pybind.cpp index b295d98..59daf19 100644 --- a/src/mapmos/pybind/mapmos_pybind.cpp +++ b/src/mapmos/pybind/mapmos_pybind.cpp @@ -68,6 +68,7 @@ PYBIND11_MODULE(mapmos_pybind, m) { .def("_remove_far_away_points", &VoxelHashMap::RemovePointsFarFromLocation, "origin"_a) .def("_point_cloud", &VoxelHashMap::Pointcloud) .def("_point_cloud_with_timestamps", &VoxelHashMap::PointcloudWithTimestamps) + .def("_voxels_with_belief", &VoxelHashMap::VoxelsWithBelief) .def("_update_belief", &VoxelHashMap::UpdateBelief, "points"_a, "updates"_a) .def("_get_belief", &VoxelHashMap::GetBelief, "points"_a); diff --git a/src/mapmos/utils/visualizer.py b/src/mapmos/utils/visualizer.py index e219920..ed25620 100644 --- a/src/mapmos/utils/visualizer.py +++ b/src/mapmos/utils/visualizer.py @@ -1,7 +1,6 @@ # MIT License # -# Copyright (c) 2022 Ignacio Vizzo, Tiziano Guadagnino, Benedikt Mersch, Cyrill -# Stachniss. +# Copyright (c) 2024 Benedikt Mersch, Luca Lobefaro, Ignazio Vizzo, Tiziano Guadagnino # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal @@ -20,21 +19,64 @@ # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. -import copy +import datetime import importlib import os from abc import ABC -from functools import partial -from typing import Callable, List import numpy as np -YELLOW = np.array([1, 0.706, 0]) -RED = np.array([128, 0, 0]) / 255.0 -BLACK = np.array([0, 0, 0]) / 255.0 -GRAY = np.array([0.4, 0.4, 0.4]) -BLUE = np.array([0.4, 0.5, 0.9]) -SPHERE_SIZE = 0.20 +# Button names +START_BUTTON = " START\n[SPACE]" +PAUSE_BUTTON = " PAUSE\n[SPACE]" +NEXT_FRAME_BUTTON = "NEXT FRAME\n\t\t [N]" +SCREENSHOT_BUTTON = "SCREENSHOT\n\t\t [S]" +LOCAL_VIEW_BUTTON = "LOCAL VIEW\n\t\t [G]" +GLOBAL_VIEW_BUTTON = "GLOBAL VIEW\n\t\t [G]" +CENTER_VIEWPOINT_BUTTON = "CENTER VIEWPOINT\n\t\t\t\t[C]" +SHOW_BELIEF_BUTTON = "SHOW BELIEF\n\t\t\t[V]" +HIDE_BELIEF_BUTTON = "HIDE BELIEF\n\t\t [V]" +QUIT_BUTTON = "QUIT\n [Q]" + +# Colors +BACKGROUND_COLOR = [0.0, 0.0, 0.0] +FRAME_COLOR = [0.8470, 0.1058, 0.3764] +MAP_COLOR = [0.0, 0.3019, 0.2509] +BELIEF_COLOR = [0.9, 0.9, 0.9] + +# Size constants +FRAME_PTS_SIZE = 0.06 +MAP_PTS_SIZE = 0.08 + +# Voxels Prototype +VOXEL_VERTICES = np.array( + [ + [0, 0, 0], + [1, 0, 0], + [1, 1, 0], + [0, 1, 0], + [0, 0, 1], + [1, 0, 1], + [1, 1, 1], + [0, 1, 1], + ] +).astype(np.float64) +VOXEL_EDGES = np.array( + [ + [0, 1], + [1, 2], + [2, 3], + [0, 3], + [0, 4], + [5, 4], + [1, 5], + [5, 6], + [2, 6], + [6, 7], + [3, 7], + [4, 7], + ] +).astype(np.int64) class StubVisualizer(ABC): @@ -47,57 +89,48 @@ def update( map_points, pred_labels_scan, pred_labels_map, - belief_labels_scan, - belief_labels_map, + belief_map, pose, ): pass + def set_voxel_size(self, voxel_size): + pass + class MapMOSVisualizer(StubVisualizer): - # Public Interaface ---------------------------------------------------------------------------- + # Public Interface ---------------------------------------------------------------------------- def __init__(self): try: - self.o3d = importlib.import_module("open3d") + self._ps = importlib.import_module("polyscope") + self._gui = self._ps.imgui except ModuleNotFoundError as err: - print(f'open3d is not installed on your system, run "pip install open3d"') + print(f'polyscope is not installed on your system, run "pip install polyscope"') exit(1) # Initialize GUI controls - self.block_vis = True - self.play_crun = False - self.reset_bounding_box = True - - # Create data - self.scan = self.o3d.geometry.PointCloud() - self.map = self.o3d.geometry.PointCloud() - self.frames = [] - - # Initialize visualizer - self.vis = self.o3d.visualization.VisualizerWithKeyCallback() - self._register_key_callbacks() + self._background_color = BACKGROUND_COLOR + self._frame_size = FRAME_PTS_SIZE + self._map_size = MAP_PTS_SIZE + self._block_execution = True + self._play_mode = False + self._toggle_frame = True + self._toggle_map = True + self._toggle_belief = False + self._global_view = False + self._last_pose = np.eye(4) + self._voxel_size = 1.0 + + # Initialize Visualizer self._initialize_visualizer() - # Visualization options - self.render_map = True - self.render_scan = True - self.visualize_belief = False - self.global_view = False - self.render_trajectory = True - # Cache the state of the visualizer - self.state = ( - self.render_map, - self.render_scan, - ) - def update( self, scan_points, map_points, pred_labels_scan, pred_labels_map, - belief_labels_scan, - belief_labels_map, + belief_map, pose, ): self._update_geometries( @@ -105,107 +138,27 @@ def update( map_points, pred_labels_scan, pred_labels_map, - belief_labels_scan, - belief_labels_map, + belief_map, pose, ) - while self.block_vis: - self.vis.poll_events() - self.vis.update_renderer() - if self.play_crun: + while self._block_execution: + self._ps.frame_tick() + if self._play_mode: break - self.block_vis = not self.block_vis + self._block_execution = not self._block_execution - # Private Interaface --------------------------------------------------------------------------- - def _initialize_visualizer(self): - w_name = self.__class__.__name__ - self.vis.create_window(window_name=w_name, width=1920, height=1080) - self.vis.add_geometry(self.scan) - self.vis.add_geometry(self.map) - self._set_black_background(self.vis) - self.vis.get_render_option().point_size = 1 - print( - f"{w_name} initialized. Press:\n" - "\t[SPACE] to pause/start\n" - "\t [ESC] to exit\n" - "\t [N] to step\n" - "\t [F] to toggle on/off the input cloud to the pipeline\n" - "\t [K] to toggle between prediction and belief\n" - "\t [M] to toggle on/off the local map\n" - "\t [V] to toggle ego/global viewpoint\n" - "\t [T] to toggle the trajectory view(only available in global view)\n" - "\t [C] to center the viewpoint\n" - "\t [W] to toggle a white background\n" - "\t [B] to toggle a black background\n" - ) + def set_voxel_size(self, voxel_size): + self._voxel_size = voxel_size - def _register_key_callback(self, keys: List, callback: Callable): - for key in keys: - self.vis.register_key_callback(ord(str(key)), partial(callback)) - - def _register_key_callbacks(self): - self._register_key_callback(["Ā", "Q", "\x1b"], self._quit) - self._register_key_callback([" "], self._start_stop) - self._register_key_callback(["N"], self._next_frame) - self._register_key_callback(["V"], self._toggle_view) - self._register_key_callback(["C"], self._center_viewpoint) - self._register_key_callback(["F"], self._toggle_scan) - self._register_key_callback(["K"], self._toggle_visualize_belief) - self._register_key_callback(["M"], self._toggle_map) - self._register_key_callback(["T"], self._toggle_trajectory) - self._register_key_callback(["B"], self._set_black_background) - self._register_key_callback(["W"], self._set_white_background) - - def _set_black_background(self, vis): - vis.get_render_option().background_color = [0.0, 0.0, 0.0] - - def _set_white_background(self, vis): - vis.get_render_option().background_color = [1.0, 1.0, 1.0] - - def _quit(self, vis): - print("Destroying Visualizer") - vis.destroy_window() - os._exit(0) - - def _next_frame(self, vis): - self.block_vis = not self.block_vis - - def _start_stop(self, vis): - self.play_crun = not self.play_crun - - def _toggle_scan(self, vis): - self.render_scan = not self.render_scan - return False - - def _toggle_visualize_belief(self, vis): - self.visualize_belief = not self.visualize_belief - return False - - def _toggle_map(self, vis): - self.render_map = not self.render_map - return False - - def _toggle_view(self, vis): - self.global_view = not self.global_view - self._trajectory_handle() - - def _center_viewpoint(self, vis): - vis.reset_view_point(True) - - def _toggle_trajectory(self, vis): - if not self.global_view: - return False - self.render_trajectory = not self.render_trajectory - self._trajectory_handle() - return False - - def _trajectory_handle(self): - if self.render_trajectory and self.global_view: - for frame in self.frames: - self.vis.add_geometry(frame, reset_bounding_box=False) - else: - for frame in self.frames: - self.vis.remove_geometry(frame, reset_bounding_box=False) + # Private Interface --------------------------------------------------------------------------- + def _initialize_visualizer(self): + self._ps.set_program_name("MapMOS Visualizer") + self._ps.init() + self._ps.set_ground_plane_mode("none") + self._ps.set_background_color(BACKGROUND_COLOR) + self._ps.set_verbosity(0) + self._ps.set_user_callback(self._main_gui_callback) + self._ps.set_build_default_gui_panels(False) def _update_geometries( self, @@ -213,53 +166,194 @@ def _update_geometries( map_points, pred_labels_scan, pred_labels_map, - belief_labels_scan, - belief_labels_map, + belief_map, pose, ): - # Scan - if self.render_scan: - self.scan.points = self.o3d.utility.Vector3dVector(scan_points) - if not self.global_view: - self.scan.transform(np.linalg.inv(pose)) - self.scan.paint_uniform_color(BLUE) - scan_colors = np.array(self.scan.colors) - if self.visualize_belief: - scan_colors[belief_labels_scan == 1] = YELLOW - else: - scan_colors[pred_labels_scan == 1] = RED - self.scan.colors = self.o3d.utility.Vector3dVector(scan_colors) - else: - self.scan.points = self.o3d.utility.Vector3dVector() - - # Map - if self.render_map: - self.map.points = self.o3d.utility.Vector3dVector(map_points) - self.map.paint_uniform_color(GRAY) - map_colors = np.array(self.map.colors) - if self.visualize_belief: - map_colors[belief_labels_map == 1] = YELLOW - else: - map_colors[pred_labels_map == 1] = RED + scan_cloud = self._ps.register_point_cloud( + "scan", + scan_points, + point_render_mode="quad", + ) + + scan_colors = np.zeros((len(pred_labels_scan), 3)) + scan_colors[pred_labels_scan == 1, :] = [1, 0, 0] + scan_cloud.set_radius(self._frame_size, relative=False) + scan_cloud.add_color_quantity("colors", scan_colors, enabled=True) + scan_cloud.set_enabled(self._toggle_frame) - self.map.colors = self.o3d.utility.Vector3dVector(map_colors) - if not self.global_view: - self.map.transform(np.linalg.inv(pose)) + map_cloud = self._ps.register_point_cloud( + "map", + map_points, + point_render_mode="quad", + ) + map_colors = np.zeros((len(pred_labels_map), 3)) + map_colors[pred_labels_map == 1, :] = [1, 0, 0] + map_cloud.set_radius(self._map_size, relative=False) + map_cloud.add_color_quantity("colors", map_colors, enabled=True) + map_cloud.set_enabled(self._toggle_map) + + # VOXEL GRID (only if toggled) + self._last_belief_map = belief_map + if self._toggle_belief: + self._register_belief() + + if self._global_view: + scan_cloud.set_transform(np.eye(4)) + map_cloud.set_transform(np.eye(4)) + else: + inv_pose = np.linalg.inv(pose) + scan_cloud.set_transform(inv_pose) + map_cloud.set_transform(inv_pose) + + self._last_pose = pose + + def _register_belief(self): + if self._last_belief_map is None: + return + voxels, belief = self._last_belief_map.voxels_with_belief() + + belief_nodes = np.zeros((voxels.shape[0] * 8, 3), dtype=np.float64) + belief_colors = np.zeros((voxels.shape[0] * 8, 3), dtype=np.float64) + belief_edges = np.zeros((voxels.shape[0] * 12, 2), dtype=np.int64) + + for idx, (voxel, belief) in enumerate(zip(voxels, belief)): + verts = np.copy(VOXEL_VERTICES) + voxel + verts = verts * self._voxel_size + edges = np.copy(VOXEL_EDGES) + (idx * 8) + belief_nodes[idx * 8 : idx * 8 + 8] = verts + belief_colors[idx * 8 : idx * 8 + 8] = [1, 0, 0] if belief > 0 else BELIEF_COLOR + belief_edges[idx * 12 : idx * 12 + 12] = edges + + belief = self._ps.register_curve_network("belief", belief_nodes, belief_edges) + belief.set_radius(0.01, relative=False) + belief.add_color_quantity("belief", belief_colors, enabled=True) + if self._global_view: + belief.set_transform(np.eye(4)) else: - self.map.points = self.o3d.utility.Vector3dVector() - - # Update always a list with all the trajectories - new_frame = self.o3d.geometry.TriangleMesh.create_sphere(SPHERE_SIZE) - new_frame.paint_uniform_color(BLUE) - new_frame.compute_vertex_normals() - new_frame.transform(pose) - self.frames.append(new_frame) - # Render trajectory, only if it make sense (global view) - if self.render_trajectory and self.global_view: - self.vis.add_geometry(self.frames[-1], reset_bounding_box=False) - - self.vis.update_geometry(self.scan) - self.vis.update_geometry(self.map) - if self.reset_bounding_box: - self.vis.reset_view_point(True) - self.reset_bounding_box = False + belief.set_transform(np.linalg.inv(self._last_pose)) + + def _unregister_belief(self): + if self._ps.has_curve_network("belief"): + self._ps.remove_curve_network("belief") + + # GUI Callbacks --------------------------------------------------------------------------- + def _start_pause_callback(self): + button_name = PAUSE_BUTTON if self._play_mode else START_BUTTON + if self._gui.Button(button_name) or self._gui.IsKeyPressed(self._gui.ImGuiKey_Space): + self._play_mode = not self._play_mode + if self._play_mode: + self._toggle_belief = False + self._unregister_belief() + self._ps.set_SSAA_factor(1) + else: + self._ps.set_SSAA_factor(4) + + def _next_frame_callback(self): + if self._gui.Button(NEXT_FRAME_BUTTON) or self._gui.IsKeyPressed(self._gui.ImGuiKey_N): + self._block_execution = not self._block_execution + + def _screenshot_callback(self): + if self._gui.Button(SCREENSHOT_BUTTON) or self._gui.IsKeyPressed(self._gui.ImGuiKey_S): + image_filename = "mapmos_" + ( + datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + ".jpg" + ) + self._ps.screenshot(image_filename) + + def _center_viewpoint_callback(self): + if self._gui.Button(CENTER_VIEWPOINT_BUTTON) or self._gui.IsKeyPressed( + self._gui.ImGuiKey_C + ): + self._ps.reset_camera_to_home_view() + + def _toggle_buttons_andslides_callback(self): + # FRAME + changed, self._frame_size = self._gui.SliderFloat( + "##frame_size", self._frame_size, v_min=0.01, v_max=0.6 + ) + if changed: + self._ps.get_point_cloud("scan").set_radius(self._frame_size, relative=False) + self._gui.SameLine() + changed, self._toggle_frame = self._gui.Checkbox("Frame Cloud", self._toggle_frame) + if changed: + self._ps.get_point_cloud("scan").set_enabled(self._toggle_frame) + + # MAP + changed, self._map_size = self._gui.SliderFloat( + "##map_size", self._map_size, v_min=0.01, v_max=0.6 + ) + if changed: + self._ps.get_point_cloud("map").set_radius(self._map_size, relative=False) + self._gui.SameLine() + changed, self._toggle_map = self._gui.Checkbox("Local Map", self._toggle_map) + if changed: + self._ps.get_point_cloud("map").set_enabled(self._toggle_map) + + def _background_color_callback(self): + changed, self._background_color = self._gui.ColorEdit3( + "Background Color", + self._background_color, + ) + if changed: + self._ps.set_background_color(self._background_color) + + def _inspection_callback(self): + if self._gui.TreeNodeEx("Inspection", self._gui.ImGuiTreeNodeFlags_DefaultOpen): + # VOXEL GRID Button + belief_button_name = HIDE_BELIEF_BUTTON if self._toggle_belief else SHOW_BELIEF_BUTTON + if self._gui.Button(belief_button_name) or self._gui.IsKeyPressed(self._gui.ImGuiKey_V): + self._toggle_belief = not self._toggle_belief + if self._toggle_belief: + self._register_belief() + else: + self._unregister_belief() + self._gui.TreePop() + + def _global_view_callback(self): + button_name = LOCAL_VIEW_BUTTON if self._global_view else GLOBAL_VIEW_BUTTON + if self._gui.Button(button_name) or self._gui.IsKeyPressed(self._gui.ImGuiKey_G): + self._global_view = not self._global_view + if self._global_view: + self._ps.get_point_cloud("scan").set_transform(np.eye(4)) + self._ps.get_point_cloud("map").set_transform(np.eye(4)) + if self._toggle_belief: + self._ps.get_curve_network("belief").set_transform(np.eye(4)) + else: + inv_pose=np.linalg.inv(self._last_pose) + self._ps.get_point_cloud("scan").set_transform(inv_pose) + self._ps.get_point_cloud("map").set_transform(inv_pose) + if self._toggle_belief: + self._ps.get_curve_network("belief").set_transform(inv_pose) + self._ps.reset_camera_to_home_view() + + def _quit_callback(self): + self._gui.SetCursorPosX( + self._gui.GetCursorPosX() + self._gui.GetContentRegionAvail()[0] - 50 + ) + if ( + self._gui.Button(QUIT_BUTTON) + or self._gui.IsKeyPressed(self._gui.ImGuiKey_Escape) + or self._gui.IsKeyPressed(self._gui.ImGuiKey_Q) + ): + print("Destroying Visualizer") + self._ps.unshow() + os._exit(0) + + def _main_gui_callback(self): + # GUI callbacks + self._start_pause_callback() + if not self._play_mode: + self._gui.SameLine() + self._next_frame_callback() + self._gui.SameLine() + self._screenshot_callback() + self._gui.Separator() + self._toggle_buttons_andslides_callback() + self._background_color_callback() + if not self._play_mode: + self._gui.Separator() + self._inspection_callback() + self._global_view_callback() + self._gui.SameLine() + self._center_viewpoint_callback() + self._gui.Separator() + self._quit_callback()