orisys-sdk/examples/viz/contact_arrow.py

"""Contact normal-force arrow overlay for finger 3D mesh examples."""

from __future__ import annotations

import numpy as np

from orisys.finger.mesh import UVSurfaceMesh, compute_uv_surface_frames, mesh_uv_to_world_scale

DEFAULT_NORMAL_FORCE_SCALE = 1.0
MIN_ARROW_LENGTH_FRACTION = 0.02
MAX_ARROW_LENGTH_FRACTION = 0.30
REFERENCE_NORMAL_FORCE = 10000.0
MIN_ARROW_THICKNESS_SCALE = 0.20
MAX_ARROW_THICKNESS_SCALE = 1.0
CONTACT_ARROW_COLOR = (1.0, 0.92, 0.10)


def _mesh_span(vertices: np.ndarray) -> float:
    extent = np.ptp(vertices, axis=0)
    return float(max(np.max(extent), 1e-6))


def _rotation_matrix_z_to(direction: np.ndarray) -> np.ndarray:
    """Rotation matrix mapping unit +Z to ``direction``."""
    direction = np.asarray(direction, dtype=np.float64).reshape(3)
    norm = np.linalg.norm(direction)
    if norm < 1e-8:
        return np.eye(3, dtype=np.float64)
    direction = direction / norm

    z_axis = np.array([0.0, 0.0, 1.0], dtype=np.float64)
    cross = np.cross(z_axis, direction)
    dot = float(np.dot(z_axis, direction))
    if dot > 0.9999:
        return np.eye(3, dtype=np.float64)
    if dot < -0.9999:
        return np.diag([1.0, -1.0, -1.0]).astype(np.float64)

    skew = np.array(
        [
            [0.0, -cross[2], cross[1]],
            [cross[2], 0.0, -cross[0]],
            [-cross[1], cross[0], 0.0],
        ],
        dtype=np.float64,
    )
    return np.eye(3, dtype=np.float64) + skew + skew @ skew * (1.0 / (1.0 + dot))


def thickness_scale_from_force(force: float) -> float:
    """Map normal force to a radial scale factor (thin at low force, thicker at high force)."""
    if force <= 0.0:
        return 0.0
    t = force / REFERENCE_NORMAL_FORCE
    return float(
        np.clip(
            MIN_ARROW_THICKNESS_SCALE + (MAX_ARROW_THICKNESS_SCALE - MIN_ARROW_THICKNESS_SCALE) * t,
            MIN_ARROW_THICKNESS_SCALE,
            MAX_ARROW_THICKNESS_SCALE,
        )
    )


def _scale_arrow_template(
    template_vertices: np.ndarray,
    *,
    length_scale: float,
    thickness_scale: float,
) -> np.ndarray:
    """Scale a +Z arrow template axially and radially."""
    scaled = np.asarray(template_vertices, dtype=np.float64).copy()
    scaled[:, 0] *= thickness_scale
    scaled[:, 1] *= thickness_scale
    scaled[:, 2] *= length_scale
    return scaled


def create_arrow_template(o3d):
    """Unit-height arrow mesh along +Z; radial size is scaled per-frame from force."""
    arrow = o3d.geometry.TriangleMesh.create_arrow(
        cylinder_radius=0.03,
        cone_radius=0.06,
        cylinder_height=0.72,
        cone_height=0.28,
    )
    arrow.compute_vertex_normals()
    return (
        np.asarray(arrow.vertices, dtype=np.float64),
        np.asarray(arrow.triangles, dtype=np.int32),
    )


def build_contact_normal_arrow_mesh(
    origin: np.ndarray,
    normal: np.ndarray,
    length: float,
    force: float,
    *,
    mesh_span: float,
    template_vertices: np.ndarray,
    template_triangles: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Build a yellow 3D arrow mesh pointing opposite the surface normal."""
    origin = np.asarray(origin, dtype=np.float64).reshape(3)
    normal = np.asarray(normal, dtype=np.float64).reshape(3)
    norm = np.linalg.norm(normal)
    color = np.tile(np.asarray(CONTACT_ARROW_COLOR, dtype=np.float64), (len(template_vertices), 1))

    radial_scale = thickness_scale_from_force(force)
    if norm < 1e-8 or length <= 0.0 or radial_scale <= 0.0 or not np.all(np.isfinite(origin)):
        return np.zeros_like(template_vertices), template_triangles, color

    direction = -normal / norm
    length = float(np.clip(length, mesh_span * MIN_ARROW_LENGTH_FRACTION, mesh_span * MAX_ARROW_LENGTH_FRACTION))
    if length <= 0.0:
        return np.zeros_like(template_vertices), template_triangles, color

    rotation = _rotation_matrix_z_to(direction)
    scaled_template = _scale_arrow_template(
        template_vertices,
        length_scale=length,
        thickness_scale=radial_scale,
    )
    vertices = scaled_template @ rotation.T + origin
    return vertices, template_triangles, color


def create_arrow_mesh_geometry(o3d, template_vertices: np.ndarray, template_triangles: np.ndarray):
    """Create a persistent Open3D mesh for the contact arrow overlay."""
    arrow_mesh = o3d.geometry.TriangleMesh()
    arrow_mesh.vertices = o3d.utility.Vector3dVector(np.zeros_like(template_vertices))
    arrow_mesh.triangles = o3d.utility.Vector3iVector(template_triangles)
    arrow_mesh.vertex_colors = o3d.utility.Vector3dVector(
        np.tile(np.asarray(CONTACT_ARROW_COLOR, dtype=np.float64), (len(template_vertices), 1))
    )
    arrow_mesh.compute_vertex_normals()
    return arrow_mesh


class ContactArrowOverlay:
    """Example-layer contact arrow driven by centroid and normal force."""

    def __init__(
        self,
        mesh: UVSurfaceMesh,
        *,
        normal_force_scale: float = DEFAULT_NORMAL_FORCE_SCALE,
        enabled: bool = True,
    ):
        import open3d as o3d

        self.mesh = mesh
        self.normal_force_scale = float(normal_force_scale)
        self.enabled = enabled
        self._mesh_span = _mesh_span(mesh.rest_vertices)
        self._pixel_to_world = mesh_uv_to_world_scale(mesh.world_xyz, mesh.valid_mask)
        _, _, self._surface_normals = compute_uv_surface_frames(mesh.world_xyz, mesh.valid_mask)
        self._template_vertices, self._template_triangles = (
            create_arrow_template(o3d) if enabled else (None, None)
        )
        self.arrow_mesh = (
            create_arrow_mesh_geometry(o3d, self._template_vertices, self._template_triangles)
            if enabled
            else None
        )

    def compute(
        self,
        centroid,
        centroid_found: bool,
        fnormal: float,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray] | None:
        if not self.enabled or self.arrow_mesh is None:
            return None

        template_vertices = self._template_vertices
        template_triangles = self._template_triangles
        color = np.tile(
            np.asarray(CONTACT_ARROW_COLOR, dtype=np.float64),
            (len(template_vertices), 1),
        )
        hidden = (np.zeros_like(template_vertices), template_triangles, color)

        if not centroid_found or centroid is None:
            return hidden

        centroid = np.asarray(centroid, dtype=np.float64).reshape(-1)
        if centroid.size < 2:
            return hidden

        cx, cy = int(round(centroid[0])), int(round(centroid[1]))
        h, w = self.mesh.shape
        if not (0 <= cy < h and 0 <= cx < w):
            return hidden
        if self.mesh.index_map[cy, cx] < 0:
            return hidden

        origin = self.mesh.world_xyz[cy, cx]
        normal = self._surface_normals[cy, cx]
        if not np.all(np.isfinite(origin)) or not np.all(np.isfinite(normal)):
            return hidden

        force = max(float(fnormal), 0.0)
        if force <= 0.0:
            return hidden

        length = force * self.normal_force_scale * self._pixel_to_world
        return build_contact_normal_arrow_mesh(
            origin,
            normal,
            length,
            force,
            mesh_span=self._mesh_span,
            template_vertices=template_vertices,
            template_triangles=template_triangles,
        )

    def apply(self, arrow_vertices: np.ndarray, arrow_colors: np.ndarray) -> None:
        import open3d as o3d

        self.arrow_mesh.vertices = o3d.utility.Vector3dVector(arrow_vertices.astype(np.float64))
        self.arrow_mesh.vertex_colors = o3d.utility.Vector3dVector(arrow_colors.astype(np.float64))
        self.arrow_mesh.compute_vertex_normals()