split example visualizaiton pipeline
This commit is contained in:
ykcheungab
@@ -0,0 +1,271 @@
|
||||
"""UV-domain finger surface mesh built from world_xyz and tactile field coloring."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
|
||||
# Default grey-white surface tone (linear RGB, 0-1).
|
||||
DEFAULT_SURFACE_COLOR = (0.88, 0.88, 0.90)
|
||||
|
||||
|
||||
class UVSurfaceMesh:
|
||||
"""Triangle mesh built from a UV flatten grid and world_xyz surface map."""
|
||||
|
||||
__slots__ = (
|
||||
"rest_vertices",
|
||||
"triangles",
|
||||
"vertex_ys",
|
||||
"vertex_xs",
|
||||
"index_map",
|
||||
"valid_mask",
|
||||
"shape",
|
||||
"world_xyz",
|
||||
"base_colors",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rest_vertices,
|
||||
triangles,
|
||||
vertex_ys,
|
||||
vertex_xs,
|
||||
index_map,
|
||||
valid_mask,
|
||||
world_xyz,
|
||||
base_colors,
|
||||
):
|
||||
self.rest_vertices = rest_vertices
|
||||
self.triangles = triangles
|
||||
self.vertex_ys = vertex_ys
|
||||
self.vertex_xs = vertex_xs
|
||||
self.index_map = index_map
|
||||
self.valid_mask = valid_mask
|
||||
self.shape = valid_mask.shape
|
||||
self.world_xyz = world_xyz
|
||||
self.base_colors = base_colors
|
||||
|
||||
|
||||
def build_uv_surface_mesh(
|
||||
world_xyz: np.ndarray,
|
||||
valid_mask: np.ndarray,
|
||||
base_color=DEFAULT_SURFACE_COLOR,
|
||||
) -> UVSurfaceMesh:
|
||||
"""Build a triangle mesh from UV grid world_xyz and valid_mask.
|
||||
|
||||
Each valid UV pixel becomes one vertex. Adjacent valid 2x2 cells form two
|
||||
triangles. Invalid or non-finite vertices are excluded.
|
||||
"""
|
||||
world_xyz = np.asarray(world_xyz, dtype=np.float64)
|
||||
valid_mask = np.asarray(valid_mask, dtype=bool)
|
||||
if world_xyz.shape[:2] != valid_mask.shape:
|
||||
raise ValueError("world_xyz and valid_mask must share the same HxW shape")
|
||||
|
||||
h, w = valid_mask.shape
|
||||
finite = np.all(np.isfinite(world_xyz), axis=2)
|
||||
vertex_mask = valid_mask & finite
|
||||
|
||||
index_map = np.full((h, w), -1, dtype=np.int32)
|
||||
ys, xs = np.nonzero(vertex_mask)
|
||||
num_vertices = len(ys)
|
||||
index_map[ys, xs] = np.arange(num_vertices, dtype=np.int32)
|
||||
|
||||
rest_vertices = world_xyz[ys, xs].astype(np.float64)
|
||||
color = np.asarray(base_color, dtype=np.float64).reshape(1, 3)
|
||||
base_colors = np.tile(color, (num_vertices, 1))
|
||||
|
||||
triangles = []
|
||||
for y in range(h - 1):
|
||||
for x in range(w - 1):
|
||||
i00 = index_map[y, x]
|
||||
i10 = index_map[y, x + 1]
|
||||
i01 = index_map[y + 1, x]
|
||||
i11 = index_map[y + 1, x + 1]
|
||||
if min(i00, i10, i01, i11) < 0:
|
||||
continue
|
||||
triangles.append((i00, i10, i01))
|
||||
triangles.append((i10, i11, i01))
|
||||
|
||||
if not triangles:
|
||||
raise ValueError("UV surface mesh has no valid triangles.")
|
||||
|
||||
return UVSurfaceMesh(
|
||||
rest_vertices=rest_vertices,
|
||||
triangles=np.asarray(triangles, dtype=np.int32),
|
||||
vertex_ys=ys.astype(np.int32),
|
||||
vertex_xs=xs.astype(np.int32),
|
||||
index_map=index_map,
|
||||
valid_mask=vertex_mask,
|
||||
world_xyz=world_xyz,
|
||||
base_colors=base_colors,
|
||||
)
|
||||
|
||||
|
||||
def compute_uv_surface_frames(world_xyz: np.ndarray, valid_mask: np.ndarray):
|
||||
"""Estimate tangent frame and normals on the UV surface grid."""
|
||||
valid_mask = np.asarray(valid_mask, dtype=bool)
|
||||
filled = np.asarray(world_xyz, dtype=np.float64).copy()
|
||||
invalid = ~valid_mask | ~np.all(np.isfinite(filled), axis=2)
|
||||
filled[invalid] = 0.0
|
||||
|
||||
d_y = np.gradient(filled, axis=0)
|
||||
d_x = np.gradient(filled, axis=1)
|
||||
normals = np.cross(d_x, d_y)
|
||||
norm = np.linalg.norm(normals, axis=2, keepdims=True)
|
||||
normals = np.divide(normals, norm, out=np.zeros_like(normals), where=norm > 1e-8)
|
||||
d_x[invalid] = 0.0
|
||||
d_y[invalid] = 0.0
|
||||
normals[invalid] = 0.0
|
||||
return d_x, d_y, normals
|
||||
|
||||
|
||||
def mesh_uv_to_world_scale(world_xyz: np.ndarray, valid_mask: np.ndarray) -> float:
|
||||
"""Approximate world-units per UV pixel for deformation scaling."""
|
||||
points = world_xyz[valid_mask]
|
||||
if len(points) == 0:
|
||||
return 1e-3
|
||||
extent = np.nanmax(points, axis=0) - np.nanmin(points, axis=0)
|
||||
mesh_span = float(np.max(extent))
|
||||
h, w = world_xyz.shape[:2]
|
||||
return mesh_span / max(h, w, 1)
|
||||
|
||||
|
||||
def lift_uv_deformation_to_3d(
|
||||
vshear: np.ndarray,
|
||||
vnormal: np.ndarray,
|
||||
tangent_x: np.ndarray,
|
||||
tangent_y: np.ndarray,
|
||||
normals: np.ndarray,
|
||||
normal_gain: float = 1.0,
|
||||
) -> np.ndarray:
|
||||
"""Lift UV-space shear/normal HHD fields to 3D tangent-plane displacement."""
|
||||
tx_norm = np.linalg.norm(tangent_x, axis=2, keepdims=True)
|
||||
ty_norm = np.linalg.norm(tangent_y, axis=2, keepdims=True)
|
||||
t1 = np.divide(tangent_x, tx_norm, out=np.zeros_like(tangent_x), where=tx_norm > 1e-8)
|
||||
t2 = np.divide(tangent_y, ty_norm, out=np.zeros_like(tangent_y), where=ty_norm > 1e-8)
|
||||
|
||||
shear_3d = vshear[:, :, 0:1] * t1 + vshear[:, :, 1:2] * t2
|
||||
normal_mag = np.linalg.norm(vnormal, axis=2, keepdims=True) * float(normal_gain)
|
||||
normal_3d = normal_mag * normals
|
||||
|
||||
normal_component = np.sum(shear_3d * normals, axis=2, keepdims=True)
|
||||
shear_tangent = shear_3d - normal_component * normals
|
||||
return shear_tangent + normal_3d
|
||||
|
||||
|
||||
def lift_uv_flow_to_3d(
|
||||
flow: np.ndarray,
|
||||
tangent_x: np.ndarray,
|
||||
tangent_y: np.ndarray,
|
||||
normals: np.ndarray,
|
||||
) -> np.ndarray:
|
||||
"""Lift raw 2D optical flow on the UV grid to tangent-plane 3D displacement."""
|
||||
tx_norm = np.linalg.norm(tangent_x, axis=2, keepdims=True)
|
||||
ty_norm = np.linalg.norm(tangent_y, axis=2, keepdims=True)
|
||||
t1 = np.divide(tangent_x, tx_norm, out=np.zeros_like(tangent_x), where=tx_norm > 1e-8)
|
||||
t2 = np.divide(tangent_y, ty_norm, out=np.zeros_like(tangent_y), where=ty_norm > 1e-8)
|
||||
flow_3d = flow[:, :, 0:1] * t1 + flow[:, :, 1:2] * t2
|
||||
normal_component = np.sum(flow_3d * normals, axis=2, keepdims=True)
|
||||
return flow_3d - normal_component * normals
|
||||
|
||||
|
||||
def sample_field_at_mesh_vertices(field: np.ndarray, mesh: UVSurfaceMesh) -> np.ndarray:
|
||||
"""Sample a per-pixel UV field at mesh vertex locations."""
|
||||
sampled = field[mesh.vertex_ys, mesh.vertex_xs]
|
||||
if sampled.ndim == 1:
|
||||
return sampled.astype(np.float64)
|
||||
return sampled.astype(np.float64)
|
||||
|
||||
|
||||
def flow_magnitude_vertex_colors(
|
||||
flow: np.ndarray | None,
|
||||
mesh: UVSurfaceMesh,
|
||||
*,
|
||||
base_color=DEFAULT_SURFACE_COLOR,
|
||||
max_magnitude: float | None = None,
|
||||
threshold: float = 0.0,
|
||||
) -> np.ndarray:
|
||||
"""Color mesh vertices from UV-space optical-flow magnitude.
|
||||
|
||||
Magnitude matches ``orisys.util.draw_arrows``: ``2 * ||flow||``.
|
||||
Vertices below ``threshold`` keep the base grey-white surface color.
|
||||
"""
|
||||
colors = mesh.base_colors.copy()
|
||||
if flow is None or flow.size == 0:
|
||||
return colors
|
||||
|
||||
flow = np.asarray(flow, dtype=np.float64)
|
||||
mag = 2.0 * np.linalg.norm(flow, axis=2)
|
||||
vertex_mag = mag[mesh.vertex_ys, mesh.vertex_xs]
|
||||
finite = np.isfinite(vertex_mag)
|
||||
active = finite & (vertex_mag >= float(threshold))
|
||||
if not np.any(active):
|
||||
return colors
|
||||
|
||||
active_mag = vertex_mag[active]
|
||||
if max_magnitude is None or max_magnitude <= 0:
|
||||
max_magnitude = max(float(np.percentile(active_mag, 95)), 1e-6)
|
||||
|
||||
value = np.clip(vertex_mag / max_magnitude, 0.0, 1.0)
|
||||
heat = np.column_stack(
|
||||
(
|
||||
value,
|
||||
np.clip(1.0 - np.abs(value - 0.5) * 2.0, 0.0, 1.0),
|
||||
1.0 - value,
|
||||
)
|
||||
).astype(np.float64)
|
||||
base = np.asarray(base_color, dtype=np.float64)
|
||||
colors[active] = 0.25 * base + 0.75 * heat[active]
|
||||
return colors
|
||||
|
||||
|
||||
def scalar_map_vertex_colors(
|
||||
scalar_map: np.ndarray | None,
|
||||
mesh: UVSurfaceMesh,
|
||||
*,
|
||||
base_color=DEFAULT_SURFACE_COLOR,
|
||||
value_min: float | None = None,
|
||||
value_max: float | None = None,
|
||||
threshold: float = 0.0,
|
||||
) -> np.ndarray:
|
||||
"""Color mesh vertices from a scalar UV map such as depth/divergence."""
|
||||
colors = mesh.base_colors.copy()
|
||||
if scalar_map is None or scalar_map.size == 0:
|
||||
return colors
|
||||
|
||||
scalar_map = np.asarray(scalar_map, dtype=np.float64)
|
||||
vertex_values = scalar_map[mesh.vertex_ys, mesh.vertex_xs]
|
||||
finite = np.isfinite(vertex_values)
|
||||
active = finite & (vertex_values > float(threshold))
|
||||
if not np.any(active):
|
||||
return colors
|
||||
|
||||
active_values = vertex_values[active]
|
||||
if value_min is None:
|
||||
value_min = float(np.percentile(active_values, 5))
|
||||
if value_max is None:
|
||||
value_max = float(np.percentile(active_values, 95))
|
||||
if value_max <= value_min:
|
||||
value_max = value_min + 1e-6
|
||||
|
||||
value = np.clip((vertex_values - value_min) / (value_max - value_min), 0.0, 1.0)
|
||||
heat = np.column_stack(
|
||||
(
|
||||
value,
|
||||
np.clip(1.0 - np.abs(value - 0.5) * 2.0, 0.0, 1.0),
|
||||
1.0 - value,
|
||||
)
|
||||
).astype(np.float64)
|
||||
base = np.asarray(base_color, dtype=np.float64)
|
||||
colors[active] = 0.25 * base + 0.75 * heat[active]
|
||||
return colors
|
||||
|
||||
|
||||
def deform_mesh_vertices(
|
||||
mesh: UVSurfaceMesh,
|
||||
displacement_map: np.ndarray,
|
||||
pixel_to_world: float,
|
||||
deform_scale: float = 1.0,
|
||||
) -> np.ndarray:
|
||||
"""Apply lifted 3D displacement to rest mesh vertices."""
|
||||
disp = sample_field_at_mesh_vertices(displacement_map, mesh)
|
||||
return mesh.rest_vertices + disp * float(pixel_to_world) * float(deform_scale)
|
||||
Reference in New Issue
Block a user