added mesh surface arrows
This commit is contained in:
ykcheungab
@@ -0,0 +1,431 @@
|
||||
"""Optical-flow arrow tubes on the finger UV surface mesh (TriangleMesh cylinders)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from orisys.finger.mesh import DEFAULT_SURFACE_COLOR, UVSurfaceMesh
|
||||
|
||||
from .contact_arrow import _rotation_matrix_z_to
|
||||
|
||||
# Grey mesh tone used to approximate opacity on backends without true alpha.
|
||||
_MESH_GREY_RGB = np.asarray(DEFAULT_SURFACE_COLOR, dtype=np.float64)
|
||||
|
||||
# Maps style.line_width to tube radius as a fraction of mesh extent.
|
||||
LINE_WIDTH_TO_SPAN_FRACTION = 0.0003
|
||||
# Slight tail extension into the head segment to avoid visible gaps.
|
||||
_TAIL_HEAD_OVERLAP_FRACTION = 0.03
|
||||
|
||||
SegmentRole = Literal["tail", "head"]
|
||||
|
||||
_CYLINDER_TEMPLATE_VERTICES: np.ndarray | None = None
|
||||
_CYLINDER_TEMPLATE_TRIANGLES: np.ndarray | None = None
|
||||
_CONE_TEMPLATE_VERTICES: np.ndarray | None = None
|
||||
_CONE_TEMPLATE_TRIANGLES: np.ndarray | None = None
|
||||
_CONE_APEX_AT_LOW_Z: bool | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FlowArrowStyle:
|
||||
"""Visual tuning for 3D mesh flow arrows (dashboard: edit MESH_FLOW_ARROW in layout.py).
|
||||
|
||||
``line_width`` / ``head_line_width`` scale tail and head tube radius relative to mesh extent.
|
||||
``tail_color`` / ``head_color`` are RGB in 0–255 (same convention as OpenCV BGR tuples, but RGB order).
|
||||
"""
|
||||
|
||||
enabled: bool = True
|
||||
threshold: float = 1.0
|
||||
grid_spacing: int = 20
|
||||
arrow_scale: float = 0.5
|
||||
length_margin: float = 1.0
|
||||
tail_color: tuple[int, int, int] = (38, 235, 89)
|
||||
head_color: tuple[int, int, int] = (13, 115, 31)
|
||||
head_fraction: float = 0.25
|
||||
opacity: float = 0.65
|
||||
line_width: float = 2.0
|
||||
head_line_width: float = 3.0
|
||||
edge_margin: int = 8
|
||||
|
||||
|
||||
DEFAULT_FLOW_ARROW_STYLE = FlowArrowStyle()
|
||||
|
||||
|
||||
def _mesh_span(vertices: np.ndarray) -> float:
|
||||
extent = np.ptp(vertices, axis=0)
|
||||
return float(max(np.max(extent), 1e-6))
|
||||
|
||||
|
||||
def segment_radius(mesh_span: float, line_width: float) -> float:
|
||||
"""World-space tube radius from mesh extent and a width knob."""
|
||||
return mesh_span * max(float(line_width), 0.1) * LINE_WIDTH_TO_SPAN_FRACTION
|
||||
|
||||
|
||||
def segment_radius_for_role(mesh_span: float, style: FlowArrowStyle, role: SegmentRole) -> float:
|
||||
width = style.line_width if role == "tail" else style.head_line_width
|
||||
return segment_radius(mesh_span, width)
|
||||
|
||||
|
||||
def _prepare_uv_mask(valid_mask, shape: tuple[int, int]) -> np.ndarray | None:
|
||||
if valid_mask is None:
|
||||
return None
|
||||
h, w = shape
|
||||
mask = np.asarray(valid_mask, dtype=bool)
|
||||
if mask.shape != (h, w):
|
||||
mask_u8 = mask.astype(np.uint8)
|
||||
if mask_u8.ndim > 2:
|
||||
mask_u8 = mask_u8[:, :, 0]
|
||||
mask = cv2.resize(mask_u8, (w, h), interpolation=cv2.INTER_NEAREST).astype(bool)
|
||||
return mask
|
||||
|
||||
|
||||
def _style_rgb01(style: FlowArrowStyle, role: SegmentRole) -> np.ndarray:
|
||||
rgb = np.asarray(style.tail_color if role == "tail" else style.head_color, dtype=np.float64).reshape(3)
|
||||
return np.clip(rgb / 255.0, 0.0, 1.0)
|
||||
|
||||
|
||||
def effective_segment_color(
|
||||
style: FlowArrowStyle,
|
||||
role: SegmentRole,
|
||||
*,
|
||||
use_alpha_blend: bool = True,
|
||||
) -> np.ndarray:
|
||||
"""Return RGB in 0–1 for a tail or head segment."""
|
||||
rgb = _style_rgb01(style, role)
|
||||
if not use_alpha_blend:
|
||||
return rgb
|
||||
opacity = float(np.clip(style.opacity, 0.0, 1.0))
|
||||
return np.clip(opacity * rgb + (1.0 - opacity) * _MESH_GREY_RGB, 0.0, 1.0)
|
||||
|
||||
|
||||
def _sample_flow_segments(
|
||||
world_xyz: np.ndarray,
|
||||
flow: np.ndarray,
|
||||
*,
|
||||
style: FlowArrowStyle = DEFAULT_FLOW_ARROW_STYLE,
|
||||
valid_mask: np.ndarray | None = None,
|
||||
) -> tuple[np.ndarray, np.ndarray, list[SegmentRole]]:
|
||||
"""Return segment start/end world points and role per segment (tail then head)."""
|
||||
empty = np.empty((0, 3), dtype=np.float64)
|
||||
if flow is None or np.asarray(flow).size == 0:
|
||||
return empty, empty, []
|
||||
|
||||
flow = np.asarray(flow, dtype=np.float64)
|
||||
world_xyz = np.asarray(world_xyz, dtype=np.float64)
|
||||
h, w = flow.shape[:2]
|
||||
if world_xyz.shape[:2] != (h, w):
|
||||
return empty, empty, []
|
||||
|
||||
grid_spacing = max(int(style.grid_spacing), 1)
|
||||
edge_margin = int(style.edge_margin)
|
||||
margin = float(style.length_margin)
|
||||
arrow_scale = max(float(style.arrow_scale), 1e-9)
|
||||
head_fraction = float(np.clip(style.head_fraction, 0.05, 0.95))
|
||||
tail_fraction = 1.0 - head_fraction
|
||||
|
||||
rr = np.arange(grid_spacing // 2, h, grid_spacing)
|
||||
cc = np.arange(grid_spacing // 2, w, grid_spacing)
|
||||
rr, cc = np.meshgrid(rr, cc, indexing="ij")
|
||||
rr, cc = rr.reshape(-1), cc.reshape(-1)
|
||||
|
||||
disp = 2.0 * flow[rr, cc, :]
|
||||
mag = np.linalg.norm(disp, axis=1)
|
||||
dx = disp[:, 0] / arrow_scale
|
||||
dy = disp[:, 1] / arrow_scale
|
||||
length = np.sqrt(dx * dx + dy * dy)
|
||||
scale = np.where(
|
||||
length > grid_spacing * margin,
|
||||
grid_spacing / np.maximum(length, 1e-9),
|
||||
1.0,
|
||||
) * margin
|
||||
dx, dy = dx * scale, dy * scale
|
||||
end_x = np.clip(np.round(cc + dx).astype(np.int32), 0, w - 1)
|
||||
end_y = np.clip(np.round(rr + dy).astype(np.int32), 0, h - 1)
|
||||
|
||||
valid = (
|
||||
(cc >= edge_margin)
|
||||
& (cc < w - edge_margin)
|
||||
& (rr >= edge_margin)
|
||||
& (rr < h - edge_margin)
|
||||
& (mag >= float(style.threshold))
|
||||
)
|
||||
|
||||
mask = _prepare_uv_mask(valid_mask, (h, w))
|
||||
if mask is not None:
|
||||
valid &= mask[rr, cc]
|
||||
valid &= mask[end_y, end_x]
|
||||
|
||||
if not np.any(valid):
|
||||
return empty, empty, []
|
||||
|
||||
starts = world_xyz[rr[valid], cc[valid]]
|
||||
ends = world_xyz[end_y[valid], end_x[valid]]
|
||||
finite = np.all(np.isfinite(starts), axis=1) & np.all(np.isfinite(ends), axis=1)
|
||||
if not np.any(finite):
|
||||
return empty, empty, []
|
||||
|
||||
starts = starts[finite]
|
||||
ends = ends[finite]
|
||||
non_zero = np.linalg.norm(ends - starts, axis=1) > 1e-9
|
||||
if not np.any(non_zero):
|
||||
return empty, empty, []
|
||||
|
||||
starts = starts[non_zero]
|
||||
ends = ends[non_zero]
|
||||
junctions = starts + tail_fraction * (ends - starts)
|
||||
|
||||
seg_starts = np.vstack((starts, junctions))
|
||||
seg_ends = np.vstack((junctions, ends))
|
||||
n = len(starts)
|
||||
roles: list[SegmentRole] = ["tail"] * n + ["head"] * n
|
||||
return seg_starts, seg_ends, roles
|
||||
|
||||
|
||||
def build_flow_arrow_lines(
|
||||
world_xyz: np.ndarray,
|
||||
flow: np.ndarray,
|
||||
*,
|
||||
style: FlowArrowStyle = DEFAULT_FLOW_ARROW_STYLE,
|
||||
valid_mask: np.ndarray | None = None,
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
"""Build line indices for flow arrows (used by tests and debugging)."""
|
||||
empty_pts = np.empty((0, 3), dtype=np.float64)
|
||||
empty_lines = np.empty((0, 2), dtype=np.int32)
|
||||
empty_colors = np.empty((0, 3), dtype=np.float64)
|
||||
|
||||
seg_starts, seg_ends, roles = _sample_flow_segments(
|
||||
world_xyz, flow, style=style, valid_mask=valid_mask
|
||||
)
|
||||
if len(seg_starts) == 0:
|
||||
return empty_pts, empty_lines, empty_colors
|
||||
|
||||
points = np.empty((len(seg_starts) * 2, 3), dtype=np.float64)
|
||||
points[0::2] = seg_starts
|
||||
points[1::2] = seg_ends
|
||||
lines = np.column_stack(
|
||||
(np.arange(0, len(points), 2, dtype=np.int32), np.arange(1, len(points), 2, dtype=np.int32))
|
||||
)
|
||||
line_colors = np.vstack(
|
||||
[effective_segment_color(style, role, use_alpha_blend=True) for role in roles]
|
||||
)
|
||||
return points, lines, line_colors
|
||||
|
||||
|
||||
def _cylinder_template(o3d):
|
||||
global _CYLINDER_TEMPLATE_VERTICES, _CYLINDER_TEMPLATE_TRIANGLES
|
||||
if _CYLINDER_TEMPLATE_VERTICES is None:
|
||||
cyl = o3d.geometry.TriangleMesh.create_cylinder(
|
||||
radius=1.0,
|
||||
height=1.0,
|
||||
resolution=8,
|
||||
split=1,
|
||||
)
|
||||
_CYLINDER_TEMPLATE_VERTICES = np.asarray(cyl.vertices, dtype=np.float64)
|
||||
_CYLINDER_TEMPLATE_TRIANGLES = np.asarray(cyl.triangles, dtype=np.int32)
|
||||
return _CYLINDER_TEMPLATE_VERTICES, _CYLINDER_TEMPLATE_TRIANGLES
|
||||
|
||||
|
||||
def _cone_apex_at_low_z(vertices: np.ndarray) -> bool:
|
||||
"""True when the cone tip is on the low-z side of the Open3D template."""
|
||||
radii = np.linalg.norm(vertices[:, :2], axis=1)
|
||||
apex_z = float(vertices[np.argmin(radii), 2])
|
||||
z_mid = 0.5 * (float(vertices[:, 2].min()) + float(vertices[:, 2].max()))
|
||||
return apex_z < z_mid
|
||||
|
||||
|
||||
def _cone_template(o3d):
|
||||
global _CONE_TEMPLATE_VERTICES, _CONE_TEMPLATE_TRIANGLES, _CONE_APEX_AT_LOW_Z
|
||||
if _CONE_TEMPLATE_VERTICES is None:
|
||||
cone = o3d.geometry.TriangleMesh.create_cone(
|
||||
radius=1.0,
|
||||
height=1.0,
|
||||
resolution=8,
|
||||
split=1,
|
||||
)
|
||||
_CONE_TEMPLATE_VERTICES = np.asarray(cone.vertices, dtype=np.float64)
|
||||
_CONE_TEMPLATE_TRIANGLES = np.asarray(cone.triangles, dtype=np.int32)
|
||||
_CONE_APEX_AT_LOW_Z = _cone_apex_at_low_z(_CONE_TEMPLATE_VERTICES)
|
||||
return _CONE_TEMPLATE_VERTICES, _CONE_TEMPLATE_TRIANGLES
|
||||
|
||||
|
||||
def _transform_axis_segment(
|
||||
template_vertices: np.ndarray,
|
||||
start: np.ndarray,
|
||||
end: np.ndarray,
|
||||
radius: float,
|
||||
*,
|
||||
apex_at_end: bool = False,
|
||||
) -> np.ndarray:
|
||||
"""Place a +Z template from ``start`` to ``end`` (axis anchored, not midpoint-centered).
|
||||
|
||||
Open3D cylinders are centered on the origin; cone apex side is detected from the template.
|
||||
``apex_at_end=True`` maps the template apex to ``end`` (flow tip); base to ``start``.
|
||||
"""
|
||||
direction = np.asarray(end, dtype=np.float64) - np.asarray(start, dtype=np.float64)
|
||||
length = float(np.linalg.norm(direction))
|
||||
if length <= 1e-9 or radius <= 0.0:
|
||||
return np.empty((0, 3), dtype=np.float64)
|
||||
|
||||
dir_hat = direction / length
|
||||
rotation = _rotation_matrix_z_to(dir_hat)
|
||||
scaled = np.asarray(template_vertices, dtype=np.float64).copy()
|
||||
|
||||
z_min = float(scaled[:, 2].min())
|
||||
z_max = float(scaled[:, 2].max())
|
||||
z_span = z_max - z_min
|
||||
if z_span <= 1e-9:
|
||||
return np.empty((0, 3), dtype=np.float64)
|
||||
|
||||
if apex_at_end:
|
||||
t = (z_max - scaled[:, 2]) / z_span * length
|
||||
else:
|
||||
t = (scaled[:, 2] - z_min) / z_span * length
|
||||
|
||||
radial = np.column_stack((scaled[:, 0] * radius, scaled[:, 1] * radius, t))
|
||||
return radial @ rotation.T + np.asarray(start, dtype=np.float64)
|
||||
|
||||
|
||||
def _transform_cylinder_segment(
|
||||
template_vertices: np.ndarray,
|
||||
start: np.ndarray,
|
||||
end: np.ndarray,
|
||||
radius: float,
|
||||
) -> np.ndarray:
|
||||
return _transform_axis_segment(template_vertices, start, end, radius)
|
||||
|
||||
|
||||
def build_flow_arrow_mesh(
|
||||
world_xyz: np.ndarray,
|
||||
flow: np.ndarray,
|
||||
mesh_span: float,
|
||||
*,
|
||||
style: FlowArrowStyle = DEFAULT_FLOW_ARROW_STYLE,
|
||||
valid_mask: np.ndarray | None = None,
|
||||
o3d=None,
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
"""Build merged mesh flow arrows: cylinder tail + cone head per arrow."""
|
||||
empty_v = np.empty((0, 3), dtype=np.float64)
|
||||
empty_t = np.empty((0, 3), dtype=np.int32)
|
||||
empty_c = np.empty((0, 3), dtype=np.float64)
|
||||
|
||||
seg_starts, seg_ends, roles = _sample_flow_segments(
|
||||
world_xyz, flow, style=style, valid_mask=valid_mask
|
||||
)
|
||||
if len(seg_starts) == 0:
|
||||
return empty_v, empty_t, empty_c
|
||||
|
||||
if o3d is None:
|
||||
import open3d as o3d
|
||||
|
||||
cyl_vertices, cyl_triangles = _cylinder_template(o3d)
|
||||
cone_vertices, cone_triangles = _cone_template(o3d)
|
||||
|
||||
vertices_parts: list[np.ndarray] = []
|
||||
triangles_parts: list[np.ndarray] = []
|
||||
colors_parts: list[np.ndarray] = []
|
||||
vertex_offset = 0
|
||||
|
||||
for start, end, role in zip(seg_starts, seg_ends, roles):
|
||||
seg_start = np.asarray(start, dtype=np.float64)
|
||||
seg_end = np.asarray(end, dtype=np.float64)
|
||||
axis = seg_end - seg_start
|
||||
seg_len = float(np.linalg.norm(axis))
|
||||
if seg_len <= 1e-9:
|
||||
continue
|
||||
|
||||
radius = segment_radius_for_role(mesh_span, style, role)
|
||||
if role == "tail":
|
||||
template_vertices, template_triangles = cyl_vertices, cyl_triangles
|
||||
apex_at_end = False
|
||||
overlap = _TAIL_HEAD_OVERLAP_FRACTION * seg_len
|
||||
seg_end = seg_end + (axis / seg_len) * overlap
|
||||
else:
|
||||
template_vertices, template_triangles = cone_vertices, cone_triangles
|
||||
apex_at_end = not bool(_CONE_APEX_AT_LOW_Z)
|
||||
tail_radius = segment_radius_for_role(mesh_span, style, "tail")
|
||||
radius = max(radius, tail_radius)
|
||||
|
||||
verts = _transform_axis_segment(
|
||||
template_vertices,
|
||||
seg_start,
|
||||
seg_end,
|
||||
radius,
|
||||
apex_at_end=apex_at_end,
|
||||
)
|
||||
if verts.size == 0:
|
||||
continue
|
||||
color = effective_segment_color(style, role, use_alpha_blend=True)
|
||||
vertices_parts.append(verts)
|
||||
triangles_parts.append(template_triangles + vertex_offset)
|
||||
colors_parts.append(np.tile(color, (len(verts), 1)))
|
||||
vertex_offset += len(verts)
|
||||
|
||||
if not vertices_parts:
|
||||
return empty_v, empty_t, empty_c
|
||||
|
||||
return (
|
||||
np.vstack(vertices_parts),
|
||||
np.vstack(triangles_parts),
|
||||
np.vstack(colors_parts),
|
||||
)
|
||||
|
||||
|
||||
def create_flow_mesh_geometry(o3d):
|
||||
"""Create a persistent Open3D mesh for the flow arrow overlay."""
|
||||
mesh = o3d.geometry.TriangleMesh()
|
||||
mesh.vertices = o3d.utility.Vector3dVector(np.empty((0, 3), dtype=np.float64))
|
||||
mesh.triangles = o3d.utility.Vector3iVector(np.empty((0, 3), dtype=np.int32))
|
||||
mesh.vertex_colors = o3d.utility.Vector3dVector(np.empty((0, 3), dtype=np.float64))
|
||||
mesh.compute_vertex_normals()
|
||||
return mesh
|
||||
|
||||
|
||||
class FlowArrowOverlay:
|
||||
"""Example-layer flow arrows as cylinder tail + cone head meshes on the UV surface."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mesh: UVSurfaceMesh,
|
||||
*,
|
||||
style: FlowArrowStyle | None = None,
|
||||
):
|
||||
import open3d as o3d
|
||||
|
||||
self.mesh = mesh
|
||||
self.style = style or DEFAULT_FLOW_ARROW_STYLE
|
||||
self._mesh_span = _mesh_span(mesh.rest_vertices)
|
||||
_cylinder_template(o3d)
|
||||
_cone_template(o3d)
|
||||
self.arrow_mesh = create_flow_mesh_geometry(o3d) if self.style.enabled else None
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return self.style.enabled
|
||||
|
||||
def compute(
|
||||
self,
|
||||
flow,
|
||||
valid_mask: np.ndarray | None = None,
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray] | None:
|
||||
if not self.style.enabled or self.arrow_mesh is None:
|
||||
return None
|
||||
|
||||
return build_flow_arrow_mesh(
|
||||
self.mesh.world_xyz,
|
||||
flow,
|
||||
self._mesh_span,
|
||||
style=self.style,
|
||||
valid_mask=valid_mask,
|
||||
)
|
||||
|
||||
def apply(self, arrow_vertices: np.ndarray, arrow_triangles: np.ndarray, arrow_colors: np.ndarray) -> None:
|
||||
"""Update the persistent flow mesh geometry in place."""
|
||||
import open3d as o3d
|
||||
|
||||
self.arrow_mesh.vertices = o3d.utility.Vector3dVector(arrow_vertices.astype(np.float64))
|
||||
self.arrow_mesh.triangles = o3d.utility.Vector3iVector(arrow_triangles.astype(np.int32))
|
||||
self.arrow_mesh.vertex_colors = o3d.utility.Vector3dVector(arrow_colors.astype(np.float64))
|
||||
self.arrow_mesh.compute_vertex_normals()
|
||||
Reference in New Issue
Block a user