# ----------------------------------------------------------------------------
# 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.
# ----------------------------------------------------------------------------
import numpy as np
import torch
from typing import Union, Dict, Any
from pathlib import Path
try:
import open3d as o3d
OPEN3D_AVAILABLE = True
except ImportError:
OPEN3D_AVAILABLE = False
try:
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
MATPLOTLIB_AVAILABLE = True
except ImportError:
MATPLOTLIB_AVAILABLE = False
from embodichain.utils import logger
from embodichain.lab.sim.motion.workspace.visualizers.base_visualizer import (
BaseVisualizer,
)
from embodichain.lab.sim.motion.workspace.configs.visualization_config import (
VisualizationType,
)
__all__ = ["AxisVisualizer"]
[docs]
class AxisVisualizer(BaseVisualizer):
"""Visualizer for coordinate axes/frames at specified poses.
This visualizer creates coordinate axes (X, Y, Z) at given transformation matrices,
useful for visualizing robot end-effector poses, workspace reference frames, etc.
Supports multiple backends:
- 'sim_manager': Uses SimulationManager.draw_marker() with MarkerCfg
- 'open3d': Creates coordinate frames using Open3D
- 'matplotlib': Draws axis lines in 3D matplotlib plot
- 'data': Returns axis data without visualization
"""
[docs]
def __init__(
self,
backend: str = "sim_manager",
axis_length: float = 0.15,
axis_size: float = 0.005,
config: Dict[str, Any] | None = None,
sim_manager: Any | None = None,
control_part_name: str | None = None,
reference_pose: Union[np.ndarray, torch.Tensor] | None = None,
arena_index: int = 0,
):
"""Initialize the axis visualizer.
Args:
backend: Visualization backend ('sim_manager', 'open3d', 'matplotlib', or 'data').
Defaults to 'sim_manager'.
axis_length: Length of each axis. Defaults to 0.15.
axis_size: Thickness/size of axes. Defaults to 0.005.
config: Optional configuration dictionary. Defaults to None.
sim_manager: SimulationManager instance for 'sim_manager' backend. Defaults to None.
control_part_name: Control part name for naming (compatibility). Defaults to None.
reference_pose: Reference pose (4x4 matrix) for orientation. Defaults to None.
arena_index: Arena index for sim_manager markers. Defaults to 0.
"""
super().__init__(backend, config)
self.axis_length = axis_length
self.axis_size = axis_size
self.sim_manager = sim_manager
self.control_part_name = control_part_name
self.reference_pose = reference_pose
self.arena_index = arena_index
[docs]
def visualize(
self,
poses: torch.Tensor | np.ndarray,
colors: torch.Tensor | np.ndarray | None = None,
**kwargs: Any,
) -> Any:
"""Visualize coordinate axes at specified poses or points.
Args:
poses: Either transformation matrices of shape (4, 4) or (N, 4, 4),
or point coordinates of shape (N, 3) which will be converted to poses.
colors: Optional colors (not used for axes but kept for interface compatibility).
**kwargs: Additional visualization parameters:
- axis_length: Override default axis length
- axis_size: Override default axis size
- name_prefix: Prefix for axis names (default: "axis")
- arena_index: Arena index for sim_manager (default: 0)
Returns:
Backend-specific axis representation.
Examples:
>>> import numpy as np
>>> visualizer = AxisVisualizer(backend='sim_manager', sim_manager=sim)
>>> # Using transformation matrix
>>> pose = np.eye(4)
>>> pose[:3, 3] = [1.0, 0.5, 1.2] # Set position
>>> result = visualizer.visualize(pose)
>>>
>>> # Using point coordinates
>>> points = np.array([[1.0, 0.5, 1.2], [2.0, 1.0, 0.8]])
>>> result = visualizer.visualize(points)
>>> visualizer.show()
"""
# Convert to numpy
poses = self._to_numpy(poses)
# Convert points to poses if needed
poses = self._convert_points_to_poses(poses)
# Validate pose dimensions
self._validate_poses(poses)
# Get parameters
axis_length = kwargs.get("axis_length", self.axis_length)
axis_size = kwargs.get("axis_size", self.axis_size)
name_prefix = kwargs.get("name_prefix", "axis")
arena_index = kwargs.get("arena_index", self.arena_index)
# Dispatch to backend implementation
if self.backend == "sim_manager":
return self._visualize_sim_manager(
poses, axis_length, axis_size, name_prefix, arena_index, **kwargs
)
elif self.backend == "open3d":
return self._visualize_open3d(
poses, axis_length, axis_size, name_prefix, **kwargs
)
elif self.backend == "matplotlib":
return self._visualize_matplotlib(
poses, axis_length, axis_size, name_prefix, **kwargs
)
elif self.backend == "data":
return self._visualize_data(
poses, axis_length, axis_size, name_prefix, **kwargs
)
else:
raise ValueError(f"Unsupported backend: {self.backend}")
def _visualize_sim_manager(
self,
poses: np.ndarray,
axis_length: float,
axis_size: float,
name_prefix: str,
arena_index: int,
**kwargs: Any,
) -> Any:
"""Visualize axes using sim_manager backend."""
if self.sim_manager is None:
raise ValueError("sim_manager is required for 'sim_manager' backend")
# Import here to avoid circular imports
from embodichain.lab.sim.cfg import MarkerCfg
axis_markers = []
# Handle single pose (4,4) or multiple poses (N,4,4)
if poses.ndim == 2:
poses = poses[np.newaxis, :, :]
for i, pose_matrix in enumerate(poses):
marker_name = f"{name_prefix}_{i}"
try:
# Create axis marker using MarkerCfg
marker_cfg = MarkerCfg(
name=marker_name,
marker_type="axis",
axis_xpos=pose_matrix, # 4x4 transformation matrix
axis_size=axis_size,
axis_len=axis_length,
arena_index=arena_index,
)
# Draw the marker
self.sim_manager.draw_marker(cfg=marker_cfg)
axis_markers.append(marker_name)
except Exception as e:
logger.log_warning(f"Failed to draw axis {marker_name}: {e}")
logger.log_info(f"Created {len(axis_markers)} coordinate axes with sim_manager")
self._last_visualization = axis_markers
return axis_markers
def _visualize_open3d(
self,
poses: np.ndarray,
axis_length: float,
axis_size: float,
name_prefix: str,
**kwargs: Any,
) -> Any:
"""Visualize axes using Open3D backend."""
if not OPEN3D_AVAILABLE:
raise RuntimeError(
"Open3D is not available. Install with: pip install open3d"
)
# Create coordinate frame geometries
frame_geometries = []
if poses.ndim == 2:
poses = poses[np.newaxis, :, :]
for i, pose_matrix in enumerate(poses):
frame = o3d.geometry.TriangleMesh.create_coordinate_frame(
size=axis_length, origin=[0, 0, 0]
)
frame.transform(pose_matrix)
frame_geometries.append(frame)
logger.log_info(
f"Created {len(frame_geometries)} coordinate frames with Open3D"
)
self._last_visualization = frame_geometries
return frame_geometries
def _visualize_matplotlib(
self,
poses: np.ndarray,
axis_length: float,
axis_size: float,
name_prefix: str,
**kwargs: Any,
) -> Any:
"""Visualize axes using matplotlib backend."""
if not MATPLOTLIB_AVAILABLE:
raise RuntimeError(
"Matplotlib is not available. Install with: pip install matplotlib"
)
# Create figure and 3D axis
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection="3d")
if poses.ndim == 2:
poses = poses[np.newaxis, :, :]
colors = ["red", "green", "blue"] # X, Y, Z axes
labels = ["X", "Y", "Z"]
for i, pose_matrix in enumerate(poses):
origin = pose_matrix[:3, 3]
for j, (color, label) in enumerate(zip(colors, labels)):
axis_direction = pose_matrix[:3, j]
axis_end = origin + axis_direction * axis_length
ax.plot3D(
[origin[0], axis_end[0]],
[origin[1], axis_end[1]],
[origin[2], axis_end[2]],
color=color,
linewidth=axis_size * 1000,
label=f"{label}_axis_{i}" if i == 0 else "",
)
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")
ax.legend()
ax.set_title(f"Coordinate Axes ({len(poses)} frames)")
logger.log_info(f"Created matplotlib plot with {len(poses)} coordinate frames")
self._last_visualization = fig
return fig
def _visualize_data(
self,
poses: np.ndarray,
axis_length: float,
axis_size: float,
name_prefix: str,
**kwargs: Any,
) -> Any:
"""Return axis data without visualization."""
axis_data = {
"poses": poses,
"axis_length": axis_length,
"axis_size": axis_size,
"name_prefix": name_prefix,
"type": "coordinate_axes",
"num_frames": len(poses) if poses.ndim == 3 else 1,
}
logger.log_info(
f"Generated axis data for {axis_data['num_frames']} coordinate frames"
)
self._last_visualization = axis_data
return axis_data
def _convert_points_to_poses(self, data: np.ndarray) -> np.ndarray:
"""Convert point coordinates to transformation matrices if needed.
Args:
data: Either points (N, 3) or poses (4, 4) or (N, 4, 4)
Returns:
Transformation matrices (4, 4) or (N, 4, 4)
"""
# Check if input is points (N, 3)
if data.ndim == 2 and data.shape[1] == 3:
# Convert points to poses
num_points = data.shape[0]
# Use reference pose if available, otherwise identity rotation
if self.reference_pose is not None:
# Convert reference pose to numpy if needed
if isinstance(self.reference_pose, torch.Tensor):
ref_pose = self.reference_pose.cpu().numpy()
else:
ref_pose = self.reference_pose
# Ensure 4x4 matrix
if ref_pose.ndim == 3:
ref_pose = ref_pose[0] # Take first pose if batch
# Create poses with reference orientation and point positions
poses = np.tile(ref_pose, (num_points, 1, 1))
poses[:, :3, 3] = data # Override translation with point positions
logger.log_debug(
f"Using reference pose orientation for {num_points} coordinate axes"
)
else:
# Fallback to identity rotation
poses = np.tile(np.eye(4), (num_points, 1, 1))
poses[:, :3, 3] = data # Set translation
logger.log_debug(
f"Using identity orientation for {num_points} coordinate axes"
)
return poses
elif data.ndim == 1 and data.shape[0] == 3:
# Single point case
if self.reference_pose is not None:
# Convert reference pose to numpy if needed
if isinstance(self.reference_pose, torch.Tensor):
pose = self.reference_pose.cpu().numpy()
else:
pose = self.reference_pose.copy()
# Ensure 2D matrix
if pose.ndim == 3:
pose = pose[0]
# Override translation
pose[:3, 3] = data
else:
# Fallback to identity
pose = np.eye(4)
pose[:3, 3] = data
return pose
else:
# Already poses, return as is
return data
def _validate_poses(self, poses: np.ndarray) -> None:
"""Validate pose array dimensions and values.
Args:
poses: Array of transformation matrices
Raises:
ValueError: If poses have invalid shape or values
"""
if poses.ndim not in [2, 3]:
raise ValueError(
f"Poses must be 2D (4,4) or 3D (N,4,4) array, got shape {poses.shape}"
)
if poses.ndim == 2:
if poses.shape != (4, 4):
raise ValueError(f"Single pose must be (4,4), got {poses.shape}")
else: # poses.ndim == 3
if poses.shape[1:] != (4, 4):
raise ValueError(f"Multiple poses must be (N,4,4), got {poses.shape}")
# Check if last row is [0, 0, 0, 1] for valid transformation matrices
if poses.ndim == 2:
poses_to_check = [poses]
else:
poses_to_check = poses
for i, pose in enumerate(poses_to_check):
expected_bottom_row = np.array([0, 0, 0, 1])
if not np.allclose(pose[3, :], expected_bottom_row, atol=1e-6):
logger.log_warning(
f"Pose {i} bottom row {pose[3, :]} != [0,0,0,1]. May not be a valid transformation matrix."
)
[docs]
def get_type_name(self) -> str:
"""Return the type name for this visualizer."""
return VisualizationType.AXIS.value
def _save_impl(self, filepath: Path, **kwargs: Any) -> None:
"""Save the visualization to file."""
if self._last_visualization is None:
raise RuntimeError("No visualization to save. Call visualize() first.")
if self.backend == "open3d":
if filepath.suffix.lower() in [".ply", ".obj", ".stl"]:
# Save combined mesh
combined_mesh = o3d.geometry.TriangleMesh()
for frame in self._last_visualization:
combined_mesh += frame
o3d.io.write_triangle_mesh(str(filepath), combined_mesh)
logger.log_info(f"Saved Open3D coordinate frames to {filepath}")
else:
logger.log_warning(
f"Unsupported file format {filepath.suffix} for Open3D backend"
)
elif self.backend == "matplotlib":
self._last_visualization.savefig(filepath, **kwargs)
logger.log_info(f"Saved matplotlib plot to {filepath}")
elif self.backend in ["sim_manager", "data"]:
# Save as numpy file
if self.backend == "sim_manager":
data_to_save = {
"axis_names": self._last_visualization,
"type": "sim_manager_axes",
}
else:
data_to_save = self._last_visualization
np.save(filepath.with_suffix(".npy"), data_to_save, allow_pickle=True)
logger.log_info(f"Saved axis data to {filepath.with_suffix('.npy')}")
def _show_impl(self, **kwargs: Any) -> None:
"""Display the visualization."""
if self._last_visualization is None:
logger.log_warning("No visualization to show. Call visualize() first.")
return
if self.backend == "open3d":
o3d.visualization.draw_geometries(self._last_visualization)
elif self.backend == "matplotlib":
plt.show()
elif self.backend == "sim_manager":
logger.log_info(
f"Axes are displayed in simulation. Marker names: {self._last_visualization}"
)
elif self.backend == "data":
logger.log_info(
f"Data backend - no visual display. Use .save() to export data."
)