Source code for embodichain.lab.sim.motion.workspace.visualizers.point_cloud_visualizer

# ----------------------------------------------------------------------------
# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------------

from __future__ import annotations

import numpy as np
import torch
from typing import Union, Any, Dict
from pathlib import Path

from embodichain.lab.sim.motion.workspace.configs import (
    VisualizationType,
)
from embodichain.lab.sim.motion.workspace.visualizers.base_visualizer import (
    BaseVisualizer,
    OPEN3D_AVAILABLE,
)

if OPEN3D_AVAILABLE:
    import open3d as o3d

from embodichain.utils import logger

__all__ = ["PointCloudVisualizer"]


[docs] class PointCloudVisualizer(BaseVisualizer): """Point cloud visualizer using SimulationManager, Viser, or local tools. Attributes: point_size: Size of points in visualization. """
[docs] def __init__( self, backend: str = "sim_manager", point_size: float = 2.0, config: Dict[str, Any] | None = None, sim_manager: Any | None = None, control_part_name: str | None = None, ): """Initialize the point cloud visualizer. Args: backend: Visualization backend ('sim_manager', 'viser', 'open3d', 'matplotlib', or 'data'). Defaults to 'sim_manager'. The 'data' backend returns raw data without visualization. point_size: Size of points in visualization. Defaults to 2.0. config: Optional configuration dictionary. Defaults to None. sim_manager: SimulationManager instance for the 'sim_manager' or 'viser' backend. Defaults to None. control_part_name: Control part name used to name the point cloud. Defaults to None. """ super().__init__(backend, config) self.point_size = point_size self.sim_manager = sim_manager self.control_part_name = control_part_name
[docs] def visualize( self, points: torch.Tensor | np.ndarray, colors: torch.Tensor | np.ndarray | None = None, **kwargs: Any, ) -> Any: """Visualize points as a point cloud. Args: points: Array of shape (N, 3) containing point positions. colors: Optional array of shape (N, 3) or (N, 4) containing colors. **kwargs: Additional visualization parameters: - point_size: Override default point size Returns: Backend-specific point-cloud handle, Viser overlay, geometry, or matplotlib figure. Examples: >>> visualizer = PointCloudVisualizer() >>> points = np.random.rand(1000, 3) >>> colors = np.random.rand(1000, 3) >>> pcd = visualizer.visualize(points, colors) >>> visualizer.show() """ # Convert to numpy points = self._to_numpy(points) self._validate_points(points) # Get visualization parameters point_size = kwargs.get("point_size", self.point_size) # Validate and prepare colors colors = self._validate_colors(colors, len(points)) if colors is None: colors = self._get_default_colors(len(points)) # Convert to RGB if RGBA if colors.shape[1] == 4: colors = colors[:, :3] if self.backend == "data": # Return raw data for user to handle data = { "points": points, "colors": colors, "point_size": point_size, "type": "point_cloud", } self._last_visualization = {"data": data} return data elif self.backend == "sim_manager": pcd_handle = self._create_sim_manager_point_cloud( points, colors, point_size ) self._last_visualization = { "point_cloud_handle": pcd_handle, "point_size": point_size, } return pcd_handle elif self.backend == "viser": overlay = self._create_viser_point_cloud(points, colors, point_size) self._last_visualization = { "point_cloud_overlay": overlay, "point_size": point_size, } return overlay elif self.backend == "open3d": pcd = self._create_open3d_point_cloud(points, colors) self._last_visualization = { "point_cloud": pcd, "point_size": point_size, } return pcd elif self.backend == "matplotlib": fig = self._create_matplotlib_point_cloud(points, colors, point_size) self._last_visualization = {"figure": fig} return fig else: raise ValueError(f"Unsupported backend: {self.backend}")
def _create_sim_manager_point_cloud( self, points: np.ndarray, colors: np.ndarray, point_size: float ) -> Any: if self.sim_manager is None: raise ValueError("sim_manager is required for 'sim_manager' backend") pcd_name = f"workspace_pcd_{self.control_part_name or 'default'}" return self.sim_manager.visualize_point_cloud( points=points, colors=colors, point_size=point_size, name=pcd_name, ) def _create_viser_point_cloud( self, points: np.ndarray, colors: np.ndarray, point_size: float ) -> Any: """Publish a persistent point-cloud overlay through SimulationManager.""" if self.sim_manager is None: raise ValueError("sim_manager is required for 'viser' backend") from embodichain.lab.visualization import PointCloudOverlay, SceneOverlays colors_uint8 = np.rint(np.clip(colors, 0.0, 1.0) * 255.0).astype(np.uint8) overlay = PointCloudOverlay( overlay_id=f"workspace_{self.control_part_name or 'default'}", points=np.asarray(points, dtype=np.float32), colors=colors_uint8, point_size=point_size, ) self.sim_manager.set_visualization_overlays( SceneOverlays(point_clouds=(overlay,)) ) logger.log_info( f"Published Viser point cloud '{overlay.overlay_id}' with " f"{len(points)} points (point_size={point_size})" ) return overlay def _create_open3d_point_cloud( self, points: np.ndarray, colors: np.ndarray ) -> "o3d.geometry.PointCloud": pcd = o3d.geometry.PointCloud() pcd.points = o3d.utility.Vector3dVector(points) pcd.colors = o3d.utility.Vector3dVector(colors) logger.log_info(f"Created point cloud with {len(points)} points") return pcd def _create_matplotlib_point_cloud( self, points: np.ndarray, colors: np.ndarray, point_size: float ): import matplotlib.pyplot as plt from mpl_toolkits.mplot3d import Axes3D fig = plt.figure(figsize=(10, 8)) ax = fig.add_subplot(111, projection="3d") ax.scatter( points[:, 0], points[:, 1], points[:, 2], c=colors, s=point_size, alpha=0.6 ) ax.set_xlabel("X") ax.set_ylabel("Y") ax.set_zlabel("Z") ax.set_title("Workspace Point Cloud") return fig def _save_impl(self, filepath: Path, **kwargs: Any) -> None: if self.backend == "data": # Save data as numpy file data = self._last_visualization["data"] np.savez(filepath, **data) elif self.backend == "open3d": pcd = self._last_visualization["point_cloud"] # Determine file format from extension suffix = filepath.suffix.lower() if suffix in [".pcd", ".ply", ".xyz", ".xyzrgb", ".pts"]: o3d.io.write_point_cloud(str(filepath), pcd) elif suffix in [".png", ".jpg", ".jpeg"]: # Render to image vis = o3d.visualization.Visualizer() vis.create_window(visible=False) vis.add_geometry(pcd) # Coordinate frame removed - implement separately if needed vis.update_geometry(pcd) vis.poll_events() vis.update_renderer() vis.capture_screen_image(str(filepath)) vis.destroy_window() else: raise ValueError( f"Unsupported file format: {suffix}. " f"Use .pcd, .ply, .xyz, .xyzrgb, .pts, .png, .jpg" ) elif self.backend == "matplotlib": fig = self._last_visualization["figure"] fig.savefig(filepath, dpi=300, bbox_inches="tight") def _show_impl(self, **kwargs: Any) -> None: if self.backend == "data": logger.log_warning( "Cannot display visualization with 'data' backend. " "Use 'open3d' or 'matplotlib' backend for interactive display." ) return elif self.backend == "open3d": geometries = [self._last_visualization["point_cloud"]] # Coordinate frame removed - implement separately if needed # Set point size in visualization vis = o3d.visualization.Visualizer() vis.create_window() for geom in geometries: vis.add_geometry(geom) render_option = vis.get_render_option() render_option.point_size = self._last_visualization.get("point_size", 2.0) vis.run() vis.destroy_window() elif self.backend == "matplotlib": import matplotlib.pyplot as plt plt.show()
[docs] def get_type_name(self) -> str: return VisualizationType.POINT_CLOUD.value