update csv export and added fixed mesh axis drawing

examples/viz/offscreen_mesh.py

@@ -84,6 +84,168 @@ def _background_rgb01(color: tuple[int, int, int]) -> np.ndarray:
     return np.asarray(_background_rgba01(color)[:3], dtype=np.float64)
 
 
+# Cached pre-rendered mesh-axis triads keyed by size_px.
+_AXIS_TRIAD_CACHE: dict[int, np.ndarray] = {}
+
+
+def _project_mesh_axis_to_screen(point_xyz: np.ndarray, *, size_px: int) -> tuple[int, int, float]:
+    """Project mesh coords to HUD pixels.
+
+    Screen mapping for the fixed mesh-axis reference:
+      +X -> down, +Y -> left, +Z -> inward (into the page).
+
+    Mild foreshortening keeps +Z visible as a short into-page axis.
+    Returns (px, py, depth) where larger depth is farther from the viewer.
+    """
+    x, y, z = float(point_xyz[0]), float(point_xyz[1]), float(point_xyz[2])
+    # Camera looks predominantly along +Z (into page): u right, v down.
+    # Soft shear so Z has a short foreshortened stub.
+    shear = 0.34
+    u = -y - shear * z
+    v = x - shear * z
+    scale = size_px * 0.30
+    # Origin slightly lower-right-of-center so labeled tips stay in-frame.
+    cx = size_px * 0.58
+    cy = size_px * 0.48
+    return int(round(cx + u * scale)), int(round(cy + v * scale)), z
+
+
+def _render_axis_triad_rgba(size_px: int) -> np.ndarray:
+    """Build (and cache) a transparent RGBA 3D-looking XYZ triad."""
+    import cv2
+
+    size_px = max(48, int(size_px))
+    cached = _AXIS_TRIAD_CACHE.get(size_px)
+    if cached is not None:
+        return cached
+
+    rgba = np.zeros((size_px, size_px, 4), dtype=np.uint8)
+    rgb = rgba[:, :, :3]
+    alpha = rgba[:, :, 3]
+
+    # Soft circular plate for contrast on bright contact heatmaps.
+    plate = np.zeros((size_px, size_px), dtype=np.uint8)
+    center = (size_px // 2, size_px // 2)
+    radius = int(size_px * 0.46)
+    cv2.circle(plate, center, radius, 255, thickness=-1, lineType=cv2.LINE_AA)
+    rgba[:, :, 0:3][plate > 0] = (18, 18, 18)
+    alpha[plate > 0] = 140
+
+    thickness = max(2, size_px // 42)
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    font_scale = max(0.35, size_px / 220.0)
+    origin = np.zeros(3, dtype=np.float64)
+    o_xy = _project_mesh_axis_to_screen(origin, size_px=size_px)[:2]
+
+    axes = (
+        ("X", np.array([1.0, 0.0, 0.0]), (230, 70, 70)),
+        ("Y", np.array([0.0, 1.0, 0.0]), (70, 210, 95)),
+        ("Z", np.array([0.0, 0.0, 1.0]), (70, 140, 255)),
+    )
+    # Paint far axes first so nearer shafts occlude correctly.
+    painted = []
+    for label, direction, color in axes:
+        tip = direction * 1.0
+        px, py, depth = _project_mesh_axis_to_screen(tip, size_px=size_px)
+        painted.append((depth, label, color, (px, py)))
+    painted.sort(key=lambda item: item[0], reverse=True)
+
+    layer = np.zeros((size_px, size_px, 3), dtype=np.uint8)
+    for _, label, color, tip_xy in painted:
+        cv2.arrowedLine(
+            layer,
+            o_xy,
+            tip_xy,
+            color,
+            thickness,
+            line_type=cv2.LINE_AA,
+            tipLength=0.28,
+        )
+        (tw, th), _ = cv2.getTextSize(label, font, font_scale, max(1, thickness - 1))
+        # Offset label slightly past the arrow tip, away from origin.
+        dx = tip_xy[0] - o_xy[0]
+        dy = tip_xy[1] - o_xy[1]
+        norm = max(1.0, float(np.hypot(dx, dy)))
+        lx = int(tip_xy[0] + 8 * dx / norm - tw // 2)
+        ly = int(tip_xy[1] + 8 * dy / norm + th // 2)
+        lx = min(max(1, lx), size_px - tw - 1)
+        ly = min(max(th + 1, ly), size_px - 2)
+        cv2.putText(
+            layer,
+            label,
+            (lx, ly),
+            font,
+            font_scale,
+            color,
+            max(1, thickness - 1),
+            cv2.LINE_AA,
+        )
+    # Origin hub on the same contiguous buffer.
+    cv2.circle(layer, o_xy, max(2, thickness), (235, 235, 235), -1, lineType=cv2.LINE_AA)
+
+    mask = np.any(layer > 0, axis=2)
+    rgb[mask] = layer[mask]
+    alpha[mask] = 255
+
+    _AXIS_TRIAD_CACHE[size_px] = rgba
+    return rgba
+
+
+def draw_axis_overlay(
+    rgb: np.ndarray,
+    *,
+    enabled: bool = True,
+    padding_px: int = 12,
+    size_px: int = 96,
+    opacity: float = 1.0,
+) -> np.ndarray:
+    """Composite a fixed top-right mesh-axis triad onto an RGB frame.
+
+    The triad is pre-rendered and cached; orientation is locked to true mesh
+    axes (+X down, +Y left, +Z inward) and does not follow free-rotate.
+    """
+    if not enabled or rgb is None or rgb.size == 0:
+        return rgb
+
+    out = np.ascontiguousarray(rgb)
+    if out.ndim != 3 or out.shape[2] < 3:
+        return out
+
+    triad = _render_axis_triad_rgba(size_px)
+    th, tw = triad.shape[:2]
+    pad = max(0, int(padding_px))
+    h, w = out.shape[:2]
+    x0 = max(0, w - tw - pad)
+    y0 = max(0, pad)
+    x1 = min(w, x0 + tw)
+    y1 = min(h, y0 + th)
+    if x1 <= x0 or y1 <= y0:
+        return out
+
+    roi = out[y0:y1, x0:x1]
+    overlay = triad[: y1 - y0, : x1 - x0]
+    alpha = (overlay[:, :, 3].astype(np.float32) / 255.0) * float(np.clip(opacity, 0.0, 1.0))
+    if alpha.ndim == 2:
+        alpha = alpha[:, :, None]
+    rgb_o = overlay[:, :, :3].astype(np.float32)
+    base = roi.astype(np.float32)
+    blended = rgb_o * alpha + base * (1.0 - alpha)
+    roi[:] = np.clip(blended, 0, 255).astype(np.uint8)
+    return out
+
+
+def _apply_axis_overlay(rgb: np.ndarray, axis_overlay) -> np.ndarray:
+    if axis_overlay is None:
+        return draw_axis_overlay(rgb)
+    return draw_axis_overlay(
+        rgb,
+        enabled=bool(getattr(axis_overlay, "enabled", True)),
+        padding_px=int(getattr(axis_overlay, "padding_px", 12)),
+        size_px=int(getattr(axis_overlay, "size_px", 96)),
+        opacity=float(getattr(axis_overlay, "opacity", 1.0)),
+    )
+
+
 class _FilamentOffscreenBackend:
     """Filament OffscreenRenderer (Linux headless / platforms with EGL)."""
 
@@ -104,6 +266,7 @@ class _FilamentOffscreenBackend:
         mesh_background: tuple[int, int, int] = (0, 0, 0),
         contact_arrow_style=None,
         flow_arrow_style=None,
+        axis_overlay=None,
     ):
         import open3d as o3d
 
@@ -119,6 +282,7 @@ class _FilamentOffscreenBackend:
         self.lookat = lookat
         self.mesh_lighting = mesh_lighting
         self.mesh_background = mesh_background
+        self.axis_overlay = axis_overlay
         self._shade_factors = (
             precompute_shade_factors(mesh.rest_vertices, mesh.triangles, mesh_lighting)
             if mesh_lighting is not None
@@ -257,7 +421,8 @@ class _FilamentOffscreenBackend:
             mesh_to_o3d=lambda v, t, c: _mesh_to_o3d(o3d, v, t, c),
         )
 
-        return _numpy_o3d_image_to_rgb(self._renderer.render_to_image())
+        rgb = _numpy_o3d_image_to_rgb(self._renderer.render_to_image())
+        return _apply_axis_overlay(rgb, self.axis_overlay)
 
     def close(self) -> None:
         self._renderer = None
@@ -283,6 +448,7 @@ class _HiddenVisualizerBackend:
         mesh_background: tuple[int, int, int] = (0, 0, 0),
         contact_arrow_style=None,
         flow_arrow_style=None,
+        axis_overlay=None,
     ):
         import open3d as o3d
 
@@ -298,6 +464,7 @@ class _HiddenVisualizerBackend:
         self.lookat = lookat
         self.mesh_lighting = mesh_lighting
         self.mesh_background = mesh_background
+        self.axis_overlay = axis_overlay
         self._shade_factors = (
             precompute_shade_factors(mesh.rest_vertices, mesh.triangles, mesh_lighting)
             if mesh_lighting is not None
@@ -441,7 +608,8 @@ class _HiddenVisualizerBackend:
 
         self.vis.poll_events()
         self.vis.update_renderer()
-        return _numpy_o3d_image_to_rgb(self.vis.capture_screen_float_buffer(do_render=False))
+        rgb = _numpy_o3d_image_to_rgb(self.vis.capture_screen_float_buffer(do_render=False))
+        return _apply_axis_overlay(rgb, self.axis_overlay)
 
     def close(self) -> None:
         if self.vis is not None:
@@ -478,6 +646,7 @@ class OffscreenFingerMeshRenderer:
         mesh_background: tuple[int, int, int] = (0, 0, 0),
         contact_arrow_style=None,
         flow_arrow_style=None,
+        axis_overlay=None,
     ):
         lookat = lookat if lookat is not None else mesh.rest_vertices.mean(axis=0)
         self._backend = _create_backend(
@@ -495,6 +664,7 @@ class OffscreenFingerMeshRenderer:
             mesh_background=mesh_background,
             contact_arrow_style=contact_arrow_style,
             flow_arrow_style=flow_arrow_style,
+            axis_overlay=axis_overlay,
         )
 
     @property