python-motion-planning 2.0.dev1__py3-none-any.whl → 2.0.1__py3-none-any.whl

python_motion_planning/common/visualizer/{visualizer.py → visualizer_2d.py}

@@ -1,11 +1,12 @@
 """
-@file: visualizer.py
+@file: visualizer_2d.py
 @author: Wu Maojia, Yang Haodong 
-@update: 2025.10.3
+@update: 2025.12.20
 """
 from typing import Union, Dict, List, Tuple, Any
 from collections import namedtuple
 import time
+import os
 
 import numpy as np
 import matplotlib
@@ -14,52 +15,60 @@ import matplotlib.colors as mcolors
 from matplotlib import animation
 import matplotlib.patheffects as path_effects
 
+from python_motion_planning.common.visualizer.base_visualizer import BaseVisualizer
 from python_motion_planning.controller import BaseController
 from python_motion_planning.common.env import TYPES, ToySimulator, Grid, CircularRobot, Node
 from python_motion_planning.common.utils import Geometry
 
-class Visualizer:
+class Visualizer2D(BaseVisualizer):
     """
     Simple visualizer for motion planning using matplotlib.
 
     Args:
         figname: Figure name (window title).
-        figsize: Figure size (width, height).
+        figsize: Figure size (width, height) (matplotlib figure size, unit: inch).
+        cmap_dict: Color map for 2d visualization.
+        zorder: Zorder for 2d matplotlib visualization.
     """
-    def __init__(self, figname: str = "", figsize: tuple = (10, 8)):
+    def __init__(self, 
+                figname: str = "", 
+                figsize: tuple = (10, 8), 
+                cmap_dict: dict = {
+                    TYPES.FREE: "#ffffff",
+                    TYPES.OBSTACLE: "#000000",
+                    TYPES.START: "#ff0000",
+                    TYPES.GOAL: "#1155cc",
+                    TYPES.INFLATION: "#ffccff",
+                    TYPES.EXPAND: "#eeeeee",
+                    TYPES.CUSTOM: "#bbbbbb",
+                },
+                zorder: dict = {
+                    'grid_map': 10,
+                    'voxels': 10,
+                    'esdf': 20,
+                    'expand_tree_edge': 30,
+                    'expand_tree_node': 40,
+                    'path_2d': 50,
+                    'path_3d': 700,
+                    'traj': 60,
+                    'lookahead_pose_node': 70,
+                    'lookahead_pose_orient': 80,
+                    'pred_traj': 90,
+                    'robot_circle': 100,
+                    'robot_orient': 110,
+                    'robot_text': 120,
+                    'env_info_text': 10000
+                }
+            ):
         self.fig = plt.figure(figname, figsize=figsize)
         self.ax = self.fig.add_subplot()
         self.ani = None
 
         # colors
-        self.cmap_dict = {
-            TYPES.FREE: "#ffffff",
-            TYPES.OBSTACLE: "#000000",
-            TYPES.START: "#ff0000",
-            TYPES.GOAL: "#1155cc",
-            TYPES.INFLATION: "#ffccff",
-            TYPES.EXPAND: "#eeeeee",
-            TYPES.CUSTOM: "#bbbbbb",
-        }
+        self.cmap_dict = cmap_dict
         # self.norm = mcolors.BoundaryNorm(list(range(len(self.cmap_dict))), len(self.cmap_dict))
 
-        self.zorder = {
-            'grid_map': 10,
-            'voxels': 10,
-            'esdf': 20,
-            'expand_tree_edge': 30,
-            'expand_tree_node': 40,
-            'path_2d': 50,
-            'path_3d': 700,
-            'traj': 60,
-            'lookahead_pose_node': 70,
-            'lookahead_pose_orient': 80,
-            'pred_traj': 90,
-            'robot_circle': 100,
-            'robot_orient': 110,
-            'robot_text': 120,
-            'env_info_text': 10000
-        }
+        self.zorder = zorder
 
         self.cmap = mcolors.ListedColormap([info for info in self.cmap_dict.values()])
         self.norm = mcolors.BoundaryNorm([i for i in range(self.cmap.N + 1)], self.cmap.N)
@@ -71,15 +80,7 @@ class Visualizer:
     def __del__(self):
         self.close()
 
-    def plot_grid_map(self, grid_map: Grid, equal: bool = False, alpha_3d: dict = {
-                            TYPES.FREE: 0.0,
-                            TYPES.OBSTACLE: 0.5,
-                            TYPES.START: 0.5,
-                            TYPES.GOAL: 0.5,
-                            TYPES.INFLATION: 0.0,
-                            TYPES.EXPAND: 0.1,
-                            TYPES.CUSTOM: 0.5,
-                        },
+    def plot_grid_map(self, grid_map: Grid, equal: bool = False,
                         show_esdf: bool = False, alpha_esdf: float = 0.5) -> None:
         '''
         Plot grid map with static obstacles.
@@ -87,82 +88,44 @@ class Visualizer:
         Args:
             map: Grid map or its type map.
             equal: Whether to set axis equal.
-            alpha_3d: Alpha of occupancy for 3d visualization.
             show_esdf: Whether to show esdf.
             alpha_esdf: Alpha of esdf.
         '''
+        if grid_map.dim != 2:
+            raise ValueError(f"Grid map dimension must be 2.")
+
         self.grid_map = grid_map
         self.dim = grid_map.dim
-        if grid_map.dim == 2:
+        type_data = grid_map.type_map.data
+
+        plt.imshow(
+            np.transpose(type_data), 
+            cmap=self.cmap, 
+            norm=self.norm, 
+            origin='lower', 
+            interpolation='nearest', 
+            extent=[*grid_map.bounds[0], *grid_map.bounds[1]],
+            zorder=self.zorder['grid_map'],
+            )
+
+        if show_esdf:   # draw esdf hotmap
             plt.imshow(
-                np.transpose(grid_map.type_map.array), 
-                cmap=self.cmap, 
-                norm=self.norm, 
-                origin='lower', 
-                interpolation='nearest', 
+                np.transpose(grid_map.esdf),
+                cmap="jet",
+                origin="lower",
+                interpolation="nearest",
                 extent=[*grid_map.bounds[0], *grid_map.bounds[1]],
-                zorder=self.zorder['grid_map'],
-                )
-
-            if show_esdf:   # draw esdf hotmap
-                plt.imshow(
-                    np.transpose(grid_map.esdf),
-                    cmap="jet",
-                    origin="lower",
-                    interpolation="nearest",
-                    extent=[*grid_map.bounds[0], *grid_map.bounds[1]],
-                    alpha=alpha_esdf,
-                    zorder=self.zorder['esdf'],
-                )
-                plt.colorbar(label="ESDF distance")
-                
-            if equal: 
-                plt.axis("equal")
-
-        elif grid_map.dim == 3:
-            self.ax = self.fig.add_subplot(projection='3d')
-
-            data = grid_map.type_map.array
-            nx, ny, nz = data.shape
-
-            filled = np.zeros_like(data, dtype=bool)
-            colors = np.zeros(data.shape + (4,), dtype=float)  # RGBA
-
-            for key, color in self.cmap_dict.items():
-                mask = (data == key)
-                if alpha_3d[key] < 1e-6:
-                    continue
-                filled |= mask
-                rgba = matplotlib.colors.to_rgba(color, alpha=alpha_3d[key])
-                colors[mask] = rgba
-
-            self.ax.voxels(filled, facecolors=colors, zorder=self.zorder['voxels'])
-
-            if show_esdf:
-                # TODO
-                raise NotImplementedError
-
-            self.ax.set_xlabel("X")
-            self.ax.set_ylabel("Y")
-            self.ax.set_zlabel("Z")
-
-            # let voxels look not stretched
-            max_range = 0
-            for d in range(grid_map.dim):
-                max_range = max(max_range, grid_map.bounds[d, 1] - grid_map.bounds[d, 0])
-            self.ax.set_xlim(grid_map.bounds[0, 0], grid_map.bounds[0, 0] + max_range)
-            self.ax.set_ylim(grid_map.bounds[1, 0], grid_map.bounds[1, 0] + max_range)
-            self.ax.set_zlim(grid_map.bounds[2, 0], grid_map.bounds[2, 0] + max_range)
-
-            if equal:
-                self.ax.set_box_aspect([1,1,1])
-
-        else:
-            raise NotImplementedError(f"Grid map with dim={grid_map.dim} not supported.")
+                alpha=alpha_esdf,
+                zorder=self.zorder['esdf'],
+            )
+            plt.colorbar(label="ESDF distance")
+            
+        if equal: 
+            plt.axis("equal")
 
     def plot_expand_tree(self, expand_tree: Dict[Union[Tuple[int, ...], Tuple[float, ...]], Node], 
-                        node_color: str = "C5", 
-                        edge_color: str = "C6", 
+                        node_color: str = "#8c564b", 
+                        edge_color: str = "#e377c2", 
                         node_size: float = 5, 
                         linewidth: float = 1.0, 
                         node_alpha: float = 1.0,
@@ -181,29 +144,16 @@ class Visualizer:
             connect_to_parent: Whether to draw parent-child connections.
             map_frame: whether path is in map frame or not (world frame)
         """
-        if self.dim == 2:
-            for coord, node in expand_tree.items():
-                current = node.current
-                if map_frame:
-                    current = self.grid_map.map_to_world(current)
+        if not isinstance(expand_tree, list):   # for multiple trees
+            expand_tree = [expand_tree]
 
-                self.ax.scatter(current[0], current[1],
-                                c=node_color, s=node_size, zorder=self.zorder['expand_tree_node'], alpha=node_alpha)
-                if connect_to_parent and node.parent is not None:
-                    parent = node.parent
-                    if map_frame:
-                        parent = self.grid_map.map_to_world(parent)
-                    self.ax.plot([parent[0], current[0]],
-                                [parent[1], current[1]],
-                                color=edge_color, linewidth=linewidth, zorder=self.zorder['expand_tree_edge'], alpha=edge_alpha)
-
-        elif self.dim == 3:
-            for coord, node in expand_tree.items():
+        for tree in expand_tree:
+            for coord, node in tree.items():
                 current = node.current
                 if map_frame:
                     current = self.grid_map.map_to_world(current)
 
-                self.ax.scatter(current[0], current[1], current[2],
+                self.ax.scatter(current[0], current[1],
                                 c=node_color, s=node_size, zorder=self.zorder['expand_tree_node'], alpha=node_alpha)
                 if connect_to_parent and node.parent is not None:
                     parent = node.parent
@@ -211,13 +161,8 @@ class Visualizer:
                         parent = self.grid_map.map_to_world(parent)
                     self.ax.plot([parent[0], current[0]],
                                 [parent[1], current[1]],
-                                [parent[2], current[2]],
                                 color=edge_color, linewidth=linewidth, zorder=self.zorder['expand_tree_edge'], alpha=edge_alpha)
 
-        else:
-            raise ValueError("Dimension must be 2 or 3")
-
-
     def plot_path(self, path: List[Union[Tuple[int, ...], Tuple[float, ...]]], 
                     style: str = "-", color: str = "#13ae00", label: str = None, 
                     linewidth: float = 3, marker: str = None, map_frame: bool = True) -> None:
@@ -227,7 +172,7 @@ class Visualizer:
 
         Args:
             path: point list of path
-            style: style of path
+            style: style of path 
             color: color of path
             label: label of path
             linewidth: linewidth of path
@@ -242,12 +187,7 @@ class Visualizer:
 
         path = np.array(path)
         
-        if self.dim == 2:
-            self.ax.plot(path[:, 0], path[:, 1], style, lw=linewidth, color=color, label=label, marker=marker, zorder=self.zorder['path_2d'])
-        elif self.dim == 3:
-            self.ax.plot(path[:, 0], path[:, 1], path[:, 2], style, lw=linewidth, color=color, label=label, marker=marker, zorder=self.zorder['path_3d'])
-        else:
-            raise ValueError("Dimension not supported")
+        self.ax.plot(path[:, 0], path[:, 1], style, lw=linewidth, color=color, label=label, marker=marker, zorder=self.zorder['path_2d'])
 
         if label:
             self.ax.legend()
@@ -270,19 +210,13 @@ class Visualizer:
         text = self.ax.text(*robot.pos, robot.text, color=robot.text_color, ha='center', va='center', 
                             fontsize=fontsize, zorder=self.zorder['robot_text'])
 
-        if robot.dim == 2:
-            theta = robot.orient[0]
-            dx = np.cos(theta) * robot.radius
-            dy = np.sin(theta) * robot.radius
-            orient_patch = self.ax.arrow(robot.pos[0], robot.pos[1], dx, dy,
-                                         head_width=0.1*robot.radius, head_length=0.2*robot.radius,
-                                         fc=robot.color, ec=robot.text_color, zorder=self.zorder['robot_orient'])
-            return patch, text, orient_patch
-        elif robot.dim == 3:
-            # TODO: quiver for 3D vector
-            return patch, text
-        else:
-            return patch, text
+        theta = robot.orient[0]
+        dx = np.cos(theta) * robot.radius
+        dy = np.sin(theta) * robot.radius
+        orient_patch = self.ax.arrow(robot.pos[0], robot.pos[1], dx, dy,
+                                        head_width=0.1*robot.radius, head_length=0.2*robot.radius,
+                                        fc=robot.color, ec=robot.text_color, zorder=self.zorder['robot_orient'])
+        return patch, text, orient_patch
 
     def render_toy_simulator(self, env: ToySimulator, controllers: Dict[str, BaseController],
             steps: int = 1000, interval: int = None,
@@ -313,7 +247,7 @@ class Visualizer:
         else:
             traj_color = {rid: traj_kwargs.get("color") for rid, robot in env.robots.items()}
 
-        # 先画静态的地图和路径
+        # Draw static map and paths
         self.ax.clear()
         self.plot_grid_map(env.obstacle_grid, **grid_kwargs)
 
@@ -411,74 +345,6 @@ class Visualizer:
             self.fig, update, frames=steps+prepare_frames, interval=interval, blit=True, repeat=False
         )
 
-    def get_traj_info(self, rid: int, goal_pose: np.ndarray, goal_dist_tol: float, goal_orient_tol: float) -> Dict[str, Any]:
-        """
-        Get trajectory information.
-
-        Args:
-            rid: Robot ID.
-            goal_pose: Goal pose.
-            goal_dist_tol: Distance tolerance for goal.
-            goal_orient_tol: Orientation tolerance for goal.
-        """
-        traj = self.trajs[rid]
-
-        info = {
-            "traj_length": 0.0,
-            "success": False,
-            "dist_success": False, 
-            "oracle_success": False,
-            "oracle_dist_success": False,
-            "success_time": None,
-            "dist_success_time": None,
-            "oracle_success_time": None,
-            "oracle_dist_success_time": None,
-        }
-
-        for i in range(len(traj["poses"])):
-            pose = traj["poses"][i]
-            time = traj["time"][i]
-            
-            pos = pose[:self.dim]
-            orient = pose[self.dim:]
-            goal_pos = goal_pose[:self.dim]
-            goal_orient = goal_pose[self.dim:]
-
-            if i > 0:
-                info["traj_length"] += np.linalg.norm(pos - traj["poses"][i-1][:self.dim])
-
-            if np.linalg.norm(pos - goal_pos) < goal_dist_tol:
-                if not info["oracle_dist_success"]:
-                    info["oracle_dist_success"] = True
-                    info["oracle_dist_success_time"] = time
-
-                if not info["dist_success"]:
-                    info["dist_success"] = True
-                    info["dist_success_time"] = time
-
-                if np.abs(Geometry.regularize_orient(orient - goal_orient)) < goal_orient_tol:
-                    if not info["oracle_success"]:
-                        info["oracle_success"] = True
-                        info["oracle_success_time"] = time  
-
-                    if not info["success"]:
-                        info["success"] = True
-                        info["success_time"] = time
-                    
-                else:
-                    info["success"] = False
-                    info["success_time"] = None
-                
-            else:
-                info["success"] = False
-                info["success_time"] = None
-                info["dist_success"] = False
-                info["dist_success_time"] = None
-
-        info["traj_length"] = float(info["traj_length"])
-        return info
-
-
     def set_title(self, title: str) -> None:
         """
         Set title.
@@ -509,7 +375,18 @@ class Visualizer:
         Update plot.
         """
         self.fig.canvas.draw_idle()
-    
+
+    def savefig(self, filename, *args, **kwargs):
+        """
+        Save figure. 
+
+        Args:
+            filename: Filename to save.
+            *args: For 2D, see matplotlib.pyplot.savefig. For 3D, see pyvista.Plotter.screenshot.
+            **kwargs: For 2D, see matplotlib.pyplot.savefig. For 3D, see pyvista.Plotter.screenshot.
+        """
+        plt.savefig(fname=filename, *args, **kwargs)
+
     def show(self):
         """
         Show plot.

python_motion_planning/common/visualizer/visualizer_3d.py

@@ -0,0 +1,242 @@
+"""
+@file: visualizer_3d.py
+@author: Wu Maojia
+@update: 2025.12.20
+"""
+from typing import Union, Dict, List, Tuple, Any
+from collections import namedtuple
+import time
+import os
+
+import numpy as np
+import pyvista as pv
+
+from python_motion_planning.common.visualizer.base_visualizer import BaseVisualizer
+from python_motion_planning.controller import BaseController
+from python_motion_planning.common.env import TYPES, ToySimulator, Grid, CircularRobot, Node
+from python_motion_planning.common.utils import Geometry
+
+class Visualizer3D(BaseVisualizer):
+    """
+    Simple 3D visualizer for motion planning using pyvista.
+
+    Args:
+        window_size: Window size (width, height) (pyvista window size, unit: pixel).
+        off_screen: `off_screen` argument for pyvista. Renders off screen when True. Useful for automated screenshots.
+        show_axes: Whether to show axes for pyvista.
+        cmap_dict: Color map for 3d voxel visualization.
+    """
+    def __init__(self,  
+                window_size: tuple = (1200, 900),
+                off_screen: bool = False,
+                show_axes: bool = True,
+                cmap_dict: dict = {
+                    TYPES.FREE: "#ffffff",
+                    TYPES.OBSTACLE: "#000000",
+                    TYPES.START: "#ff0000",
+                    TYPES.GOAL: "#1155cc",
+                    TYPES.INFLATION: "#ffccff",
+                    TYPES.EXPAND: "#eeeeee",
+                    TYPES.CUSTOM: "#bbbbbb",
+                }
+            ):
+        super().__init__()
+        self.pv_plotter = pv.Plotter(window_size=list(window_size), off_screen=off_screen)
+        if show_axes: 
+            self.pv_plotter.show_axes()
+        self.pv_actors = {} 
+
+        # colors
+        self.cmap_dict = cmap_dict
+
+    def plot_grid_map(self, grid_map: Grid, equal: bool = False, alpha_3d: dict = {
+                            TYPES.FREE: 0.0,
+                            TYPES.OBSTACLE: 0.1,
+                            TYPES.START: 0.5,
+                            TYPES.GOAL: 0.5,
+                            TYPES.INFLATION: 0.0,
+                            TYPES.EXPAND: 0.01,
+                            TYPES.CUSTOM: 0.1,
+                        }) -> None:
+        '''
+        Plot grid map with static obstacles.
+
+        Args:
+            map: Grid map or its type map.
+            equal: Whether to set axis equal.
+            alpha_3d: Alpha of occupancy for 3d visualization.
+        '''
+        if grid_map.dim != 3:
+            raise ValueError(f"Grid map dimension must be 3.")
+        
+        self.grid_map = grid_map
+        self.dim = grid_map.dim
+        type_data = grid_map.type_map.data
+
+        nx, ny, nz = type_data.shape
+
+        for key, color in self.cmap_dict.items():
+            alpha = alpha_3d.get(key, 0.0)
+            if alpha < 1e-6:
+                continue
+
+            mask = (type_data == key)
+            if not np.any(mask):
+                continue
+
+            # voxels
+            points = np.argwhere(mask)
+
+            # map → world
+            points = np.array([
+                self.grid_map.map_to_world(p)
+                for p in points
+            ])
+
+            cloud = pv.PolyData(points)
+            glyph = cloud.glyph(
+                geom=pv.Cube(),
+                scale=False,
+                factor=self.grid_map.resolution
+            )
+
+            actor = self.pv_plotter.add_mesh(
+                glyph,
+                color=color,
+                opacity=alpha,
+                show_edges=False
+            )
+
+            self.pv_actors[f"voxels_{key}"] = actor
+
+    def plot_expand_tree(self, expand_tree: Dict[Union[Tuple[int, ...], Tuple[float, ...]], Node], 
+                        edge_color: str = "#e377c2", 
+                        linewidth: float = 1.0, 
+                        node_alpha: float = 1.0,
+                        edge_alpha: float = 1.0,
+                        map_frame: bool = True) -> None:
+        """
+        Visualize an expand tree (e.g. RRT).
+        
+        Args:
+            expand_tree: Dict mapping coordinate tuple -> Node (world frame).
+            edge_color: Color of the edges (parent -> child).
+            linewidth: Line width of edges.
+            map_frame: whether path is in map frame or not (world frame)
+        """
+        if not isinstance(expand_tree, list):
+            expand_tree = [expand_tree]
+
+        points = []
+        lines = []
+
+        idx = 0
+        for tree in expand_tree:
+            for _, node in tree.items():
+                cur = node.current
+                if map_frame:
+                    cur = self.grid_map.map_to_world(cur)
+
+                points.append(cur)
+
+                if node.parent is not None:
+                    parent = node.parent
+                    if map_frame:
+                        parent = self.grid_map.map_to_world(parent)
+
+                    points.append(parent)
+                    lines.append([idx, idx + 1])
+                    idx += 2
+                else:
+                    idx += 1
+
+        if not points:
+            return
+
+        points = np.array(points)
+        poly = pv.PolyData(points)
+
+        if lines:
+            cells = np.hstack([[2, l[0], l[1]] for l in lines])
+            poly.lines = cells
+
+        self.pv_plotter.add_mesh(
+            poly,
+            color=edge_color,
+            line_width=linewidth
+        )
+
+
+    def plot_path(self, path: List[Union[Tuple[int, ...], Tuple[float, ...]]], 
+                    color: str = "#13ae00", 
+                    linewidth: float = 5, map_frame: bool = True) -> None:
+        '''
+        Plot path-like information.
+        The meaning of parameters are similar to pyvista.Plotter.add_mesh (https://docs.pyvista.org/api/plotting/_autosummary/pyvista.plotter.add_mesh#pyvista.Plotter.add_mesh).
+
+        Args:
+            path: point list of path
+            color: color of path
+            linewidth: linewidth of path
+            map_frame: whether path is in map frame or not (world frame)
+        '''
+        if len(path) == 0:
+            return
+
+        if map_frame:
+            path = [self.grid_map.map_to_world(point) for point in path]
+
+        path = np.array(path)
+
+        path_line = pv.lines_from_points(path)
+        self.pv_plotter.add_mesh(
+            path_line,
+            color=color,
+            line_width=linewidth
+        )
+
+    def set_title(self, title: str) -> None:
+        """
+        Set title.
+
+        Args:
+            title: Title.
+        """
+        self.pv_plotter.add_text(title, position='upper_edge', font_size=14, color='black')
+
+    def clean(self):
+        """
+        Clean plot.
+        """
+        self.pv_plotter.clear()
+        self.pv_actors = {}
+
+    def update(self):
+        """
+        Update plot.
+        """
+        self.pv_plotter.render()
+
+    def savefig(self, filename, *args, **kwargs):
+        """
+        Save figure. 
+
+        Args:
+            filename: Filename to save.
+            *args: See pyvista.Plotter.screenshot.
+            **kwargs: See pyvista.Plotter.screenshot.
+        """
+        self.pv_plotter.screenshot(filename=filename, *args, **kwargs)
+
+    def show(self):
+        """
+        Show plot.
+        """
+        self.pv_plotter.reset_camera()
+        self.pv_plotter.show(interactive=True, auto_close=False)
+
+    def close(self):
+        """
+        Close plot.
+        """
+        self.pv_plotter.close()