ykcheungab
2e5c047d1e
2. centralixzed arrow config in layout.py 3. seperate mesh render and display resolution to speed up (~40 fps)
307 lines
10 KiB
Python
307 lines
10 KiB
Python
"""Contact normal-force arrow overlay for finger 3D mesh examples."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, replace
|
|
from typing import Literal
|
|
|
|
import numpy as np
|
|
|
|
from orisys.finger.mesh import UVSurfaceMesh, compute_uv_surface_frames
|
|
|
|
ForceStrengthMode = Literal["log", "power", "linear"]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ContactArrowStyle:
|
|
"""Visual tuning for the 3D contact normal arrow (dashboard: edit MESH_CONTACT_ARROW in layout.py)."""
|
|
|
|
normal_force_scale: float = 1.0
|
|
min_length_fraction: float = 0.02
|
|
max_length_fraction: float = 0.30
|
|
min_thickness_scale: float = 0.20
|
|
max_thickness_scale: float = 1.0
|
|
reference_normal_force: float = 10000.0
|
|
force_strength_mode: ForceStrengthMode = "log"
|
|
force_strength_gamma: float = 0.5
|
|
color: tuple[float, float, float] = (1.0, 0.92, 0.10)
|
|
|
|
|
|
DEFAULT_CONTACT_ARROW_STYLE = ContactArrowStyle()
|
|
|
|
# Backward-compatible aliases (match DEFAULT_CONTACT_ARROW_STYLE).
|
|
DEFAULT_NORMAL_FORCE_SCALE = DEFAULT_CONTACT_ARROW_STYLE.normal_force_scale
|
|
MIN_ARROW_LENGTH_FRACTION = DEFAULT_CONTACT_ARROW_STYLE.min_length_fraction
|
|
MAX_ARROW_LENGTH_FRACTION = DEFAULT_CONTACT_ARROW_STYLE.max_length_fraction
|
|
MIN_ARROW_THICKNESS_SCALE = DEFAULT_CONTACT_ARROW_STYLE.min_thickness_scale
|
|
MAX_ARROW_THICKNESS_SCALE = DEFAULT_CONTACT_ARROW_STYLE.max_thickness_scale
|
|
REFERENCE_NORMAL_FORCE = DEFAULT_CONTACT_ARROW_STYLE.reference_normal_force
|
|
FORCE_STRENGTH_MODE = DEFAULT_CONTACT_ARROW_STYLE.force_strength_mode
|
|
FORCE_STRENGTH_GAMMA = DEFAULT_CONTACT_ARROW_STYLE.force_strength_gamma
|
|
CONTACT_ARROW_COLOR = DEFAULT_CONTACT_ARROW_STYLE.color
|
|
|
|
|
|
def force_visual_strength(
|
|
force: float,
|
|
*,
|
|
style: ContactArrowStyle = DEFAULT_CONTACT_ARROW_STYLE,
|
|
) -> float:
|
|
"""Map normal force Fn to a unitless strength in [0, 1] for length and thickness."""
|
|
reference = style.reference_normal_force
|
|
if force <= 0.0 or reference <= 0.0:
|
|
return 0.0
|
|
|
|
ratio = force / reference
|
|
mode = style.force_strength_mode
|
|
|
|
if mode == "linear":
|
|
strength = ratio
|
|
elif mode == "power":
|
|
strength = ratio ** style.force_strength_gamma
|
|
else: # log
|
|
strength = np.log1p(ratio) / np.log1p(1.0)
|
|
|
|
return float(np.clip(strength, 0.0, 1.0))
|
|
|
|
|
|
def thickness_scale_from_force(
|
|
force: float,
|
|
*,
|
|
style: ContactArrowStyle = DEFAULT_CONTACT_ARROW_STYLE,
|
|
) -> float:
|
|
"""Radial scale from Fn using the shared force_strength curve."""
|
|
strength = force_visual_strength(force, style=style)
|
|
span = style.max_thickness_scale - style.min_thickness_scale
|
|
return float(style.min_thickness_scale + span * strength)
|
|
|
|
|
|
def length_from_force(
|
|
force: float,
|
|
mesh_span: float,
|
|
*,
|
|
style: ContactArrowStyle = DEFAULT_CONTACT_ARROW_STYLE,
|
|
) -> float:
|
|
"""Arrow length in world units from Fn and mesh extent."""
|
|
strength = force_visual_strength(force, style=style)
|
|
frac_span = style.max_length_fraction - style.min_length_fraction
|
|
fraction = style.min_length_fraction + frac_span * strength
|
|
return float(mesh_span * fraction * style.normal_force_scale)
|
|
|
|
|
|
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 _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.15,
|
|
cone_radius=0.3,
|
|
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,
|
|
style: ContactArrowStyle = DEFAULT_CONTACT_ARROW_STYLE,
|
|
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
"""Build a 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(style.color, dtype=np.float64), (len(template_vertices), 1))
|
|
|
|
radial_scale = thickness_scale_from_force(force, style=style)
|
|
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 * style.min_length_fraction,
|
|
mesh_span * style.max_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,
|
|
*,
|
|
style: ContactArrowStyle = DEFAULT_CONTACT_ARROW_STYLE,
|
|
):
|
|
"""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(style.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,
|
|
*,
|
|
style: ContactArrowStyle | None = None,
|
|
normal_force_scale: float | None = None,
|
|
enabled: bool = True,
|
|
):
|
|
import open3d as o3d
|
|
|
|
self.mesh = mesh
|
|
base_style = style or DEFAULT_CONTACT_ARROW_STYLE
|
|
if normal_force_scale is not None:
|
|
base_style = replace(base_style, normal_force_scale=float(normal_force_scale))
|
|
self.style = base_style
|
|
self.enabled = enabled
|
|
self._mesh_span = _mesh_span(mesh.rest_vertices)
|
|
_, _, 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,
|
|
style=self.style,
|
|
)
|
|
if enabled
|
|
else None
|
|
)
|
|
|
|
@property
|
|
def normal_force_scale(self) -> float:
|
|
return self.style.normal_force_scale
|
|
|
|
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
|
|
|
|
style = self.style
|
|
template_vertices = self._template_vertices
|
|
template_triangles = self._template_triangles
|
|
color = np.tile(np.asarray(style.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 = length_from_force(force, self._mesh_span, style=style)
|
|
return build_contact_normal_arrow_mesh(
|
|
origin,
|
|
normal,
|
|
length,
|
|
force,
|
|
mesh_span=self._mesh_span,
|
|
template_vertices=template_vertices,
|
|
template_triangles=template_triangles,
|
|
style=style,
|
|
)
|
|
|
|
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()
|