diff --git a/.github/actions/setup-uv/action.yml b/.github/actions/setup-uv/action.yml index ac0cbb95b7..6187201c87 100644 --- a/.github/actions/setup-uv/action.yml +++ b/.github/actions/setup-uv/action.yml @@ -12,4 +12,5 @@ runs: uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 with: # Pinning explicitly avoids setup-uv version-resolution API calls and fallback paths. - version: "0.10.6" + # Required for package-scoped dependency exclusions in pyproject.toml. + version: "0.12.13" diff --git a/pyproject.toml b/pyproject.toml index 49fc75eda3..85b4719325 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,13 @@ dev = [ ] ui = [ "imgui[glfw]>=2.0.0", + "numpy>=1.23.0", + # Dockable panels and Trimesh 5 compatibility require Viser >=1.1.1. + # https://github.com/viser-project/viser/pull/769 + "viser>=1.1.1", + "yourdfpy", + "trimesh", + "rtree", ] retargeters = [ "dex-retargeting>=0.4.6,<0.6.0", @@ -109,6 +116,8 @@ result = "{major}.{minor}+local" # [build-system], not [tool.uv]). [tool.uv] package = false +# Ignore Viser's websockets requirement while preserving CloudXR's requirement. +exclude-dependencies = [{ package = { name = "viser" }, dependencies = ["websockets"] }] [tool.uv.sources] isaacteleop = { path = "src/compat" } diff --git a/scripts/tactile_metadata/README.md b/scripts/tactile_metadata/README.md new file mode 100644 index 0000000000..9811027c44 --- /dev/null +++ b/scripts/tactile_metadata/README.md @@ -0,0 +1,254 @@ + + +# Tactile Metadata Editor + +Tactile Metadata Editor is a local Viser application for creating, editing, and +previewing glove tactile metadata. It loads URDF models and +[`TactileMetadata`](../../src/core/schema/fbs/tactile_metadata.fbs) JSON files. +All edits remain local; live connections only receive Isaac Capture data. + +> [!WARNING] +> The editor is designed for completing all work in one run. Saved metadata +> contains only schema fields, so ending the session loses editor-specific +> configuration such as URDF paths, alignments, and URDF joint values. + +**Todo**: add cross-run project saving-loading. + +## Set up and launch + +Use an existing Isaac Capture environment with Python 3.11 or newer. Install +the shared UI requirements into that environment: + +```bash +cd scripts/tactile_metadata +python -m pip install -r ../../src/core/python/requirements-ui.txt +python -m tactile_metadata_editor.cli +``` + +Open [http://127.0.0.1:8080](http://127.0.0.1:8080). To use another port: + +```bash +python -m tactile_metadata_editor.cli --port 8090 +``` + +Metadata, side, and URDF inputs are optional: + +```bash +python -m tactile_metadata_editor.cli \ + --metadata \ + --side \ + --urdf +``` + +`--side` defaults to `left` and initializes **Hand Side**. Paths beginning with +`~` are expanded. + +Offline editing uses Viser, NumPy, yourdfpy, trimesh, and rtree. Live reception +also requires the Isaac Capture SDK, a compatible OpenXR runtime, and an +available HMD system in the same Python environment. + +## Data and coordinates + +The editor maintains one metadata document and file path. **Hand Side** selects +the default poses used by **New** and the hand stream used by **Connect Live**. +Changing it disconnects Live and clears received metadata without switching, +mirroring, or replacing the current document or its Saved and Preview datasets. +General uses a separate right-docked menu above the Editor/Preview menu. The +Selected Object panel is always present as a floating window on the left. +The **Mode** section uses bold, mutually exclusive **Edit** and **Preview** +checkbox labels; one mode is always selected. +Focused numeric fields preserve partial decimal input such as `1.` until the +next digit is entered or the field loses focus. + +Three pose datasets are displayed: + +- **Editing** contains editable segment and taxel binding poses plus Edit-mode + URDF joint values. +- **Saved** is the last loaded or saved metadata snapshot. Enable **Show Saved + Binding Poses** to render it at opacity 0.3. +- **Preview** contains current segment poses and Preview-mode URDF joint values. + Preview taxel poses are derived from the selected binding metadata. + +All UI values and scene interactions use visualization-space coordinates. +**Segments Alignment** maps metadata wrist coordinates into visualization space. +Loading applies that transform; saving applies its inverse. Changing the +alignment transforms the visible editor values without changing the underlying +metadata coordinates. + +A three-axis indicator in the top-right shows visualization space relative to +the camera. + +## Edit metadata + +The **Edit** mode contains three tabs. + +### URDF + +1. Enter an absolute **URDF Path** and select **Load URDF**. +2. Use **URDF Alignment** to translate, rotate, and scale the model. +3. Expand **URDF Joints** to adjust each movable joint. + +URDF joints are direct sliders. Segment and Taxel lists start collapsed, and +their displayed indices start at 0. + +URDF configuration is session-only and is not saved in tactile metadata. The +model renders at opacity 0.2. + +### Segments + +- Use **Segments Alignment** to align all metadata poses. +- Add a segment from a HandJoint and XYZ position, or select **Add to URDF + Surface** and click the model. +- Segments are flat-shaded cubes with side length 0.004 m. Smooth color + gradients and cast/received shadows are disabled. +- Expand **All Segments** to edit HandJoint, position, and XYZ rotation. + +HandJoint values must be unique. Taxel bindings use HandJoint names internally, +so removing or renaming a referenced segment leaves the affected taxel field +invalid and outlined in red. + +### Taxels + +Set Rows, Columns, Area Width, Area Height, Center Position, and Rotation. +**Capture URDF Surface Pose** copies a clicked surface position and orientation +into those fields without creating taxels. **Capture Segment Pose** copies the +complete pose of a clicked Editing segment. **Add Taxels** then creates the +row-major planar grid and copies the four Segment/Weight fields above the +capture buttons into every new Taxel. Taxels are squares with side length +0.002 m. A taxel with any validation error uses `ERROR_COLOR`; +`SELECTED_COLOR` overrides it while that taxel is selected. + +To bind taxels: + +1. Select **Start Binding to Segments** and select taxels in the scene. +2. Select **Select Segments** and choose one to four segments. +3. Select **Complete**. The first Segment receives weight one; later Segments + receive weight zero. + +At most one binding-step button is enabled at a time. + +Expand **All Taxels** to edit 25 rows at a time. Page status shows the visible +global Taxel range. Use **Jump to page**, **Last Page** (previous page), or +**Next Page** from the controls above or below the rows. The unlabeled input and +Jump button share one line; Last/Next share the following line. Each visible row +contains pose, four segment slots, and four weights. `None` is an editor-only +inactive slot with weight zero. Selecting the first active segment sets its +weight to one; later segment selections start at zero. Editing a weight changes +only that slot. Page changes reuse the same 25 row controls instead of rebuilding +the menu. Validation reports rows whose weights do not sum to one. + +## Scene selection + +- Click an object to select it. +- Shift-click adds objects of the same type. +- Shift-drag selects a screen rectangle and temporarily disables camera dragging. +- A selection cannot contain both segments and taxels. If an initial rectangle + contains both, taxels take priority. +- Clicking empty scene space clears the selection. + +Selected objects keep `SELECTED_COLOR`. Related objects use the separate +`SIDE_HIGHLIGHT_COLOR`: selecting segments highlights taxels with positive +weights to those HandJoints, while selecting taxels highlights their +positive-weight segments. Zero-weight slots are ignored. + +The **Selected Object** panel edits position and rotation, deletes selected +objects, and scales multi-selection positions about their center. Edit mode +changes Editing data; Preview mode changes Preview segments. Preview taxels are +derived and cannot be selected. With no selection, the panel remains visible, +shows Type `None` and Count `0`, and disables its input fields. +Exactly one selected Segment reveals its HandJoint editor; exactly one selected +Taxel reveals its four Segment/Weight binding slots. These editors reuse one +control set rather than creating controls for every Taxel. + +The selected-object 3-axis control supports axis-only translation. Dragging an +axis moves every selected object by the same delta. Single selections use the +object orientation; multi-selections use visualization-space axes. Rotation, +plane-translation, and scale handles are disabled. + +Action and validation errors appear in the fixed **Info** area at the bottom of +the 3D viewport. It cannot be moved. + +## Save and load + +Enter **Metadata Path**, then use **Load** or **Save**. Saving validates the +current document. If errors remain, the editor asks whether to continue; invalid +drafts can still be written. + +**New** resets the document and initializes all OpenXR HandJoint segments in the +neutral pose for **Hand Side**. The left-hand defaults are mirrored before +conversion to visualization space. + +The JSON contains only schema fields: + +- `geometry_mapping_type` and `geometry_mapping` +- `segment_source_type` and `segment_source` + +Editor HandJoint references are converted to `taxel_segment_indices` when +saving and restored when loading. Pose orientations are XYZW quaternions, and +taxel array order is the tactile force order. URDF paths, alignments, joint +values, selections, and live settings are not persisted. + +The tool does not provide SDK binary export, a separate project format, or +standalone check/export CLI commands. + +## Preview and live data + +Use **Reset Values** to restore Preview segment poses and URDF joint values from +Editing. +Choose the mutually exclusive checkbox in **Binding Metadata Source**: + +- **Editor** uses the current editor bindings and is selected by default. +- **Live** is disabled until metadata is received from the tactile stream. + +The Preview Static section lists Segments before URDF Joints. + +Preview provides four independent apply operations: + +- **Apply Segments** replaces HandJoint/segment poses. +- **Apply Taxels** replaces Editing taxels with the derived Preview taxels; + invalid rows retain their original binding poses. +- **Apply URDF** replaces URDF joint values. +- **Apply ALL** applies Segments and URDF plus the same taxels used by Apply + Taxels. + +Each Preview taxel is evaluated independently. Invalid rows are reported and +omitted from the scene while valid taxels continue to render. Live force rows +remain matched by their original taxel indices. Invalid taxels are reported in +Info. Apply Taxels and Apply ALL preserve all source rows in their original +order to match the forces; no Apply button is disabled. + +For live reception: + +1. Set **Tactile Collection ID**. +2. Optionally enable **Use JointState to Drive URDF**, set its collection ID, + and provide mappings: + + ```json + {"source_joint": ["urdf_joint", 1.0, 0.0]} + ``` + +3. Select **Connect Live**. + +HandPose replaces Preview segments with only its valid joints, retained in +HandJoint order, and is the sole source of Live Preview segment identities and +poses. Invalid joints are absent. If the hand is missing or WRIST is invalid, +Preview retains the last valid skeleton, which can still be applied to Editing, +and the status displays **Data Invalid** until a valid hand arrives. Selected +binding metadata affects only taxel deformation; it never seeds Preview +segments. Enabled JointState mappings update Preview URDF joints. Live frames +overwrite manual Static values while connected. +Tactile force magnitude maps from blue at 0 N to red at 30 N, with values above +30 N clamped. Missing or inactive force data uses the default taxel color. + +## Verify changes + +From `scripts/tactile_metadata`: + +```bash +ruff format --check tactile_metadata_editor +ruff check tactile_metadata_editor +python -m compileall -q tactile_metadata_editor +``` diff --git a/scripts/tactile_metadata/tactile_metadata_editor/__init__.py b/scripts/tactile_metadata/tactile_metadata_editor/__init__.py new file mode 100644 index 0000000000..52a7a9daf0 --- /dev/null +++ b/scripts/tactile_metadata/tactile_metadata_editor/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/scripts/tactile_metadata/tactile_metadata_editor/cli.py b/scripts/tactile_metadata/tactile_metadata_editor/cli.py new file mode 100644 index 0000000000..d50cf4defa --- /dev/null +++ b/scripts/tactile_metadata/tactile_metadata_editor/cli.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import argparse +import sys +from pathlib import Path + +from . import storage +from .document import State + + +def main() -> int: + parser = argparse.ArgumentParser(description="Local tactile metadata editor") + parser.add_argument("--port", type=int, default=8080) + parser.add_argument("--metadata", help="Metadata JSON path") + parser.add_argument("--side", choices=("left", "right"), default="left") + parser.add_argument("--urdf", help="URDF path") + args = parser.parse_args() + try: + from .ui import Editor + + state = State() + if args.metadata: + path = Path(args.metadata).expanduser().resolve() + state.replace(storage.load(path), str(path)) + Editor(state, args).run() + return 0 + except (ImportError, OSError, RuntimeError, TypeError, ValueError) as exc: + print(f"{type(exc).__name__}: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/tactile_metadata/tactile_metadata_editor/document.py b/scripts/tactile_metadata/tactile_metadata_editor/document.py new file mode 100644 index 0000000000..9e5d08f554 --- /dev/null +++ b/scripts/tactile_metadata/tactile_metadata_editor/document.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import copy +import math +from dataclasses import dataclass, field + +JOINTS = """PALM WRIST THUMB_METACARPAL THUMB_PROXIMAL THUMB_DISTAL THUMB_TIP +INDEX_METACARPAL INDEX_PROXIMAL INDEX_INTERMEDIATE INDEX_DISTAL INDEX_TIP +MIDDLE_METACARPAL MIDDLE_PROXIMAL MIDDLE_INTERMEDIATE MIDDLE_DISTAL MIDDLE_TIP +RING_METACARPAL RING_PROXIMAL RING_INTERMEDIATE RING_DISTAL RING_TIP +LITTLE_METACARPAL LITTLE_PROXIMAL LITTLE_INTERMEDIATE LITTLE_DISTAL LITTLE_TIP""".split() +IDENTITY_POSE = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0] +IDENTITY_MATRIX = [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0], +] +# Neutral right-hand OpenXR poses in wrist coordinates; left-hand mirroring is +# applied before conversion into visualization space. +# fmt: off +DEFAULT_OPENXR_SEGMENT_POSES = ( + (-0.00244, 0.00122, -0.05917, 0.010296699, 0.02060433, -0.000212213, 0.999734662), + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0), + (-0.022, -0.012, -0.028, -0.00692186, 0.30239121, 0.275661445, 0.912426661), + (-0.047208158, -0.020249943, -0.065583071, -0.00692186, 0.30239121, 0.275661445, 0.912426661), + (-0.065292271, -0.02616838, -0.09254484, -0.00692186, 0.30239121, 0.275661445, 0.912426661), + (-0.080088363, -0.031010737, -0.114604468, -0.00692186, 0.30239121, 0.275661445, 0.912426661), + (-0.00528, 0.0, -0.02068, 0.0, 0.124664128, 0.0, 0.992199), + (-0.024, 0.0, -0.094, -0.017438446, 0.039983244, 0.00069791, 0.999047924), + (-0.027194641, -0.00139598, -0.133847779, -0.043584497, 0.039951273, 0.00174431, 0.99824909), + (-0.029105295, -0.003487718, -0.157679985, -0.060999707, 0.039914746, 0.00244129, 0.997336397), + (-0.030691675, -0.005925105, -0.177467419, -0.060999707, 0.039914746, 0.00244129, 0.997336397), + (-0.00088, 0.00044, -0.02134, 0.010296699, 0.02060433, -0.000212213, 0.999734662), + (-0.004, 0.002, -0.097, -0.017452406, 0.0, 0.0, 0.999847695), + (-0.004, 0.000429523, -0.141972587, -0.043619387, 0.0, 0.0, 0.999048222), + (-0.004, -0.002010838, -0.169866039, -0.06104854, 0.0, 0.0, 0.998134798), + (-0.004, -0.004691964, -0.191702054, -0.06104854, 0.0, 0.0, 0.998134798), + (0.0033, 0.0, -0.02002, 0.0, -0.081592271, 0.0, 0.996665792), + (0.015, 0.0, -0.091, -0.017444553, -0.029990932, -0.000523494, 0.999397797), + (0.017516954, -0.001465779, -0.132898883, -0.04359976, -0.029966951, -0.001308385, 0.998598684), + (0.019070086, -0.003731828, -0.158753338, -0.06102107, -0.029939553, -0.001831181, 0.997685671), + (0.020319944, -0.006291084, -0.1795593, -0.06102107, -0.029939553, -0.001831181, 0.997685671), + (0.00704, -0.00088, -0.01804, -0.022312091, -0.184914996, -0.004199362, 0.982492229), + (0.032, -0.004, -0.082, -0.017415551, -0.064944346, -0.001133608, 0.99773626), + (0.036145766, -0.005116784, -0.113710652, -0.043527274, -0.064892417, -0.002833264, 0.996938475), + (0.038599442, -0.006772743, -0.132478637, -0.06091962, -0.064833086, -0.003965361, 0.996026981), + (0.040915463, -0.008966391, -0.150193714, -0.06091962, -0.064833086, -0.003965361, 0.996026981), +) +# fmt: on + + +@dataclass +class Segment: + joint: str + pose: list[float] = field(default_factory=lambda: IDENTITY_POSE.copy()) + + +@dataclass +class Taxel: + pose: list[float] = field(default_factory=lambda: IDENTITY_POSE.copy()) + segments: list[str | None] = field(default_factory=lambda: [None] * 4) + weights: list[float] = field(default_factory=lambda: [0.0] * 4) + # Loaded references survive zero-weight slots being displayed as None. + stored_segments: list[str | None] | None = None + + def set_segment(self, slot: int, joint: str | None) -> None: + self.segments[slot] = joint + if self.stored_segments is not None: + self.stored_segments[slot] = joint + + +def pose_error(values: object) -> str: + if not isinstance(values, list) or len(values) != 7: + return "Expected XYZ and XYZW" + if any( + type(value) not in (int, float) or not math.isfinite(value) for value in values + ): + return "Pose components must be finite numbers" + if abs(sum(value * value for value in values[3:]) - 1.0) > 1e-4: + return "Orientation must be a unit quaternion" + return "" + + +class Document: + """Editor-native data. Poses and semantic bindings never use FlatBuffer row indices.""" + + def __init__(self): + self.segments: list[Segment] = [] + self.taxels: list[Taxel] = [] + self.q: dict[str, float] = {} + self.alignment: list[list[float]] = [row.copy() for row in IDENTITY_MATRIX] + + def clone(self) -> Document: + return copy.deepcopy(self) + + def add_segment(self, joint: str, pose: list[float]) -> int: + if joint not in JOINTS: + raise ValueError("Choose a HandJoint") + self.segments.append(Segment(joint, list(pose))) + return len(self.segments) - 1 + + def add_taxels(self, poses: list[list[float]]) -> list[int]: + first = len(self.taxels) + self.taxels.extend(Taxel(list(pose)) for pose in poses) + return list(range(first, len(self.taxels))) + + def delete(self, kind: str, rows: list[int]) -> None: + values = self.segments if kind == "segments" else self.taxels + for row in sorted(set(rows), reverse=True): + if 0 <= row < len(values): + values.pop(row) + + def assign(self, taxels: list[int], segment_rows: list[int]) -> None: + if not 1 <= len(segment_rows) <= 4: + raise ValueError("Select one to four segments") + joints: list[str] = [] + for row in segment_rows: + if not 0 <= row < len(self.segments): + raise ValueError("Selected segment no longer exists") + joints.append(self.segments[row].joint) + weights = [1.0] + [0.0] * (len(joints) - 1) + for row in taxels: + if 0 <= row < len(self.taxels): + for slot, joint in enumerate(joints + [None] * (4 - len(joints))): + self.taxels[row].set_segment(slot, joint) + self.taxels[row].weights = weights + [0.0] * (4 - len(joints)) + + def set_weight(self, taxel_row: int, slot: int, value: float) -> None: + self.set_taxel_weight(self.taxels[taxel_row], slot, value) + + @staticmethod + def set_taxel_weight(taxel: Taxel, slot: int, value: float) -> None: + if taxel.segments[slot] is None: + raise ValueError("Select a segment before editing its weight") + value = float(value) + if not math.isfinite(value) or not 0.0 <= value <= 1.0: + raise ValueError("Weight must be from 0 to 1") + taxel.weights[slot] = value + + def validate(self) -> list[dict]: + issues: list[dict] = [] + + def issue(kind: str, row: int | None, field: str, message: str) -> None: + issues.append(dict(kind=kind, row=row, field=field, message=message)) + + duplicate_joints = { + joint + for joint in {segment.joint for segment in self.segments} + if sum(segment.joint == joint for segment in self.segments) > 1 + } + joints = {segment.joint for segment in self.segments} + for row, segment in enumerate(self.segments): + if segment.joint not in JOINTS: + issue("segments", row, "joint", "Invalid HandJoint") + elif segment.joint in duplicate_joints: + issue("segments", row, "joint", "HandJoint is already used") + if error := pose_error(segment.pose): + issue("segments", row, "pose", error) + for row, taxel in enumerate(self.taxels): + if error := pose_error(taxel.pose): + issue("taxels", row, "pose", error) + if len(taxel.segments) != 4 or len(taxel.weights) != 4: + issue("taxels", row, "binding", "Expected four binding slots") + continue + for slot, (joint, weight) in enumerate(zip(taxel.segments, taxel.weights)): + valid_weight = type(weight) in (int, float) and math.isfinite(weight) + if joint is None: + if not valid_weight or weight != 0.0: + issue( + "taxels", row, f"weight{slot}", "None weight must be zero" + ) + elif joint not in joints: + issue("taxels", row, f"segment{slot}", f"Missing segment {joint}") + elif not valid_weight or weight < 0.0: + issue( + "taxels", + row, + f"weight{slot}", + "Selected segment weight must be nonnegative", + ) + if not any( + joint in joints for joint in taxel.segments if joint is not None + ): + issue("taxels", row, "binding", "Taxel is not bound to a segment") + numeric = all( + type(value) in (int, float) and math.isfinite(value) + for value in taxel.weights + ) + if not numeric or abs(sum(taxel.weights) - 1.0) > 1e-5: + issue("taxels", row, "weights", "Weights must sum to one") + return issues + + +class State: + def __init__(self): + self.document = Document() + self.saved: Document | None = None + self.path = "" + + def replace(self, document: Document, path: str | None = None) -> None: + self.document = document + self.saved = document.clone() + self.saved.q = {} + if path is not None: + self.path = path diff --git a/scripts/tactile_metadata/tactile_metadata_editor/editor_interactions.js b/scripts/tactile_metadata/tactile_metadata_editor/editor_interactions.js new file mode 100644 index 0000000000..a2a54cb052 --- /dev/null +++ b/scripts/tactile_metadata/tactile_metadata_editor/editor_interactions.js @@ -0,0 +1,365 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +(() => { + "use strict"; + + const key = "__tactileMetadataEditorInteractions"; + globalThis[key]?.stop(); + + let observer = null; + let infoInput = null; + let infoElement = null; + let selectedNamesTextarea = null; + let raf = 0; + const numberDrafts = new Map(); + const nativeInputValueSetter = Object.getOwnPropertyDescriptor( + HTMLInputElement.prototype, + "value", + )?.set; + + const findBridge = (name) => { + const label = [...document.querySelectorAll("label")].find( + (candidate) => candidate.textContent?.trim() === name, + ); + let input = label?.htmlFor ? document.getElementById(label.htmlFor) : null; + let row = label?.parentElement; + while (!input && row) { + input = row.querySelector("input"); + row = row.parentElement; + } + if (!(input instanceof HTMLInputElement)) return null; + input.style.position = "absolute"; + input.style.width = "1px"; + input.style.height = "1px"; + input.style.opacity = "0"; + input.style.pointerEvents = "none"; + input.tabIndex = -1; + label.style.display = "none"; + return input; + }; + + // Place Viser's logo above the fixed bottom Info area. + const positionViserLogo = () => { + const logo = document.querySelector('a[title="About Viser"]'); + if (!logo) return; + logo.style.top = "auto"; + logo.style.left = "16px"; + logo.style.bottom = "144px"; + }; + + const syncInfo = () => { + const parent = globalThis.__viserMutable?.canvas?.parentElement; + if (!parent) return; + if (!infoElement) { + infoElement = document.createElement("div"); + Object.assign(infoElement.style, { + position: "absolute", + left: "12px", + right: "12px", + bottom: "12px", + minHeight: "64px", + maxHeight: "120px", + overflow: "auto", + zIndex: "20", + padding: "10px 14px", + border: "1px solid rgba(0,0,0,0.25)", + borderRadius: "6px", + background: "rgba(255,255,255,0.9)", + color: "#202020", + boxSizing: "border-box", + }); + } + if (infoElement.parentElement !== parent) parent.appendChild(infoElement); + if (!infoInput?.isConnected) { + infoInput = findBridge("__Tactile Info"); + } + infoElement.innerHTML = `Info
${infoInput?.value || " "}
`; + }; + + const frame = () => { + for (const [input, draft] of numberDrafts) { + if (!input.isConnected || document.activeElement !== input) { + numberDrafts.delete(input); + } else if (input.value !== draft) { + nativeInputValueSetter?.call(input, draft); + input.dispatchEvent(new Event("input", { bubbles: true })); + input.setSelectionRange(draft.length, draft.length); + } + } + if ( + selectedNamesTextarea?.isConnected && + selectedNamesTextarea.dataset.tactileResizing !== "true" + ) { + const height = selectedNamesTextarea.dataset.tactileHeight; + if (height && selectedNamesTextarea.style.height !== height) { + selectedNamesTextarea.style.height = height; + } + } + syncInfo(); + raf = requestAnimationFrame(frame); + }; + + const boldChoiceLabels = () => { + for (const label of document.querySelectorAll("label")) { + if (["Edit", "Preview", "Editor", "Live"].includes(label.textContent?.trim())) { + label.style.fontWeight = "700"; + } + } + }; + + const directChild = (parent, descendant) => { + let child = descendant; + while (child?.parentElement && child.parentElement !== parent) { + child = child.parentElement; + } + return child?.parentElement === parent ? child : null; + }; + + const layoutSelectedNames = () => { + for (const label of document.querySelectorAll("label")) { + if (label.textContent?.trim() !== "__Selected Object Names") continue; + let row = label.parentElement; + while (row && !row.querySelector("textarea")) row = row.parentElement; + if (!row) continue; + const textarea = row.querySelector("textarea"); + const labelColumn = directChild(row, label); + const valueColumn = directChild(row, textarea); + if (!labelColumn || !valueColumn) continue; + labelColumn.style.display = "none"; + valueColumn.style.flex = "1 1 100%"; + valueColumn.style.width = "100%"; + row.style.width = "100%"; + selectedNamesTextarea = textarea; + if (textarea.dataset.tactileFixedRows !== "true") { + const style = getComputedStyle(textarea); + const fontSize = Number.parseFloat(style.fontSize); + const lineHeight = + Number.parseFloat(style.lineHeight) || (Number.isFinite(fontSize) ? fontSize * 1.5 : 20); + const chrome = ["paddingTop", "paddingBottom", "borderTopWidth", "borderBottomWidth"] + .map((name) => Number.parseFloat(style[name]) || 0) + .reduce((sum, value) => sum + value, 0); + textarea.rows = 2; + const height = `${Math.ceil(2 * lineHeight + chrome)}px`; + textarea.style.height = height; + textarea.style.overflowY = "auto"; + textarea.dataset.tactileHeight = height; + textarea.dataset.tactileFixedRows = "true"; + textarea.addEventListener("pointerdown", (event) => { + const rect = textarea.getBoundingClientRect(); + if (event.clientX < rect.right - 20 || event.clientY < rect.bottom - 20) { + return; + } + textarea.dataset.tactileResizing = "true"; + const finish = () => { + textarea.dataset.tactileHeight = `${textarea.getBoundingClientRect().height}px`; + delete textarea.dataset.tactileResizing; + window.removeEventListener("pointerup", finish); + window.removeEventListener("pointercancel", finish); + }; + window.addEventListener("pointerup", finish, { once: true }); + window.addEventListener("pointercancel", finish, { once: true }); + }); + } + } + }; + + const layoutTaxelBindings = () => { + const names = [ + "Segment 0", + "Segment 1", + "Segment 2", + "Segment 3", + "Weight 0", + "Weight 1", + "Weight 2", + "Weight 3", + ]; + const clean = (label) => label.textContent?.trim().replace(/^⚠\s*/, ""); + for (const first of document.querySelectorAll("label")) { + if (clean(first) !== names[0]) continue; + let container = first.parentElement; + let fields = []; + while (container) { + const labels = [...container.querySelectorAll("label")]; + fields = names.map((name) => labels.find((label) => clean(label) === name)); + if (fields.every(Boolean)) break; + container = container.parentElement; + } + if (!container) continue; + const children = fields.map((field) => directChild(container, field)); + if (children.some((child) => !child) || new Set(children).size !== names.length) { + continue; + } + const title = container.querySelector("[data-tactile-binding-title]"); + const titleChild = directChild(container, title); + if (!title || !titleChild) continue; + const bindingChildren = new Set(children); + Object.assign(container.style, { + display: "grid", + gridTemplateColumns: "6em repeat(4, minmax(0, 1fr))", + gap: "0", + alignItems: "end", + }); + // Only bindings use the four columns; all other taxel controls stay full-width. + for (const child of container.children) { + if (child === titleChild) { + child.style.gridColumn = "1"; + child.style.gridRow = "span 2"; + child.style.alignSelf = "center"; + child.style.margin = "0"; + child.style.padding = "0 0.5em"; + } else if (bindingChildren.has(child)) { + child.style.gridColumn = "auto"; + child.style.minWidth = "0"; + child.style.margin = "0"; + child.style.padding = "0"; + } else { + child.style.gridColumn = "1 / -1"; + } + } + Object.assign(title.style, { + color: "var(--mantine-color-dimmed)", + fontSize: "0.875em", + fontWeight: "450", + }); + for (const field of fields) { + let row = field.parentElement; + while (row && !row.querySelector("input, button")) row = row.parentElement; + if (!row) continue; + const control = row.querySelector("input, button"); + const labelColumn = directChild(row, field); + const valueColumn = directChild(row, control); + if (!labelColumn || !valueColumn) continue; + // Hiding only the label text leaves Viser's fixed-width label column empty. + labelColumn.style.display = "none"; + valueColumn.style.flex = "1 1 100%"; + valueColumn.style.width = "100%"; + valueColumn.style.minWidth = "0"; + row.style.width = "100%"; + row.style.margin = "0"; + control.style.borderRadius = "0"; + } + if (title.dataset.tactileBindingTitle === "add") { + const capture = [...container.querySelectorAll("button")].find( + (button) => button.textContent?.trim() === "Capture URDF Surface Pose", + ); + const captureChild = directChild(container, capture); + if (captureChild) captureChild.style.marginTop = "0.5em"; + } + } + }; + + const layoutTaxelPagination = () => { + const containers = new Set(); + for (const label of document.querySelectorAll("label")) { + if (!label.textContent?.trim().startsWith("__Taxel Page Jump ")) continue; + let container = label.parentElement; + let controls = []; + while (container) { + const jumpChild = directChild(container, label); + if (jumpChild) { + controls = [jumpChild]; + for (let index = 0; index < 3; ++index) { + controls.push( + controls[controls.length - 1]?.nextElementSibling ?? null, + ); + } + if (controls.every(Boolean)) break; + } + container = container.parentElement; + } + if (!container || controls.some((control) => !control)) continue; + containers.add(container); + const columns = ["1", "2", "1", "2"]; + controls.forEach((control, index) => { + control.dataset.tactilePaginationControl = "true"; + control.style.gridColumn = columns[index]; + control.style.minWidth = "0"; + }); + const input = controls[0].querySelector("input"); + if (label && input) { + const row = directChild(controls[0], input); + const labelColumn = row ? directChild(row, label) : null; + const valueColumn = row ? directChild(row, input) : null; + if (labelColumn && valueColumn) { + labelColumn.style.display = "none"; + valueColumn.style.flex = "1 1 100%"; + valueColumn.style.width = "100%"; + } + } + } + for (const container of containers) { + Object.assign(container.style, { + display: "grid", + gridTemplateColumns: "minmax(0, 1fr) auto", + columnGap: "0.5em", + rowGap: "0.5em", + alignItems: "center", + }); + for (const child of container.children) { + if (child.dataset.tactilePaginationControl !== "true") { + child.style.gridColumn = "1 / -1"; + } + } + } + }; + + const preservePartialNumbers = () => { + for (const input of document.querySelectorAll("input")) { + if (input.dataset.tactileNumberDraft === "true") continue; + if (!["decimal", "numeric"].includes(input.inputMode) && input.type !== "number") { + continue; + } + input.dataset.tactileNumberDraft = "true"; + input.addEventListener("input", () => { + const value = input.value; + const parsed = Number(value); + if ( + value === "" || + !Number.isFinite(parsed) || + String(parsed) !== value + ) { + numberDrafts.set(input, value); + } else { + numberDrafts.delete(input); + } + }); + input.addEventListener("blur", () => numberDrafts.delete(input)); + } + }; + + const bind = () => { + findBridge("__Tactile Info"); + positionViserLogo(); + boldChoiceLabels(); + layoutSelectedNames(); + preservePartialNumbers(); + layoutTaxelBindings(); + layoutTaxelPagination(); + for (const label of document.querySelectorAll("label")) { + if (!label.textContent?.trim().startsWith("⚠")) continue; + const field = label.parentElement?.querySelector("input, button"); + if (field) { + field.style.outline = "1px solid #d32f2f"; + field.style.outlineOffset = "1px"; + } + } + }; + + observer = new MutationObserver(bind); + observer.observe(document.body, { childList: true, subtree: true }); + bind(); + frame(); + globalThis[key] = { + stop() { + observer?.disconnect(); + observer = null; + cancelAnimationFrame(raf); + infoElement?.remove(); + infoElement = null; + infoInput = null; + selectedNamesTextarea = null; + numberDrafts.clear(); + }, + }; +})(); diff --git a/scripts/tactile_metadata/tactile_metadata_editor/geometry.py b/scripts/tactile_metadata/tactile_metadata_editor/geometry.py new file mode 100644 index 0000000000..cb958ab4bf --- /dev/null +++ b/scripts/tactile_metadata/tactile_metadata_editor/geometry.py @@ -0,0 +1,370 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from collections import OrderedDict +from collections.abc import Iterable +from pathlib import Path + +import numpy as np +import trimesh +from numpy.typing import NDArray +from trimesh import transformations as tf +from yourdfpy import URDF + +from .document import ( + DEFAULT_OPENXR_SEGMENT_POSES, + Document, + JOINTS, + Segment, + Taxel, +) + +VISUALIZATION_FROM_HAND = np.array( + [ + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, -1.0, 0.0], + [0.0, -1.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=float, +) +VISUALIZATION_FROM_URDF = np.array( + [ + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=float, +) + + +def matrix(values: list[float]) -> NDArray[np.float64]: + array = np.asarray(values, dtype=np.float64) + if ( + array.shape != (7,) + or not np.isfinite(array).all() + or abs(array[3:] @ array[3:] - 1.0) > 1e-4 + ): + raise ValueError("Invalid XYZ/XYZW pose") + result = tf.quaternion_matrix(array[[6, 3, 4, 5]]) + result[:3, 3] = array[:3] + return result + + +def pose(transform: np.ndarray) -> list[float]: + quaternion = tf.quaternion_from_matrix(transform) + return [ + *np.asarray(transform[:3, 3], dtype=float).tolist(), + *quaternion[[1, 2, 3, 0]].tolist(), + ] + + +def default_segments(side: str) -> list[Segment]: + poses = [list(value) for value in DEFAULT_OPENXR_SEGMENT_POSES] + if side == "left": + for value in poses: + value[0], value[4], value[5] = -value[0], -value[4], -value[5] + return [ + Segment(joint, pose(VISUALIZATION_FROM_HAND @ matrix(value))) + for joint, value in zip(JOINTS, poses) + ] + + +def transform(position, rotation_degrees) -> np.ndarray: + result = tf.euler_matrix(*np.deg2rad(rotation_degrees), axes="sxyz") + result[:3, 3] = np.asarray(position, dtype=float) + return result + + +def components(value: np.ndarray) -> tuple[tuple[float, ...], tuple[float, ...]]: + return ( + tuple(value[:3, 3]), + tuple(np.rad2deg(tf.euler_from_matrix(value, axes="sxyz"))), + ) + + +def unit(value) -> np.ndarray: + value = np.asarray(value, dtype=float) + norm = np.linalg.norm(value) + if not np.isfinite(norm) or norm < 1e-10: + raise ValueError("Expected a finite nonzero direction") + return value / norm + + +def surface_frame(origin, normal) -> np.ndarray: + z = unit(normal) + x = np.array([1.0, 0.0, 0.0]) + x -= z * (x @ z) + if np.linalg.norm(x) < 1e-6: + x = np.array([0.0, 1.0, 0.0]) + x -= z * (x @ z) + x = unit(x) + result = np.eye(4) + result[:3, :3] = np.column_stack((x, np.cross(z, x), z)) + result[:3, 3] = origin + return result + + +def realign(document: Document, alignment: np.ndarray) -> np.ndarray: + old = np.asarray(document.alignment, dtype=float) + delta = alignment @ np.linalg.inv(old) + for value in (*document.segments, *document.taxels): + value.pose = pose(delta @ matrix(value.pose)) + document.alignment = alignment.tolist() + return delta + + +def apply_delta(values, rows: Iterable[int], delta: np.ndarray) -> None: + """Apply the transform once per distinct row.""" + for row in set(rows): + if 0 <= row < len(values): + values[row].pose = pose(delta @ matrix(values[row].pose)) + + +def scale_positions(values, rows: list[int], factor: float) -> None: + if not np.isfinite(factor) or factor <= 0: + raise ValueError("Scale must be finite and positive") + points = np.asarray([values[row].pose[:3] for row in rows], dtype=float) + center = points.mean(axis=0) + for row in rows: + values[row].pose[:3] = ( + center + (values[row].pose[:3] - center) * factor + ).tolist() + + +def grid_poses( + center, width: float, height: float, rows: int, columns: int, frame=None +) -> list[list[float]]: + if rows < 1 or columns < 1: + raise ValueError("Rows and Columns must be positive") + if not np.isfinite(width) or not np.isfinite(height) or width < 0 or height < 0: + raise ValueError("Grid dimensions must be finite and nonnegative") + basis = np.eye(4) if frame is None else np.asarray(frame, dtype=float).copy() + basis[:3, 3] = np.asarray(center, dtype=float) + result = [] + for row in range(rows): + y = 0.0 if rows == 1 else (row / (rows - 1) - 0.5) * height + for column in range(columns): + x = 0.0 if columns == 1 else (column / (columns - 1) - 0.5) * width + current = basis.copy() + current[:3, 3] += basis[:3, :3] @ [x, y, 0.0] + result.append(pose(current)) + return result + + +class Model: + def __init__(self, path: str): + self.path = Path(path).resolve() + self.handles: dict[str, object] = {} + self._fk_cache: OrderedDict[tuple[float, ...], dict[str, np.ndarray]] = ( + OrderedDict() + ) + + def resource(fname): + name = str(fname) + if name.startswith("package://"): + package, relative = name[10:].split("/", 1) + root = self.path.parent + while root != root.parent and root.name != package: + root = root.parent + if root.name == package: + return str((root / relative).resolve()) + for root in self.path.parents: + candidate = root / relative + if candidate.is_file(): + return str(candidate) + raise ValueError(f"Cannot resolve package resource: {name}") + elif "://" in name: + raise ValueError(f"Only local mesh resources are supported: {name}") + return str((self.path.parent / name).resolve()) + + self.urdf = URDF.load(str(self.path), filename_handler=resource) + self.joints = self.urdf.actuated_joints + self.links = list(self.urdf.link_map) + self.surfaces: dict[str, tuple[str, trimesh.Trimesh]] = {} + scene = self.urdf.scene + parents = {child: parent for parent, child, _ in scene.graph.to_edgelist()} + for node in scene.graph.nodes_geometry: + link = node + while link not in self.links: + link = parents[link] + local, geometry_name = scene.graph.get(frame_to=node, frame_from=link) + mesh = scene.geometry[geometry_name].copy() + mesh.apply_transform(local) + for index, part in enumerate(mesh.split(only_watertight=False)): + self.surfaces[f"{link}/{node}/{index}"] = (link, part) + + def raw_fk(self, q: dict[str, float]) -> dict[str, np.ndarray]: + configuration = tuple(float(q.get(joint.name, 0.0)) for joint in self.joints) + cached = self._fk_cache.pop(configuration, None) + if cached is None: + self.urdf.update_cfg( + dict(zip((j.name for j in self.joints), configuration)) + ) + cached = {link: self.urdf.get_transform(link).copy() for link in self.links} + self._fk_cache[configuration] = cached + if len(self._fk_cache) > 4: + self._fk_cache.popitem(last=False) + return cached + + @staticmethod + def _rigid(raw: np.ndarray, alignment: np.ndarray, scale: float) -> np.ndarray: + result = alignment.copy() + result[:3, :3] = alignment[:3, :3] @ raw[:3, :3] + result[:3, 3] = alignment[:3, :3] @ (raw[:3, 3] * scale) + alignment[:3, 3] + return result + + def draw( + self, + scene, + q: dict[str, float], + alignment: np.ndarray, + scale: float, + opacity: float, + highlighted: bool = False, + ) -> None: + raw = self.raw_fk(q) + for name, (link, mesh) in self.surfaces.items(): + if name not in self.handles: + self.handles[name] = scene.add_mesh_simple( + "/urdf/" + name, + mesh.vertices, + mesh.faces, + color=(170, 175, 180), + opacity=opacity, + side="double", + ) + current = self._rigid(raw[link], alignment, scale) + handle = self.handles[name] + handle.position = current[:3, 3] + handle.wxyz = tf.quaternion_from_matrix(current) + handle.scale = scale + handle.opacity = opacity + handle.color = (255, 210, 45) if highlighted else (170, 175, 180) + handle.visible = True + + def hit( + self, + origin, + direction, + q: dict[str, float], + alignment: np.ndarray, + scale: float, + ): + nearest, transforms = None, self.raw_fk(q) + for link, mesh in self.surfaces.values(): + raw = transforms[link] + similarity = alignment @ np.diag([scale, scale, scale, 1.0]) @ raw + inverse = np.linalg.inv(similarity) + local_origin = (inverse @ np.r_[origin, 1.0])[:3] + local_direction = inverse[:3, :3] @ direction + points, _, faces = mesh.ray.intersects_location( + [local_origin], [local_direction] + ) + for point, face in zip(points, faces): + world = (similarity @ np.r_[point, 1.0])[:3] + distance = np.linalg.norm(world - origin) + if nearest is None or distance < nearest[0]: + normal = ( + np.linalg.inv(similarity[:3, :3]).T @ mesh.face_normals[face] + ) + nearest = (distance, world, unit(normal)) + return nearest + + def remove(self) -> None: + for handle in self.handles.values(): + handle.remove() + self.handles.clear() + + +def deform( + binding: Document, current: list[Segment] +) -> tuple[list[Taxel | None], list[str]]: + bind_segments = { + segment.joint: matrix(segment.pose) for segment in binding.segments + } + current_segments = {segment.joint: matrix(segment.pose) for segment in current} + if len(bind_segments) != len(binding.segments) or len(current_segments) != len( + current + ): + raise ValueError("Segment HandJoint values must be unique") + invalid_taxels: dict[int, list[str]] = {} + for issue in binding.validate(): + if issue["kind"] == "taxels" and issue["row"] is not None: + invalid_taxels.setdefault(issue["row"], []).append(issue["message"]) + result: list[Taxel | None] = [] + errors = [] + for row, taxel in enumerate(binding.taxels): + if row in invalid_taxels: + result.append(None) + errors.append( + f"Taxel {row}: " + "; ".join(dict.fromkeys(invalid_taxels[row])) + ) + continue + try: + active = [ + (joint, weight) + for joint, weight in zip(taxel.segments, taxel.weights) + if joint is not None and weight > 0 + ] + if not active: + raise ValueError("no active segment") + transforms = [] + for joint, weight in active: + if joint not in bind_segments or joint not in current_segments: + raise ValueError(f"missing segment {joint}") + transforms.append( + ( + weight, + current_segments[joint] @ np.linalg.inv(bind_segments[joint]), + ) + ) + bind_taxel = matrix(taxel.pose) + position = sum( + weight * (delta @ np.r_[bind_taxel[:3, 3], 1.0])[:3] + for weight, delta in transforms + ) + if len({joint for joint, _ in active}) == 1: + rotation = transforms[0][1][:3, :3] @ bind_taxel[:3, :3] + else: + blend = sum(weight * delta[:3, :3] for weight, delta in transforms) + u, singular, vh = np.linalg.svd(blend) + if singular[-1] < 1e-8: + rotation = bind_taxel[:3, :3] + else: + correction = np.eye(3) + correction[2, 2] = np.linalg.det(u @ vh) + rotation = u @ correction @ vh @ bind_taxel[:3, :3] + value = np.eye(4) + value[:3, :3], value[:3, 3] = rotation, position + result.append(Taxel(pose(value), list(taxel.segments), list(taxel.weights))) + except ( + ValueError, + TypeError, + KeyError, + IndexError, + np.linalg.LinAlgError, + ) as exc: + result.append(None) + errors.append(f"Taxel {row}: {exc}") + return result, errors + + +def hand_segments(hand: dict, alignment: np.ndarray) -> list[Segment]: + values = [] + for position, orientation, valid in zip( + hand["positions"], hand["orientations"], hand["is_valid"] + ): + values.append(matrix([*position, *orientation]) if valid else None) + if len(values) != len(JOINTS): + raise ValueError("HandTracker joint array has an invalid length") + if values[1] is None: + return [] + wrist_inverse = np.linalg.inv(values[1]) + return [ + Segment(joint, pose(alignment @ wrist_inverse @ tracked)) + for joint, tracked in zip(JOINTS, values) + if tracked is not None + ] diff --git a/scripts/tactile_metadata/tactile_metadata_editor/sdk.py b/scripts/tactile_metadata/tactile_metadata_editor/sdk.py new file mode 100644 index 0000000000..51a420e2b2 --- /dev/null +++ b/scripts/tactile_metadata/tactile_metadata_editor/sdk.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import json +import multiprocessing +import sys +import threading +import time + +from .document import JOINTS + + +class Receiver: + def __init__(self, config: dict): + context = multiprocessing.get_context("spawn") + context.set_executable(sys.executable) + self.connection, child = context.Pipe(duplex=False) + self.latest, self.status, self.active = {}, "Connecting", False + self.process = context.Process(target=live, args=(child, config), daemon=True) + self.process.start() + child.close() + self.thread = threading.Thread(target=self._receive, daemon=True) + self.thread.start() + + def _receive(self) -> None: + try: + while True: + value = json.loads(self.connection.recv_bytes(16 * 1024 * 1024)) + self.status = value.get("status", "Live") + if "frame" in value: + self.latest = value["frame"] + self.active = True + except (EOFError, OSError, ValueError): + self.active = False + if self.status in ("Connecting", "Live"): + self.status = "Disconnected" + + def close(self) -> None: + self.process.terminate() + self.process.join(timeout=2) + if self.process.is_alive(): + self.process.kill() + self.process.join() + self.connection.close() + self.thread.join(timeout=2) + + +def _pose_rows(values) -> list[dict]: + return [ + { + "position": dict(zip("xyz", row[:3])), + "orientation": dict(zip("xyzw", row[3:])), + } + for row in values.copy().tolist() + ] + + +def _metadata(value) -> dict | None: + if value is None or value.geometry_mapping is None or value.segment_source is None: + return None + mapping = value.geometry_mapping + return { + "geometry_mapping_type": "LBS", + "geometry_mapping": { + "segment_bind_poses": _pose_rows(mapping.segment_bind_poses), + "taxel_bind_poses": _pose_rows(mapping.taxel_bind_poses), + "taxel_segment_indices": [ + {"values": row} for row in mapping.taxel_segment_indices.copy().tolist() + ], + "taxel_segment_weights": [ + {"values": row} for row in mapping.taxel_segment_weights.copy().tolist() + ], + }, + "segment_source_type": "HandSkeletonRef", + "segment_source": { + "joint_ids": [ + JOINTS[int(joint)] for joint in value.segment_source.joint_ids + ] + }, + } + + +def _hand(value) -> dict | None: + if value is None: + return None + joints = value.joints + return { + name: getattr(joints, name).copy().tolist() + for name in ("positions", "orientations", "is_valid") + } + + +def live(connection, config: dict) -> None: + try: + from isaaccapture.deviceio_session import DeviceIOSession + from isaaccapture.deviceio_trackers import ( + HandTracker, + JointStateTracker, + TactileTracker, + ) + from isaaccapture.oxr import OpenXRSession + + hand = HandTracker() + tactile = TactileTracker(config["collection"]) + joint = ( + JointStateTracker(config["joint_state"]) + if config.get("joint_state") + else None + ) + trackers = [hand, tactile, *([joint] if joint else [])] + with OpenXRSession( + "TactileMetadataEditor", DeviceIOSession.get_required_extensions(trackers) + ) as oxr: + with DeviceIOSession.run(trackers, oxr.get_handles()) as session: + while True: + session.update() + metadata = tactile.get_metadata(session) + data = tactile.get_data(session) + hand_value = getattr(hand, f"get_{config['side']}_hand")(session) + joint_value = joint.get_data(session) if joint else None + frame = { + "metadata": _metadata(metadata), + "forces": data.taxel_forces.copy().tolist() + if data is not None + else None, + "hand": _hand(hand_value), + "joints": { + item.name: [item.position, item.valid] + for item in joint_value.joints + } + if joint_value is not None + else {}, + "read_time": time.monotonic(), + } + connection.send_bytes(json.dumps({"frame": frame}).encode()) + time.sleep(1 / 30) + except Exception as exc: + connection.send_bytes( + json.dumps({"status": f"{type(exc).__name__}: {exc}"}).encode() + ) + finally: + connection.close() diff --git a/scripts/tactile_metadata/tactile_metadata_editor/storage.py b/scripts/tactile_metadata/tactile_metadata_editor/storage.py new file mode 100644 index 0000000000..79d09e0818 --- /dev/null +++ b/scripts/tactile_metadata/tactile_metadata_editor/storage.py @@ -0,0 +1,169 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import json +import os +import tempfile +from pathlib import Path + +import numpy as np + +from .document import Document, JOINTS, Segment, State, Taxel +from .geometry import VISUALIZATION_FROM_HAND, matrix, pose + + +def _pose_row(value: object) -> list[float]: + if not isinstance(value, dict): + raise ValueError("Pose row must be an object") + position, orientation = value.get("position"), value.get("orientation") + if not isinstance(position, dict) or not isinstance(orientation, dict): + raise ValueError("Pose row requires position and orientation") + return [ + *[position.get(key) for key in "xyz"], + *[orientation.get(key) for key in "xyzw"], + ] + + +def _rows(section: dict, name: str) -> list: + value = section.get(name) + if not isinstance(value, list): + raise ValueError(f"{name}: expected an array") + return value + + +def _joint_name(value: object) -> str: + if isinstance(value, str): + return value + if type(value) is int and 0 <= value < len(JOINTS): + return JOINTS[value] + return f"INVALID_HAND_JOINT_{value}" + + +def parse( + value: object, + alignment: np.ndarray | list[list[float]] | None = None, +) -> Document: + if not isinstance(value, dict): + raise ValueError("Expected a TactileMetadata object") + geometry, source = value.get("geometry_mapping"), value.get("segment_source") + if not isinstance(geometry, dict) or not isinstance(source, dict): + raise ValueError("Expected geometry_mapping and segment_source objects") + joints = source.get("joint_ids") + if not isinstance(joints, list): + raise ValueError("joint_ids: expected an array") + segment_rows = _rows(geometry, "segment_bind_poses") + taxel_rows = _rows(geometry, "taxel_bind_poses") + index_rows = _rows(geometry, "taxel_segment_indices") + weight_rows = _rows(geometry, "taxel_segment_weights") + if len(segment_rows) != len(joints): + raise ValueError("segment_bind_poses must match joint_ids") + if len({len(taxel_rows), len(index_rows), len(weight_rows)}) != 1: + raise ValueError("All taxel arrays must have the same length") + + transform = np.asarray( + VISUALIZATION_FROM_HAND if alignment is None else alignment, dtype=float + ) + document = Document() + document.alignment = transform.tolist() + names = [_joint_name(joint) for joint in joints] + document.segments = [ + Segment(joint, pose(transform @ matrix(_pose_row(row)))) + for joint, row in zip(names, segment_rows) + ] + for raw_pose, raw_indices, raw_weights in zip(taxel_rows, index_rows, weight_rows): + indices = raw_indices.get("values") if isinstance(raw_indices, dict) else None + weights = raw_weights.get("values") if isinstance(raw_weights, dict) else None + if not isinstance(indices, list) or not isinstance(weights, list): + raise ValueError("Taxel indices and weights require values arrays") + if len(indices) != 4 or len(weights) != 4: + raise ValueError("Taxel bindings require four slots") + bindings = [ + names[index] + if type(index) is int and 0 <= index < len(names) + else f"INVALID_SEGMENT_INDEX_{index}" + for index in indices + ] + document.taxels.append( + Taxel( + pose(transform @ matrix(_pose_row(raw_pose))), + [ + None if type(weight) in (int, float) and weight == 0 else joint + for joint, weight in zip(bindings, weights) + ], + list(weights), + bindings, + ) + ) + return document + + +def _json_pose(values: list[float]) -> dict: + return { + "position": dict(zip("xyz", values[:3])), + "orientation": dict(zip("xyzw", values[3:])), + } + + +def metadata(document: Document) -> dict: + inverse = np.linalg.inv(np.asarray(document.alignment, dtype=float)) + joints = [segment.joint for segment in document.segments] + indices = {joint: row for row, joint in enumerate(joints)} + segment_poses = [ + _json_pose(pose(inverse @ matrix(segment.pose))) + for segment in document.segments + ] + taxel_poses = [ + _json_pose(pose(inverse @ matrix(taxel.pose))) for taxel in document.taxels + ] + binding_indices = [ + { + "values": [ + indices.get(joint, 0 if taxel.segments[slot] is None else 0xFFFFFFFF) + for slot, joint in enumerate( + taxel.segments + if taxel.stored_segments is None + else taxel.stored_segments + ) + ] + } + for taxel in document.taxels + ] + binding_weights = [{"values": list(taxel.weights)} for taxel in document.taxels] + return { + "geometry_mapping_type": "LBS", + "geometry_mapping": { + "segment_bind_poses": segment_poses, + "taxel_bind_poses": taxel_poses, + "taxel_segment_indices": binding_indices, + "taxel_segment_weights": binding_weights, + }, + "segment_source_type": "HandSkeletonRef", + "segment_source": {"joint_ids": joints}, + } + + +def load( + path: str | Path, + alignment: np.ndarray | list[list[float]] | None = None, +) -> Document: + with Path(path).expanduser().resolve().open() as stream: + return parse(json.load(stream), alignment) + + +def write(path: str | Path, value: object) -> None: + path = Path(path).expanduser().resolve() + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary = tempfile.mkstemp(dir=path.parent, prefix=".tactile-") + try: + with os.fdopen(fd, "w") as stream: + json.dump(value, stream, indent=2) + stream.write("\n") + os.replace(temporary, path) + finally: + Path(temporary).unlink(missing_ok=True) + + +def save(state: State, path: str | Path) -> None: + write(path, metadata(state.document)) + state.replace(state.document, str(Path(path).expanduser().resolve())) diff --git a/scripts/tactile_metadata/tactile_metadata_editor/ui.py b/scripts/tactile_metadata/tactile_metadata_editor/ui.py new file mode 100644 index 0000000000..faf5204614 --- /dev/null +++ b/scripts/tactile_metadata/tactile_metadata_editor/ui.py @@ -0,0 +1,2346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import copy +import html +import json +import logging +import math +import queue +import time +from dataclasses import replace +from pathlib import Path + +import numpy as np +import trimesh +import viser +from trimesh import transformations as tf +from viser import _messages + +from . import geometry as geo, sdk, storage +from .document import Document, IDENTITY_POSE, JOINTS, Segment, State, Taxel + +LOG = logging.getLogger("isaaccapture.tactile_metadata.editor") + +SEGMENT_CUBE_SIDE_M = 0.004 +TAXEL_SIDE_M = 0.0015 +URDF_OPACITY = 0.2 +SEGMENT_FLAT_SHADING = True +VISUALIZER_CAST_SHADOW = False +VISUALIZER_RECEIVE_SHADOW = False +SEGMENT_COLOR = (200, 120, 0) +TAXEL_COLOR = (200, 120, 0) +SELECTED_COLOR = (0, 255, 0) +SIDE_HIGHLIGHT_COLOR = (0, 180, 255) +INACTIVE_COLOR = (200, 200, 200) +ERROR_COLOR = (255, 0, 0) +FORCE_MIN_COLOR = (0, 0, 255) +FORCE_MAX_COLOR = (255, 0, 0) +FORCE_MAX_N = 30.0 +SELECTION_GIZMO_SCALE_M = 0.01 +SELECTION_GIZMO_LINE_WIDTH = 2.5 +ACTIVE_DATASET_OPACITY = 1.0 +BINDING_OVERLAY_OPACITY = 0.4 +SAVED_BINDING_OPACITY = 0.3 +TAXEL_PAGE_SIZE = 25 + + +class Editor: + CAMERA_POSITION = (0.3, 0.3, -0.4) + CAMERA_LOOK_AT = (0.0, 0.08, 0.0) + + def __init__(self, state: State, args): + self.state = state + self.document.alignment = geo.VISUALIZATION_FROM_HAND.tolist() + self.preview_document = self._preview_from(self.document) + self.live_document: Document | None = None + self.live_metadata_key = "" + self.receiver: sdk.Receiver | None = None + self.model: geo.Model | None = None + self.urdf_alignment = geo.VISUALIZATION_FROM_URDF.copy() + self.urdf_scale = 1.0 + self.events: queue.SimpleQueue = queue.SimpleQueue() + self.scene_handles: dict[str, object] = {} + self.segment_row_folder = None + self.taxel_row_folder = None + self.taxel_page_container = None + self.edit_joint_folder = None + self.preview_joint_folder = None + self.preview_segment_folder = None + self.selection_folder = None + self.segment_widgets = [] + self.taxel_widgets = {} + self.segment_item_folders = [] + self.segment_error_widgets = [] + self.taxel_error_widgets = {} + self.taxel_page_rows = [] + self.taxel_page_jumps = [] + self.taxel_previous_buttons = [] + self.taxel_next_buttons = [] + self.taxel_page_info = None + self.edit_joint_widgets, self.preview_joint_widgets = {}, {} + self.preview_segment_widgets = [] + self.new_taxel_binding = Taxel() + self.new_taxel_binding_widgets = [] + self.taxel_page = 0 + self.selected_kind: str | None = None + self.selected_rows: list[int] = [] + self.binding_phase: str | None = None + self.binding_taxels: list[int] = [] + self.binding_segments: list[int] = [] + self.pick_mode: str | None = None + self.selection_rotation = np.zeros(3) + self.selection_scale = 1.0 + self.selection_gizmo_dragging = False + self.selection_gizmo_last_position: np.ndarray | None = None + self.redraw = True + self.server = viser.ViserServer(host="0.0.0.0", port=args.port) + self.server.scene.set_up_direction("+y") + for script in ("editor_interactions.js", "viewport_axes.js"): + source = Path(__file__).with_name(script).read_text() + self.server._websock_server.queue_message( + _messages.RunJavascriptMessage(source=source) + ) + self.server.initial_camera.position = self.CAMERA_POSITION + self.server.initial_camera.look_at = self.CAMERA_LOOK_AT + self.server.gui.configure_theme(control_width="medium", show_share_button=False) + self.main_panel = self.server.gui.main_panel + self.main_panel.dock_right() + self._build_gui(args.side) + self._register_scene_events() + self.selection_gizmo = self.server.scene.add_transform_controls( + "/selection/translation", + scale=SELECTION_GIZMO_SCALE_M, + line_width=SELECTION_GIZMO_LINE_WIDTH, + disable_sliders=True, + disable_rotations=True, + depth_test=False, + visible=False, + ) + self.selection_gizmo.on_update(self._selection_gizmo_event) + self._sync_document() + if self.state.path: + self.info.value = "Metadata loaded." + if args.urdf: + self.urdf_path.value = str(Path(args.urdf).expanduser().resolve()) + self._load_urdf() + + @property + def document(self) -> Document: + return self.state.document + + @staticmethod + def _preview_from(document: Document) -> Document: + preview = document.clone() + preview.taxels = [] + return preview + + def _build_gui(self, hand_side: str) -> None: + gui = self.server.gui + self.general_panel = gui.add_panel() + self.general_tab = self.general_panel.add_tab("General") + with self.general_tab: + self.hand_side = gui.add_dropdown( + "Hand Side", ("left", "right"), initial_value=hand_side + ) + self._watch(self.hand_side, self._hand_side_changed) + self.metadata_path = gui.add_text("Metadata Path", "") + self._action_button("New", self._new) + self._action_button("Save", self._save) + self._action_button("Load", self._load) + with gui.add_folder("Mode"): + self.mode_edit = gui.add_checkbox("Edit", True) + self.mode_preview = gui.add_checkbox("Preview", False) + self._action_button("Reset Camera", self._reset_camera) + self.show_saved = gui.add_checkbox("Show Saved Binding Poses", False) + self._watch(self.mode_edit, lambda checked: self._set_mode("Edit", checked)) + self._watch( + self.mode_preview, lambda checked: self._set_mode("Preview", checked) + ) + self._watch(self.show_saved, lambda _: self._request_draw()) + self.general_panel.dock_above(self.main_panel) + + self.edit_tabs = gui.add_tab_group() + self.urdf_tab = self.edit_tabs.add_tab("URDF") + with self.urdf_tab: + self.urdf_path = gui.add_text("URDF Path", "") + self._action_button("Load URDF", self._load_urdf) + with gui.add_folder("URDF Alignment"): + self.urdf_position = self._numbers("Position", (0.0, 0.0, 0.0), 0.001) + self.urdf_rotation = self._sliders("Rotation", (0.0, 0.0, 0.0)) + self.urdf_scale_widget = gui.add_number( + "Scale", 1.0, min=1e-6, step=0.01 + ) + for widget in (*self.urdf_position, *self.urdf_rotation): + self._watch(widget, lambda _: self._update_urdf_alignment()) + self._watch( + self.urdf_scale_widget, lambda _: self._update_urdf_alignment() + ) + self.edit_joint_folder = gui.add_folder( + "URDF Joints", expand_by_default=False + ) + + self.segments_tab = self.edit_tabs.add_tab("Segments") + with self.segments_tab: + with gui.add_folder("Segments Alignment"): + self.segment_position = self._numbers( + "Position", (0.0, 0.0, 0.0), 0.001 + ) + self.segment_rotation = self._sliders("Rotation", (0.0, 0.0, 0.0)) + for widget in (*self.segment_position, *self.segment_rotation): + self._watch(widget, lambda _: self._update_segment_alignment()) + with gui.add_folder("Add Segment"): + self.new_segment_joint = gui.add_dropdown("Name", tuple(JOINTS)) + self.new_segment_position = self._numbers( + "Position", (0.0, 0.0, 0.0), 0.001 + ) + self._action_button("Add with Input", self._add_segment_input) + self.segment_click_button = self._action_button( + "Add to URDF Surface", self._pick_segment + ) + self.segment_row_folder = gui.add_folder( + "All Segments", expand_by_default=False + ) + + self.taxels_tab = self.edit_tabs.add_tab("Taxels") + with self.taxels_tab: + with gui.add_folder("Add Taxels"): + self.grid_columns = gui.add_number("Columns", 4, min=1, step=1) + self.grid_rows = gui.add_number("Rows", 4, min=1, step=1) + self.grid_width = gui.add_number( + "Area Width (m)", 0.03, min=0.0, step=0.001 + ) + self.grid_height = gui.add_number( + "Area Height (m)", 0.03, min=0.0, step=0.001 + ) + self.grid_center = self._numbers( + "Center Position", (0.0, 0.0, 0.0), 0.001 + ) + self.grid_rotation = self._sliders("Rotation", (0.0, 0.0, 0.0)) + gui.add_html('Binding') + options = self._taxel_segment_options() + segment_fields = [ + self.server.gui.add_dropdown( + f"Segment {slot}", options, initial_value="None" + ) + for slot in range(4) + ] + for slot, segment in enumerate(segment_fields): + weight = self._weight_widget(0.0, True, False, slot) + self._watch( + segment, + lambda value, slot=slot: self._set_new_taxel_segment( + slot, value + ), + ) + self._watch( + weight, + lambda value, slot=slot: self._set_new_taxel_weight( + slot, value + ), + ) + self.new_taxel_binding_widgets.append((segment, weight)) + self.taxel_capture_button = self._action_button( + "Capture URDF Surface Pose", self._capture_taxel_pose + ) + self.segment_capture_button = self._action_button( + "Capture Segment Pose", self._capture_segment_pose + ) + self._action_button("Add Taxels", self._add_taxels_input) + with gui.add_folder("Bind Taxels to Segments"): + self.binding_start = gui.add_button("Start Binding to Segments") + gui.add_html( + '
↓
' + ) + self.binding_next = gui.add_button("Select Segments", disabled=True) + gui.add_html( + '
↓
' + ) + self.binding_complete = gui.add_button("Complete", disabled=True) + self.binding_info = gui.add_markdown("") + self._button(self.binding_start, self._start_binding) + self._button(self.binding_next, self._select_binding_segments) + self._button(self.binding_complete, self._complete_binding) + self.taxel_row_folder = gui.add_folder( + "All Taxels", expand_by_default=False + ) + + self.preview_panel = gui.add_folder("Preview", visible=False) + with self.preview_panel: + with gui.add_folder("Binding Metadata Source"): + self.binding_source_editor = gui.add_checkbox("Editor", True) + self.binding_source_live = gui.add_checkbox( + "Live", False, disabled=True + ) + self.show_binding = gui.add_checkbox("Show Binding Poses", False) + self.collection = gui.add_text("Tactile Collection ID", "tactile") + self._action_button("Apply Segments", self._apply_preview_segments) + self._action_button("Apply Taxels", self._apply_preview_taxels) + self._action_button("Apply URDF", self._apply_preview_urdf) + self._action_button("Apply ALL", self._apply_preview_all) + gui.add_markdown( + "Apply ALL fully overrides the editor's segments, taxels, bindings, and URDF joint values." + ) + self._watch( + self.binding_source_editor, + lambda checked: self._set_binding_source("Editor", checked), + ) + self._watch( + self.binding_source_live, + lambda checked: self._set_binding_source("Live", checked), + ) + self._watch(self.show_binding, lambda _: self._request_draw()) + with gui.add_folder("Live"): + self.live_status = gui.add_markdown("Current Statue: Disconnected") + self.joint_enabled = gui.add_checkbox( + "Use JointState to Drive URDF", False + ) + self.joint_collection = gui.add_text( + "JointState Collection ID", "", disabled=True + ) + self.joint_mapping = gui.add_text( + "Joint Mappings JSON", "{}", multiline=True, disabled=True + ) + gui.add_markdown( + 'Format: `{"source_joint": ["urdf_joint", scale, shift]}`; ' + "`value = raw * scale + shift`." + ) + self._watch(self.joint_enabled, self._joint_enabled) + self.connect_button = self._action_button("Connect Live", self._connect) + gui.add_markdown("Live input overrides values in the Static section.") + self.static_folder = gui.add_folder("Static") + with self.static_folder: + self._action_button( + "Reset Values", + self._restore_preview, + ) + gui.add_markdown("The following values affect Preview only.") + self.preview_segment_folder = gui.add_folder( + "Segments", expand_by_default=False + ) + self.preview_joint_folder = gui.add_folder( + "URDF Joints", expand_by_default=False + ) + + self.info = gui.add_text("__Tactile Info", "") + self.preview_error = self.info + + self.selection_panel = gui.add_panel() + self.selection_tab = self.selection_panel.add_tab("Selected Object") + self.selection_panel.float(x=10, y=140, width=320, height=520) + + def _numbers(self, prefix: str, values, step: float): + return tuple( + self.server.gui.add_number(f"{prefix} {axis}", float(value), step=step) + for axis, value in zip("XYZ", values) + ) + + def _sliders(self, prefix: str, values): + return tuple( + self.server.gui.add_slider( + f"{prefix} {axis} (deg)", -180.0, 180.0, 1.0, float(value) + ) + for axis, value in zip("XYZ", values) + ) + + def _watch(self, widget, callback) -> None: + async def updated(event): + if event.client is not None: + value = copy.deepcopy(event.target.value) + self.events.put(lambda: self._run_action(callback, value)) + + widget.on_update(updated) + + def _button(self, widget, callback) -> None: + async def clicked(_): + self.events.put(lambda: self._run_action(callback)) + + widget.on_click(clicked) + + def _action_button(self, label: str, callback, *, order=None): + button = self.server.gui.add_button(label, order=order) + + async def clicked(_): + self.events.put(lambda: self._run_action(callback)) + + button.on_click(clicked) + return button + + def _update_interruption_listener(self) -> None: + if self.binding_phase is not None: + allowed = (self.binding_start, self.binding_next, self.binding_complete) + elif self.pick_mode is not None: + allowed = ( + { + "segment": self.segment_click_button, + "taxel_pose": self.taxel_capture_button, + "segment_pose": self.segment_capture_button, + }[self.pick_mode], + ) + else: + allowed = () + # Viser shares this list for accepted clicks and value changes. Replace + # it to preserve in-flight dispatches; prepend to queue interruption first. + for widget in self.server.gui._gui_input_handle_from_uuid.values(): + callbacks = [ + callback + for callback in widget._impl.update_cb + if callback != self._procedure_interaction + ] + if allowed and all(widget is not button for button in allowed): + callbacks.insert(0, self._procedure_interaction) + widget._impl.update_cb = callbacks + + async def _procedure_interaction(self, event) -> None: + if event.client is not None and not event.target.disabled: + self.events.put(self._interrupt_procedure) + + def _interrupt_procedure(self) -> None: + if self.binding_phase is not None: + self._cancel_binding(report=True) + elif self.pick_mode is not None: + self._toggle_pick_mode(self.pick_mode) + else: + return + self.redraw = True + + def _run_action(self, callback, value=None) -> None: + self.info.value = "" + try: + callback() if value is None else callback(value) + except Exception as exc: + LOG.exception("Editor action failed") + self.info.value = ( + f'Error: {html.escape(str(exc))}' + ) + else: + message = self.info.value + if self.mode_value == "Preview": + self._update_preview_error() + errors = self.info.value + else: + errors = self._document_error_html() + if message and message != errors: + self.info.value = f"{message}
{errors}" if errors else message + elif errors: + self.info.value = errors + self.redraw = True + + def _request_draw(self) -> None: + self.redraw = True + + @property + def mode_value(self) -> str: + return "Edit" if self.mode_edit.value else "Preview" + + def _set_mode(self, mode: str, checked: bool) -> None: + selected = self.mode_edit if mode == "Edit" else self.mode_preview + other = self.mode_preview if mode == "Edit" else self.mode_edit + if not checked: + if not other.value: + selected.value = True + return + other.value = False + self.edit_tabs.visible = mode == "Edit" + self.preview_panel.visible = mode == "Preview" + self._set_pick_mode(None) + self.selected_kind, self.selected_rows = None, [] + self._cancel_binding() + self._sync_selection() + if mode == "Preview": + self._update_preview_error() + + def _hand_side_changed(self, _: str) -> None: + self._reset_live_metadata() + self._disconnect() + + def _reset_live_metadata(self) -> None: + self.live_document, self.live_metadata_key = None, "" + self.binding_source_editor.value = True + self.binding_source_live.value = False + self.binding_source_live.disabled = True + + def _sync_document(self) -> None: + self._reset_live_metadata() + self._set_pick_mode(None) + self.metadata_path.value = self.state.path + position, rotation = geo.components( + np.asarray(self.document.alignment, dtype=float) + ) + self._set_values(self.segment_position, position) + self._set_values(self.segment_rotation, rotation) + urdf_position, urdf_rotation = geo.components(self.urdf_alignment) + self._set_values(self.urdf_position, urdf_position) + self._set_values(self.urdf_rotation, urdf_rotation) + self.urdf_scale_widget.value = self.urdf_scale + self.selected_kind, self.selected_rows = None, [] + self._cancel_binding() + self._rebuild_edit_lists() + self._rebuild_joint_controls() + self._rebuild_preview_segments() + self.redraw = True + + @staticmethod + def _set_values(widgets, values) -> None: + for widget, value in zip(widgets, values): + widget.value = float(value) + + def _new(self) -> None: + document = Document() + document.alignment = geo.VISUALIZATION_FROM_HAND.tolist() + document.segments = geo.default_segments(self.hand_side.value) + self.state.document = document + self.state.saved = None + self.state.path = "" + self.preview_document = self._preview_from(document) + self._sync_document() + self.info.value = "New metadata created." + + def _load(self) -> None: + path = self.metadata_path.value.strip() + if not path: + raise ValueError("Enter Metadata Path") + document = storage.load(path, self.document.alignment) + self.state.replace(document, path) + self.preview_document = self._preview_from(document) + self._sync_document() + self.info.value = "Metadata loaded." + LOG.info("Metadata loaded from %s", path) + + def _save(self) -> None: + path = self.metadata_path.value.strip() + if not path: + raise ValueError("Enter Metadata Path") + issues = self.document.validate() + if not issues: + self._write_save(path) + return + modal = self.server.gui.add_modal("Save Metadata") + with modal: + self.server.gui.add_markdown( + f"Metadata contain errors, are you sure to save? ({len(issues)} error(s))" + ) + save_button = self.server.gui.add_button("Continue") + cancel_button = self.server.gui.add_button("Cancel") + + async def save(_): + self.events.put( + lambda: self._run_action( + lambda: (self._write_save(path), modal.close()) + ) + ) + + async def cancel(_): + def cancel_save(): + modal.close() + self.info.value = "Save canceled." + + self.events.put(lambda: self._run_action(cancel_save)) + + save_button.on_click(save) + cancel_button.on_click(cancel) + + def _write_save(self, path: str) -> None: + storage.save(self.state, path) + self.metadata_path.value = self.state.path + self.info.value = "Metadata saved." + self.redraw = True + + def _reset_camera(self) -> None: + for client in self.server.get_clients().values(): + client.camera.position = self.CAMERA_POSITION + client.camera.look_at = self.CAMERA_LOOK_AT + + def _load_urdf(self) -> None: + path = self.urdf_path.value.strip() + if not path: + raise ValueError("Enter URDF Path") + if self.model is not None: + self.model.remove() + self.model = geo.Model(path) + self._rebuild_joint_controls() + self.info.value = ( + f"{self.info.value}
URDF loaded." if self.info.value else "URDF loaded." + ) + + def _update_urdf_alignment(self) -> None: + self.urdf_alignment = geo.transform( + [widget.value for widget in self.urdf_position], + [widget.value for widget in self.urdf_rotation], + ) + self.urdf_scale = float(self.urdf_scale_widget.value) + if not math.isfinite(self.urdf_scale) or self.urdf_scale <= 0: + raise ValueError("URDF Scale must be positive") + self.redraw = True + + def _update_segment_alignment(self) -> None: + alignment = geo.transform( + [widget.value for widget in self.segment_position], + [widget.value for widget in self.segment_rotation], + ) + delta = geo.realign(self.document, alignment) + saved = self.state.saved + if saved is not None: + geo.apply_delta(saved.segments, list(range(len(saved.segments))), delta) + geo.apply_delta(saved.taxels, list(range(len(saved.taxels))), delta) + saved.alignment = alignment.tolist() + preview = self.preview_document + geo.apply_delta(preview.segments, list(range(len(preview.segments))), delta) + preview.alignment = alignment.tolist() + if self.live_document is not None: + geo.apply_delta( + self.live_document.segments, + list(range(len(self.live_document.segments))), + delta, + ) + geo.apply_delta( + self.live_document.taxels, + list(range(len(self.live_document.taxels))), + delta, + ) + self.live_document.alignment = alignment.tolist() + for segment, (_, positions, rotations) in zip( + self.document.segments, self.segment_widgets + ): + position, rotation = geo.components(geo.matrix(segment.pose)) + self._set_values(positions, position) + self._set_values(rotations, rotation) + for row, (positions, rotations, _) in self.taxel_widgets.items(): + taxel = self.document.taxels[row] + position, rotation = geo.components(geo.matrix(taxel.pose)) + self._set_values(positions, position) + self._set_values(rotations, rotation) + self._sync_preview_values() + values = self._selection_values() + rows = [row for row in self.selected_rows if 0 <= row < len(values)] + if rows: + center = np.asarray([values[row].pose[:3] for row in rows]).mean(axis=0) + self._set_values(self.selection_position_widgets, center) + if len(rows) == 1: + _, rotation = geo.components(geo.matrix(values[rows[0]].pose)) + self._set_values(self.selection_rotation_widgets, rotation) + self.selection_rotation = np.asarray(rotation, dtype=float) + self._update_preview_error() + self.redraw = True + + def _rebuild_joint_controls(self) -> None: + if self.edit_joint_folder is not None: + self.edit_joint_folder.remove() + if self.preview_joint_folder is not None: + self.preview_joint_folder.remove() + self.edit_joint_widgets, self.preview_joint_widgets = {}, {} + with self.urdf_tab: + self.edit_joint_folder = self.server.gui.add_folder( + "URDF Joints", expand_by_default=False + ) + with self.edit_joint_folder: + self._add_joint_sliders(self.document, self.edit_joint_widgets) + with self.static_folder: + self.preview_joint_folder = self.server.gui.add_folder( + "URDF Joints", expand_by_default=False + ) + with self.preview_joint_folder: + self._add_joint_sliders( + self.preview_document, self.preview_joint_widgets + ) + + def _add_joint_sliders(self, document: Document, output: dict) -> None: + if self.model is None: + self.server.gui.add_markdown("Load a URDF to edit joints.") + return + for index, joint in enumerate(self.model.joints): + lower = ( + joint.limit.lower + if joint.limit and joint.limit.lower is not None + else -math.pi + ) + upper = ( + joint.limit.upper + if joint.limit and joint.limit.upper is not None + else math.pi + ) + value = min(upper, max(lower, document.q.get(joint.name, 0.0))) + document.q[joint.name] = float(value) + widget = self.server.gui.add_slider( + f"{index}: {joint.name}", + float(lower), + float(upper), + 0.001, + float(value), + ) + self._watch( + widget, + lambda current, name=joint.name, document=document: ( + document.q.__setitem__(name, float(current)) + ), + ) + output[joint.name] = widget + + @staticmethod + def _sync_joint_widget_values(document: Document, widgets: dict) -> None: + for name, widget in widgets.items(): + value = float(document.q.get(name, 0.0)) + widget.value = min(float(widget.max), max(float(widget.min), value)) + + def _add_segment_input(self) -> None: + pose = IDENTITY_POSE.copy() + pose[:3] = [float(widget.value) for widget in self.new_segment_position] + row = self.document.add_segment(self.new_segment_joint.value, pose) + self.selected_kind, self.selected_rows = "segments", [row] + self._rebuild_edit_lists() + self.info.value = "Segment added." + + def _pick_segment(self) -> None: + self._toggle_pick_mode("segment") + + def _toggle_pick_mode(self, mode: str) -> None: + if self.pick_mode == mode: + self._set_pick_mode(None) + self.info.value = { + "segment": "Add Segment canceled.", + "taxel_pose": "Capture URDF Surface Pose canceled.", + "segment_pose": "Capture Segment Pose canceled.", + }[mode] + LOG.warning("%s", self.info.value) + return + self._set_pick_mode(mode) + + def _set_pick_mode(self, mode: str | None) -> None: + self.pick_mode = None if self.pick_mode == mode else mode + self.segment_click_button.label = ( + "● Add to URDF Surface" + if self.pick_mode == "segment" + else "Add to URDF Surface" + ) + self.taxel_capture_button.label = ( + "● Capture URDF Surface Pose" + if self.pick_mode == "taxel_pose" + else "Capture URDF Surface Pose" + ) + self.segment_capture_button.label = ( + "● Capture Segment Pose" + if self.pick_mode == "segment_pose" + else "Capture Segment Pose" + ) + if self.pick_mode is not None: + self.selected_kind, self.selected_rows = None, [] + self._sync_selection() + self._update_interruption_listener() + + def _grid_values(self): + center = [float(widget.value) for widget in self.grid_center] + return ( + center, + float(self.grid_width.value), + float(self.grid_height.value), + int(self.grid_rows.value), + int(self.grid_columns.value), + geo.transform(center, [widget.value for widget in self.grid_rotation]), + ) + + def _sync_new_taxel_binding_widgets(self) -> None: + options = self._taxel_segment_options() + for slot, (segment, weight) in enumerate(self.new_taxel_binding_widgets): + selected = self.new_taxel_binding.segments[slot] + if selected not in options: + selected = None + self.new_taxel_binding.segments[slot] = None + self.new_taxel_binding.weights[slot] = 0.0 + segment.options = options + segment.value = selected or "None" + weight.value = self.new_taxel_binding.weights[slot] + weight.disabled = selected is None + + def _set_new_taxel_segment(self, slot: int, value: str) -> None: + first = value != "None" and not any(self.new_taxel_binding.segments) + self.new_taxel_binding.segments[slot] = None if value == "None" else value + if self.new_taxel_binding.segments[slot] is None: + self.new_taxel_binding.weights[slot] = 0.0 + elif first: + self.new_taxel_binding.weights[slot] = 1.0 + self._sync_new_taxel_binding_widgets() + + def _set_new_taxel_weight(self, slot: int, value) -> None: + try: + Document.set_taxel_weight(self.new_taxel_binding, slot, float(value)) + finally: + self._sync_new_taxel_binding_widgets() + + def _add_taxels_input(self) -> None: + center, width, height, rows, columns, frame = self._grid_values() + selected = self.document.add_taxels( + geo.grid_poses(center, width, height, rows, columns, frame) + ) + for row in selected: + self.document.taxels[row].segments = copy.deepcopy( + self.new_taxel_binding.segments + ) + self.document.taxels[row].weights = copy.deepcopy( + self.new_taxel_binding.weights + ) + self.selected_kind, self.selected_rows = "taxels", selected + if selected: + self.taxel_page = selected[0] // TAXEL_PAGE_SIZE + self._rebuild_edit_lists() + noun = "taxel" if len(selected) == 1 else "taxels" + self.info.value = f"{len(selected)} {noun} added." + + def _capture_taxel_pose(self) -> None: + self._toggle_pick_mode("taxel_pose") + + def _capture_segment_pose(self) -> None: + self._toggle_pick_mode("segment_pose") + + def _rebuild_edit_lists(self) -> None: + self._rebuild_segment_list() + self._sync_taxel_page() + self._sync_new_taxel_binding_widgets() + self._sync_selection() + self._update_preview_error() + self.redraw = True + + def _rebuild_segment_list(self) -> None: + if self.segment_row_folder is not None: + self.segment_row_folder.remove() + self.segment_widgets = [] + self.segment_item_folders = [] + self.segment_error_widgets = [] + issues = self.document.validate() + with self.segments_tab: + self.segment_row_folder = self.server.gui.add_folder( + "All Segments", expand_by_default=False + ) + with self.segment_row_folder: + for row, segment in enumerate(self.document.segments): + row_issues = [ + issue + for issue in issues + if issue["kind"] == "segments" and issue["row"] == row + ] + folder = self.server.gui.add_folder( + f"{row}: {segment.joint}", expand_by_default=False + ) + with folder: + joint = self.server.gui.add_dropdown( + "⚠ Name" + if any(issue["field"] == "joint" for issue in row_issues) + else "Name", + tuple(JOINTS), + initial_value=segment.joint, + ) + position, rotation = geo.components(geo.matrix(segment.pose)) + positions = self._numbers("Position", position, 0.001) + rotations = self._sliders("Rotation", rotation) + self._watch( + joint, + lambda value, row=row: self._set_segment_joint(row, value), + ) + for axis, widget in enumerate(positions): + self._watch( + widget, + lambda value, row=row, axis=axis: ( + self._set_pose_component( + "segments", row, axis, value + ) + ), + ) + for axis, widget in enumerate(rotations): + self._watch( + widget, + lambda value, row=row, axis=axis: ( + self._set_pose_rotation( + "segments", row, axis, value + ) + ), + ) + messages = [issue["message"] for issue in row_issues] + error = self.server.gui.add_html( + '
' + + "
".join(map(html.escape, messages)) + + "
", + visible=bool(messages), + ) + self.segment_item_folders.append(folder) + self.segment_error_widgets.append(error) + self.segment_widgets.append((joint, positions, rotations)) + + def _taxel_segment_options(self) -> tuple[str, ...]: + # Dropdown values must stay unique even when the draft contains duplicate joints. + return ( + "None", + *dict.fromkeys(segment.joint for segment in self.document.segments), + ) + + def _taxel_page_count(self) -> int: + return max(1, math.ceil(len(self.document.taxels) / TAXEL_PAGE_SIZE)) + + def _change_taxel_page(self, delta: int) -> None: + page = min( + self._taxel_page_count() - 1, + max(0, self.taxel_page + delta), + ) + if page == self.taxel_page: + return + self.taxel_page = page + self._sync_taxel_page() + + def _jump_taxel_page(self, value) -> None: + page = min( + self._taxel_page_count() - 1, + max(0, int(value) - 1), + ) + if page == self.taxel_page: + return + self.taxel_page = page + self._sync_taxel_page() + + def _add_taxel_page_controls(self, page_count: int, location: str) -> None: + jump_page = self.server.gui.add_number( + f"__Taxel Page Jump {location}", + self.taxel_page + 1, + min=1, + max=page_count, + step=1, + ) + self._action_button( + "Jump", + lambda jump_page=jump_page: self._jump_taxel_page(jump_page.value), + ) + previous = self._action_button("Last Page", lambda: self._change_taxel_page(-1)) + next_page = self._action_button("Next Page", lambda: self._change_taxel_page(1)) + self.taxel_page_jumps.append(jump_page) + self.taxel_previous_buttons.append(previous) + self.taxel_next_buttons.append(next_page) + + def _build_taxel_page(self) -> None: + self.taxel_widgets = {} + self.taxel_error_widgets = {} + self.taxel_page_rows = [] + self.taxel_page_jumps = [] + self.taxel_previous_buttons = [] + self.taxel_next_buttons = [] + issues = self.document.validate() + options = self._taxel_segment_options() + page_count = self._taxel_page_count() + self.taxel_page = min(self.taxel_page, page_count - 1) + first = self.taxel_page * TAXEL_PAGE_SIZE + last = min(first + TAXEL_PAGE_SIZE, len(self.document.taxels)) + with self.taxel_row_folder: + self.taxel_page_container = self.server.gui.add_folder(None) + with self.taxel_page_container: + shown = f"{first}-{last - 1}" if last > first else "None" + self.taxel_page_info = self.server.gui.add_markdown( + f"Page **{self.taxel_page + 1} / {page_count}**" + f"  **{TAXEL_PAGE_SIZE} per Page** \n" + f"Taxels: **{shown}** of **{len(self.document.taxels)}**" + ) + self._add_taxel_page_controls(page_count, "top") + for page_slot in range(TAXEL_PAGE_SIZE): + row = first + page_slot + visible = row < last + taxel = self.document.taxels[row] if visible else Taxel() + row_issues = ( + [ + issue + for issue in issues + if issue["kind"] == "taxels" and issue["row"] == row + ] + if visible + else [] + ) + folder = self.server.gui.add_folder( + f"Taxel {row}", expand_by_default=False, visible=visible + ) + with folder: + position, rotation = geo.components(geo.matrix(taxel.pose)) + positions = self._numbers("Position", position, 0.001) + rotations = self._sliders("Rotation", rotation) + for axis, widget in enumerate(positions): + self._watch( + widget, + lambda value, page_slot=page_slot, axis=axis: ( + self._set_paged_taxel_pose_component( + page_slot, axis, value + ) + ), + ) + for axis, widget in enumerate(rotations): + self._watch( + widget, + lambda value, page_slot=page_slot, axis=axis: ( + self._set_paged_taxel_pose_rotation( + page_slot, axis, value + ) + ), + ) + self.server.gui.add_html( + 'Binding' + ) + slots = [] + segment_fields = [] + for slot in range(4): + selected = ( + taxel.segments[slot] + if taxel.segments[slot] in options + else "None" + ) + segment = self.server.gui.add_dropdown( + f"⚠ Segment {slot}" + if any( + issue["field"] in (f"segment{slot}", "binding") + for issue in row_issues + ) + else f"Segment {slot}", + options, + initial_value=selected, + ) + segment_fields.append((segment, selected)) + for slot, (segment, selected) in enumerate(segment_fields): + raw_weight = taxel.weights[slot] + display_weight = ( + raw_weight + if type(raw_weight) in (int, float) + and math.isfinite(raw_weight) + else 0.0 + ) + weight = self._weight_widget( + display_weight, + selected == "None", + any( + issue["field"] in (f"weight{slot}", "weights") + for issue in row_issues + ), + slot, + ) + self._watch( + segment, + lambda value, page_slot=page_slot, slot=slot: ( + self._set_paged_taxel_segment( + page_slot, slot, value + ) + ), + ) + self._watch( + weight, + lambda value, page_slot=page_slot, slot=slot: ( + self._set_paged_taxel_weight(page_slot, slot, value) + ), + ) + slots.append((segment, weight)) + messages = [issue["message"] for issue in row_issues] + error = self.server.gui.add_html( + '
' + + "
".join(map(html.escape, dict.fromkeys(messages))) + + "
", + visible=bool(messages), + ) + self.taxel_page_rows.append( + (folder, positions, rotations, slots, error) + ) + if visible: + self.taxel_error_widgets[row] = error + self.taxel_widgets[row] = (positions, rotations, slots) + self._add_taxel_page_controls(page_count, "bottom") + + def _paged_taxel_row(self, page_slot: int) -> int | None: + row = self.taxel_page * TAXEL_PAGE_SIZE + page_slot + return row if 0 <= row < len(self.document.taxels) else None + + def _set_paged_taxel_pose_component(self, page_slot: int, axis: int, value) -> None: + row = self._paged_taxel_row(page_slot) + if row is not None: + self._set_pose_component("taxels", row, axis, value) + + def _set_paged_taxel_pose_rotation(self, page_slot: int, axis: int, value) -> None: + row = self._paged_taxel_row(page_slot) + if row is not None: + self._set_pose_rotation("taxels", row, axis, value) + + def _set_paged_taxel_segment(self, page_slot: int, slot: int, value: str) -> None: + row = self._paged_taxel_row(page_slot) + if row is not None: + self._set_taxel_segment(row, slot, value) + + def _set_paged_taxel_weight(self, page_slot: int, slot: int, value) -> None: + row = self._paged_taxel_row(page_slot) + if row is not None: + self._set_taxel_weight(row, slot, value) + + def _sync_taxel_page(self) -> None: + page_count = self._taxel_page_count() + self.taxel_page = min(self.taxel_page, page_count - 1) + if self.taxel_page_container is None: + self._build_taxel_page() + first = self.taxel_page * TAXEL_PAGE_SIZE + last = min(first + TAXEL_PAGE_SIZE, len(self.document.taxels)) + shown = f"{first}-{last - 1}" if last > first else "None" + self.taxel_page_info.content = ( + f"Page **{self.taxel_page + 1} / {page_count}**" + f"  **{TAXEL_PAGE_SIZE} per Page** \n" + f"Taxels: **{shown}** of **{len(self.document.taxels)}**" + ) + for jump in self.taxel_page_jumps: + jump.min = 1 + jump.max = page_count + jump.value = self.taxel_page + 1 + for button in self.taxel_previous_buttons: + button.disabled = self.taxel_page == 0 + for button in self.taxel_next_buttons: + button.disabled = self.taxel_page >= page_count - 1 + + self.taxel_widgets = {} + self.taxel_error_widgets = {} + for page_slot, (folder, positions, rotations, slots, error) in enumerate( + self.taxel_page_rows + ): + row = first + page_slot + visible = row < last + folder.visible = visible + if not visible: + continue + folder.label = f"Taxel {row}" + taxel = self.document.taxels[row] + position, rotation = geo.components(geo.matrix(taxel.pose)) + self._set_values(positions, position) + self._set_values(rotations, rotation) + self.taxel_widgets[row] = (positions, rotations, slots) + self.taxel_error_widgets[row] = error + self._sync_taxel_binding_widgets(range(first, last)) + + def _weight_widget(self, value, disabled: bool, invalid: bool, slot: int): + label = f"{'⚠ ' if invalid else ''}Weight {slot}" + if type(value) in (int, float) and math.isfinite(value): + limits = {"min": 0.0, "max": 1.0} if 0.0 <= value <= 1.0 else {} + return self.server.gui.add_number( + label, float(value), step=0.01, disabled=disabled, **limits + ) + return self.server.gui.add_text(label, str(value), disabled=disabled) + + @staticmethod + def _set_error_widget(widget, messages) -> None: + widget.content = ( + '
' + + "
".join(map(html.escape, dict.fromkeys(messages))) + + "
" + ) + widget.visible = bool(messages) + + def _sync_segment_validation_widgets(self) -> None: + issues = self.document.validate() + for row, segment in enumerate(self.document.segments): + row_issues = [ + issue + for issue in issues + if issue["kind"] == "segments" and issue["row"] == row + ] + self.segment_item_folders[row].label = f"{row}: {segment.joint}" + self.segment_widgets[row][0].label = ( + "⚠ Name" + if any(issue["field"] == "joint" for issue in row_issues) + else "Name" + ) + self._set_error_widget( + self.segment_error_widgets[row], + [issue["message"] for issue in row_issues], + ) + + def _sync_taxel_binding_widgets(self, rows) -> None: + issues = self.document.validate() + options = self._taxel_segment_options() + for row in dict.fromkeys(rows): + if not 0 <= row < len(self.document.taxels): + continue + taxel = self.document.taxels[row] + widgets = self.taxel_widgets.get(row) + if widgets is None: + continue + row_issues = [ + issue + for issue in issues + if issue["kind"] == "taxels" and issue["row"] == row + ] + for slot, (segment, weight) in enumerate(widgets[2]): + selected = ( + taxel.segments[slot] if taxel.segments[slot] in options else "None" + ) + segment.options = options + segment.value = selected + segment.label = ( + f"⚠ Segment {slot}" + if any( + issue["field"] in (f"segment{slot}", "binding") + for issue in row_issues + ) + else f"Segment {slot}" + ) + raw_weight = taxel.weights[slot] + weight.value = ( + float(raw_weight) + if type(raw_weight) in (int, float) and math.isfinite(raw_weight) + else 0.0 + ) + weight.disabled = selected == "None" + weight.label = ( + f"⚠ Weight {slot}" + if any( + issue["field"] in (f"weight{slot}", "weights") + for issue in row_issues + ) + else f"Weight {slot}" + ) + self._set_error_widget( + self.taxel_error_widgets[row], + [issue["message"] for issue in row_issues], + ) + if self.mode_value == "Preview": + self._update_preview_error() + self.redraw = True + + def _set_segment_joint(self, row: int, value: str) -> None: + self.document.segments[row].joint = value + self._sync_segment_validation_widgets() + self._sync_taxel_binding_widgets(range(len(self.document.taxels))) + self._sync_new_taxel_binding_widgets() + self._sync_selection() + + def _values(self, kind: str): + return self.document.segments if kind == "segments" else self.document.taxels + + def _set_pose_component(self, kind: str, row: int, axis: int, value) -> None: + self._values(kind)[row].pose[axis] = float(value) + if ( + self.mode_value == "Edit" + and self.selected_kind == kind + and row in self.selected_rows + ): + self._sync_selection() + self.redraw = True + + def _set_pose_rotation(self, kind: str, row: int, axis: int, value) -> None: + current = geo.matrix(self._values(kind)[row].pose) + position, rotation = geo.components(current) + rotation = list(rotation) + rotation[axis] = float(value) + self._values(kind)[row].pose = geo.pose(geo.transform(position, rotation)) + if ( + self.mode_value == "Edit" + and self.selected_kind == kind + and row in self.selected_rows + ): + self._sync_selection() + self.redraw = True + + def _set_taxel_segment(self, row: int, slot: int, value: str) -> None: + taxel = self.document.taxels[row] + first = value != "None" and not any(taxel.segments) + taxel.set_segment(slot, None if value == "None" else value) + if taxel.segments[slot] is None: + taxel.weights[slot] = 0.0 + elif first: + taxel.weights[slot] = 1.0 + self._sync_taxel_binding_widgets([row]) + self._sync_selection() + + def _set_taxel_weight(self, row: int, slot: int, value) -> None: + try: + self.document.set_weight(row, slot, float(value)) + finally: + self._sync_taxel_binding_widgets([row]) + self._sync_selection() + + def _start_binding(self) -> None: + self._set_pick_mode(None) + self.binding_phase = "taxels" + self.binding_taxels = ( + self.selected_rows.copy() if self.selected_kind == "taxels" else [] + ) + self.binding_segments = [] + self.selected_kind, self.selected_rows = None, [] + self._sync_selection() + self._set_binding_button("next" if self.binding_taxels else None) + self.binding_info.content = ( + "Select taxel(s) to bind in 3D scene. Hold Shift to multi-select. " + 'Click "Select Segments" after selecting taxels. Click another area to cancel.' + ) + self.info.value = "Binding started." + self._update_interruption_listener() + + def _select_binding_segments(self) -> None: + if not self.binding_taxels: + raise ValueError("Select at least one taxel") + self.binding_phase = "segments" + self._set_binding_button(None) + self.binding_info.content = ( + 'Select 1-4 segment(s). Hold Shift to multi-select. Click "Complete" ' + "after selecting segments. Click another area to cancel." + ) + self.info.value = "Taxels selected. Select segments." + + def _complete_binding(self) -> None: + if not 1 <= len(self.binding_segments) <= 4: + self._cancel_binding() + self.info.value = ( + 'Too many segments selected. ' + 'Select 1-4 segments and click "Start Binding to Segments" ' + "to start over." + ) + return + rows = list(self.binding_taxels) + self.document.assign(rows, self.binding_segments) + self._cancel_binding() + self.selected_kind, self.selected_rows = "taxels", rows + self._sync_taxel_binding_widgets(rows) + self._sync_selection() + noun = "taxel" if len(rows) == 1 else "taxels" + self.info.value = f"Binding completed for {len(rows)} {noun}." + + def _set_binding_button(self, active: str | None) -> None: + self.binding_start.disabled = active != "start" + self.binding_next.disabled = active != "next" + self.binding_complete.disabled = active != "complete" + + def _cancel_binding(self, report: bool = False) -> None: + self.binding_phase = None + self.binding_taxels, self.binding_segments = [], [] + if hasattr(self, "binding_next"): + self._set_binding_button("start") + self.binding_info.content = "" + if report: + self.info.value = "Binding canceled." + LOG.warning("Binding canceled.") + self._update_interruption_listener() + + def _joint_enabled(self, enabled: bool) -> None: + self.joint_collection.disabled = not enabled + self.joint_mapping.disabled = not enabled + + def _connect(self) -> None: + if self.receiver is not None: + self._disconnect(report=True) + return + config = { + "side": self.hand_side.value, + "collection": self.collection.value, + "joint_state": self.joint_collection.value + if self.joint_enabled.value + else "", + } + self.receiver = sdk.Receiver(config) + self._reset_live_metadata() + self.connect_button.label = "Disconnect Live" + self.live_status.content = "Current Status: Connecting" + self.info.value = "Connecting Live." + + def _disconnect(self, report: bool = False) -> None: + if self.receiver is not None: + self.receiver.close() + self.receiver = None + self.connect_button.label = "Connect Live" + self.live_status.content = "Current Status: Disconnected" + self._update_preview_error() + if report: + self.info.value = "Live disconnected." + self.redraw = True + + @property + def binding_source_value(self) -> str: + return "Live" if self.binding_source_live.value else "Editor" + + def _set_binding_source(self, source: str, checked: bool) -> None: + selected = ( + self.binding_source_editor + if source == "Editor" + else self.binding_source_live + ) + other = ( + self.binding_source_live + if source == "Editor" + else self.binding_source_editor + ) + if not checked: + if not other.value: + selected.value = True + return + other.value = False + self._update_preview_error() + self.redraw = True + + def _binding_document(self) -> Document: + if self.binding_source_value == "Live": + if self.live_document is None: + raise ValueError("Live metadata is unavailable") + return self.live_document + return self.document + + def _preview_taxels(self) -> tuple[list[Taxel | None], list[str]]: + return geo.deform(self._binding_document(), self.preview_document.segments) + + def _rendered_preview_taxels( + self, + ) -> tuple[list[tuple[int, Taxel]], list[str], int]: + try: + rows, errors = self._preview_taxels() + except (ValueError, TypeError, KeyError, IndexError) as exc: + return [], [f"Binding error: {exc}"], 0 + return ( + [(row, taxel) for row, taxel in enumerate(rows) if taxel is not None], + errors, + len(rows), + ) + + def _update_preview_error(self) -> None: + _, errors, _ = self._rendered_preview_taxels() + self.preview_error.value = ( + '' + + "
".join(map(html.escape, errors)) + + "" + if errors + else "" + ) + + def _document_error_html(self) -> str: + messages = [] + labels = {"segments": "Segment", "taxels": "Taxel"} + for issue in self.document.validate(): + label = labels.get(issue["kind"], issue["kind"].title()) + location = f"{label} {issue['row']}" if issue["row"] is not None else label + messages.append(html.escape(f"{location}: {issue['message']}")) + return ( + '' + "
".join(messages) + "" + if messages + else "" + ) + + def _apply_preview_segments(self) -> None: + self.document.segments = copy.deepcopy(self.preview_document.segments) + self._rebuild_edit_lists() + self.info.value = "Preview segments applied to Edit data." + + def _applied_preview_taxels(self) -> list[Taxel]: + # Invalid rows keep their binding metadata so row order still matches forces. + rows, _ = self._preview_taxels() + sources = self._binding_document().taxels + return copy.deepcopy( + [ + source if row is None else replace(source, pose=row.pose) + for source, row in zip(sources, rows) + ] + ) + + def _apply_preview_taxels(self) -> None: + self.document.taxels = self._applied_preview_taxels() + self._rebuild_edit_lists() + self.info.value = "Preview taxels applied to Edit data." + + def _apply_preview_urdf(self) -> None: + self.document.q = copy.deepcopy(self.preview_document.q) + self._sync_joint_widget_values(self.document, self.edit_joint_widgets) + self.info.value = "Preview URDF joints applied to Edit data." + + def _apply_preview_all(self) -> None: + document = self.document + taxels = self._applied_preview_taxels() + document.segments = copy.deepcopy(self.preview_document.segments) + document.taxels = taxels + document.q = copy.deepcopy(self.preview_document.q) + self._rebuild_edit_lists() + self._sync_joint_widget_values(document, self.edit_joint_widgets) + self.info.value = "Preview data applied to Edit data." + + def _restore_preview(self) -> None: + preview = self.preview_document + replacement = self._preview_from(self.document) + old_joints = [segment.joint for segment in preview.segments] + preview.alignment = replacement.alignment + preview.segments = replacement.segments + preview.taxels = replacement.taxels + preview.q = replacement.q + if old_joints != [segment.joint for segment in preview.segments]: + self._rebuild_preview_segments() + if self.mode_value == "Preview": + self.selected_kind, self.selected_rows = None, [] + if self.mode_value == "Preview": + self._sync_selection() + self._sync_preview_values() + self._update_preview_error() + self.info.value = "Preview data reset from Edit data." + + def _rebuild_preview_segments(self) -> None: + if self.preview_segment_folder is not None: + self.preview_segment_folder.remove() + self.preview_segment_widgets = [] + with self.static_folder: + self.preview_segment_folder = self.server.gui.add_folder( + "Segments", expand_by_default=False + ) + with self.preview_segment_folder: + for row, segment in enumerate(self.preview_document.segments): + with self.server.gui.add_folder( + f"{row}: {segment.joint}", expand_by_default=False + ): + self.server.gui.add_markdown(f"HandJoint: `{segment.joint}`") + position, rotation = geo.components(geo.matrix(segment.pose)) + positions = self._numbers("Position", position, 0.001) + rotations = self._sliders("Rotation", rotation) + for axis, widget in enumerate(positions): + self._watch( + widget, + lambda value, row=row, axis=axis: ( + self._set_preview_pose(row, axis, value, False) + ), + ) + for axis, widget in enumerate(rotations): + self._watch( + widget, + lambda value, row=row, axis=axis: ( + self._set_preview_pose(row, axis, value, True) + ), + ) + self.preview_segment_widgets.append((positions, rotations)) + + def _set_preview_pose( + self, row: int, axis: int, value: float, rotation_field: bool + ) -> None: + segment = self.preview_document.segments[row] + if not rotation_field: + segment.pose[axis] = float(value) + else: + position, rotation = geo.components(geo.matrix(segment.pose)) + rotation = list(rotation) + rotation[axis] = float(value) + segment.pose = geo.pose(geo.transform(position, rotation)) + if self.mode_value == "Preview" and row in self.selected_rows: + self._sync_selection() + self._update_preview_error() + self.redraw = True + + def _update_live(self) -> None: + if self.receiver is None: + return + self.live_status.content = "Current Status: " + self.receiver.status + if not self.receiver.active: + return + frame = self.receiver.latest + try: + raw_metadata = frame.get("metadata") + if raw_metadata is not None: + key = json.dumps(raw_metadata, sort_keys=True) + if key != self.live_metadata_key: + self.live_document = storage.parse( + raw_metadata, self.document.alignment + ) + self.live_metadata_key = key + self.binding_source_live.disabled = False + hand = frame.get("hand") + updated = ( + geo.hand_segments( + hand, np.asarray(self.document.alignment, dtype=float) + ) + if hand is not None + else [] + ) + if updated: + joints = [segment.joint for segment in updated] + signature = [ + segment.joint for segment in self.preview_document.segments + ] + self.preview_document.segments = updated + if signature != joints: + self._rebuild_preview_segments() + if self.mode_value == "Preview": + self.selected_kind, self.selected_rows = None, [] + self._sync_selection() + else: + self.live_status.content = "Current Status: Data Invalid" + if self.joint_enabled.value: + mapping = json.loads(self.joint_mapping.value or "{}") + if not isinstance(mapping, dict): + raise ValueError("Joint Mappings JSON must be an object") + for source, rule in mapping.items(): + raw, valid = frame.get("joints", {}).get(source, (None, False)) + if valid and isinstance(rule, list) and len(rule) == 3: + target, scale, offset = rule + self.preview_document.q[target] = raw * scale + offset + self._sync_preview_values() + self._update_preview_error() + except (ValueError, TypeError, KeyError, IndexError) as exc: + self.preview_error.value = f'Live error: {html.escape(str(exc))}' + self.redraw = True + + def _sync_preview_values(self) -> None: + for row, widgets in enumerate(self.preview_segment_widgets): + if row >= len(self.preview_document.segments): + break + position, rotation = geo.components( + geo.matrix(self.preview_document.segments[row].pose) + ) + self._set_values(widgets[0], position) + self._set_values(widgets[1], rotation) + for name, widget in self.preview_joint_widgets.items(): + widget.value = self.preview_document.q.get(name, 0.0) + + async def _selection_gizmo_event(self, event) -> None: + if event.client is None: + return + phase = event.phase + position = np.asarray(event.target.position, dtype=float) + self.events.put(lambda: self._apply_selection_gizmo_event(phase, position)) + + def _apply_selection_gizmo_event(self, phase: str, position: np.ndarray) -> None: + values = self._selection_values() + rows = [row for row in self.selected_rows if 0 <= row < len(values)] + if not rows: + self.selection_gizmo_dragging = False + self.selection_gizmo_last_position = None + return + if phase == "start": + self.selection_gizmo_dragging = True + self.selection_gizmo_last_position = position + return + if ( + not self.selection_gizmo_dragging + or self.selection_gizmo_last_position is None + ): + return + delta = position - self.selection_gizmo_last_position + if np.any(delta): + transform = np.eye(4) + transform[:3, 3] = delta + geo.apply_delta(values, rows, transform) + center = np.asarray([values[row].pose[:3] for row in rows]).mean(axis=0) + self._set_values(self.selection_position_widgets, center) + self._sync_selected_list_values() + self.selection_gizmo_last_position = position + if phase == "end": + self.selection_gizmo_dragging = False + self.selection_gizmo_last_position = None + self.redraw = True + + def _register_scene_events(self) -> None: + self.server.scene.on_click()(self._scene_click) + self.server.scene.on_click(modifier="shift")(self._scene_click) + self.server.scene.on_rect_select(modifier="shift")(self._scene_box) + + async def _scene_click(self, event) -> None: + self.events.put( + lambda: self._run_action(lambda: self._handle_scene_event(event, False)) + ) + + async def _scene_box(self, event) -> None: + self.events.put( + lambda: self._run_action(lambda: self._handle_scene_event(event, True)) + ) + + def _display_values(self): + if self.mode_value == "Preview": + return self.preview_document.segments, [] + return self.document.segments, self.document.taxels + + def _handle_scene_event(self, event, box: bool) -> None: + if self.pick_mode is not None: + # Captures consume clicks; rectangles must not select editable objects. + if not box: + self._handle_pick(event) + return + segments, taxels = self._display_values() + candidates = { + "segments": list(enumerate(segments)), + "taxels": list(enumerate(taxels)), + } + if self.binding_phase is not None: + candidates = { + key: values + for key, values in candidates.items() + if key == self.binding_phase + } + chosen = { + kind: self._box_rows(event, values) + if box + else self._ray_rows(event, kind, values) + for kind, values in candidates.items() + } + chosen = {kind: rows for kind, rows in chosen.items() if rows} + if not chosen: + if self.binding_phase is not None: + self._cancel_binding(report=True) + else: + self.selected_kind, self.selected_rows = None, [] + self._sync_selection() + self.redraw = True + return + kind = ( + self.selected_kind + if self.selected_kind in chosen + else "taxels" + if "taxels" in chosen + else "segments" + ) + rows = chosen[kind] + if self.binding_phase == "taxels": + self.binding_taxels = ( + sorted(set(self.binding_taxels) | set(rows)) + if event.modifier == "shift" + else rows + ) + self._set_binding_button("next" if self.binding_taxels else None) + elif self.binding_phase == "segments": + self.binding_segments = ( + sorted(set(self.binding_segments) | set(rows)) + if event.modifier == "shift" + else rows + ) + self._set_binding_button("complete" if self.binding_segments else None) + elif event.modifier == "shift": + if self.selected_kind not in (None, kind): + return + self.selected_kind = kind + self.selected_rows = sorted(set(self.selected_rows) | set(rows)) + self._sync_selection() + elif self.selected_kind != kind or self.selected_rows != rows: + self.selected_kind, self.selected_rows = kind, rows + self._sync_selection() + self.redraw = True + + @staticmethod + def _ray_rows(event, kind: str, values) -> list[int]: + origin = np.asarray(event.ray_origin) + direction = np.asarray(event.ray_direction) + hits = [] + for row, value in values: + inverse = np.linalg.inv(geo.matrix(value.pose)) + local_origin = (inverse @ np.r_[origin, 1.0])[:3] + local_direction = inverse[:3, :3] @ direction + if kind == "segments": + distance = Editor._ray_box_distance( + local_origin, local_direction, SEGMENT_CUBE_SIDE_M / 2 + ) + if distance is not None: + hits.append((distance, row)) + continue + if abs(local_direction[2]) < 1e-10: + continue + distance = -local_origin[2] / local_direction[2] + local_hit = local_origin + distance * local_direction + if distance > 0 and np.max(np.abs(local_hit[:2])) <= TAXEL_SIDE_M / 2: + hits.append((distance, row)) + return [min(hits)[1]] if hits else [] + + @staticmethod + def _ray_box_distance(origin, direction, half_extent: float) -> float | None: + near, far = -math.inf, math.inf + for axis in range(3): + if abs(direction[axis]) < 1e-10: + if abs(origin[axis]) > half_extent: + return None + continue + first = (-half_extent - origin[axis]) / direction[axis] + second = (half_extent - origin[axis]) / direction[axis] + near = max(near, min(first, second)) + far = min(far, max(first, second)) + if near > far: + return None + if far <= 0: + return None + return near if near > 0 else far + + @staticmethod + def _box_rows(event, values) -> list[int]: + camera = event.client.camera + rotation = tf.quaternion_matrix(camera.wxyz)[:3, :3] + origin = np.asarray(camera.position) + rows = [] + for row, value in values: + point = np.asarray(value.pose[:3]) + local = rotation.T @ (point - origin) + if local[2] <= 0: + continue + uv = ( + 0.5 + + local[0] / local[2] / (2 * np.tan(camera.fov / 2) * camera.aspect), + 0.5 + local[1] / local[2] / (2 * np.tan(camera.fov / 2)), + ) + if all( + lower <= value <= upper + for lower, value, upper in zip(event.screen_min, uv, event.screen_max) + ): + rows.append(row) + return rows + + def _handle_pick(self, event) -> None: + mode = self.pick_mode + self._set_pick_mode(None) + if mode == "segment_pose": + rows = self._ray_rows( + event, "segments", list(enumerate(self.document.segments)) + ) + if not rows: + self.info.value = "Capture Segment Pose canceled." + return + position, rotation = geo.components( + geo.matrix(self.document.segments[rows[0]].pose) + ) + self._set_values(self.grid_center, position) + self._set_values(self.grid_rotation, rotation) + self.info.value = "Segment pose captured." + return + if self.model is None: + raise ValueError("Load a URDF before using URDF Click") + hit = self.model.hit( + np.asarray(event.ray_origin), + np.asarray(event.ray_direction), + self.document.q, + self.urdf_alignment, + self.urdf_scale, + ) + if hit is None: + self.info.value = ( + "Add Segment canceled." + if mode == "segment" + else "Capture URDF Surface Pose canceled." + ) + return + if mode == "segment": + value = IDENTITY_POSE.copy() + value[:3] = hit[1].tolist() + row = self.document.add_segment(self.new_segment_joint.value, value) + self.selected_kind, self.selected_rows = "segments", [row] + self._rebuild_edit_lists() + self.info.value = "Segment added." + else: + frame = geo.surface_frame(hit[1], hit[2]) + position, rotation = geo.components(frame) + self._set_values(self.grid_center, position) + self._set_values(self.grid_rotation, rotation) + self.info.value = "URDF surface pose captured." + + def _selection_values(self): + if self.selected_kind is None: + return [] + if self.mode_value == "Preview": + return self.preview_document.segments + return self._values(self.selected_kind) + + def _sync_selection(self) -> None: + values = self._selection_values() + rows = [row for row in self.selected_rows if 0 <= row < len(values)] + self.selected_rows = rows + if self.selection_folder is None: + with self.selection_tab: + self.selection_folder = self.server.gui.add_folder("Selection") + with self.selection_folder: + self.selection_summary = self.server.gui.add_markdown("") + self.selection_names = self.server.gui.add_text( + "__Selected Object Names", "", multiline=True, disabled=True + ) + self.selection_position_widgets = self._numbers( + "Position", np.zeros(3), 0.001 + ) + self.selection_rotation_widgets = self._sliders("Rotation", np.zeros(3)) + for axis, widget in enumerate(self.selection_position_widgets): + self._watch( + widget, + lambda value, axis=axis: self._set_selection_position( + axis, value + ), + ) + for axis, widget in enumerate(self.selection_rotation_widgets): + self._watch( + widget, + lambda value, axis=axis: self._set_selection_rotation( + axis, value + ), + ) + self.selection_segment_editor = self.server.gui.add_folder( + "Segment Editor", visible=False + ) + with self.selection_segment_editor: + self.selection_segment_joint = self.server.gui.add_dropdown( + "Name", tuple(JOINTS) + ) + self._watch( + self.selection_segment_joint, + self._set_selected_segment_joint, + ) + self.selection_taxel_editor = self.server.gui.add_folder( + "Taxel Editor", visible=False + ) + self.selection_taxel_binding_widgets = [] + with self.selection_taxel_editor: + self.server.gui.add_html( + 'Binding' + ) + options = self._taxel_segment_options() + segment_fields = [ + self.server.gui.add_dropdown( + f"Segment {slot}", options, initial_value="None" + ) + for slot in range(4) + ] + for slot, segment in enumerate(segment_fields): + weight = self._weight_widget(0.0, True, False, slot) + self._watch( + segment, + lambda value, slot=slot: self._set_selected_taxel_segment( + slot, value + ), + ) + self._watch( + weight, + lambda value, slot=slot: self._set_selected_taxel_weight( + slot, value + ), + ) + self.selection_taxel_binding_widgets.append((segment, weight)) + self.selection_scale_widget = self.server.gui.add_number( + "Scale", 1.0, min=1e-6, step=0.01 + ) + self._watch(self.selection_scale_widget, self._set_selection_scale) + self.selection_delete = self._action_button( + "Delete", self._delete_selection + ) + + self.selection_summary.content = ( + f"Type: **{self.selected_kind if rows else 'None'}** \n" + f"Count: **{len(rows)}**" + ) + names = [ + values[row].joint if self.selected_kind == "segments" else f"Taxel {row}" + for row in rows + ] + self.selection_names.value = "\n".join(names) + center = ( + np.asarray([values[row].pose[:3] for row in rows]).mean(axis=0) + if rows + else np.zeros(3) + ) + rotation = ( + geo.components(geo.matrix(values[rows[0]].pose))[1] + if len(rows) == 1 + else np.zeros(3) + ) + self._set_values(self.selection_position_widgets, center) + self._set_values(self.selection_rotation_widgets, rotation) + self.selection_rotation = np.asarray(rotation, dtype=float) + for widget in ( + *self.selection_position_widgets, + *self.selection_rotation_widgets, + ): + widget.disabled = not rows + single_segment = self.selected_kind == "segments" and len(rows) == 1 + single_taxel = self.selected_kind == "taxels" and len(rows) == 1 + self.selection_segment_editor.visible = single_segment + self.selection_taxel_editor.visible = single_taxel + if single_segment: + self.selection_segment_joint.value = values[rows[0]].joint + if single_taxel: + taxel = values[rows[0]] + options = self._taxel_segment_options() + for slot, (segment, weight) in enumerate( + self.selection_taxel_binding_widgets + ): + selected = ( + taxel.segments[slot] if taxel.segments[slot] in options else "None" + ) + segment.options = options + segment.value = selected + weight.value = taxel.weights[slot] + weight.disabled = selected == "None" + self.selection_scale = 1.0 + self.selection_scale_widget.value = 1.0 + self.selection_scale_widget.disabled = len(rows) < 2 + self.selection_delete.disabled = not rows + + def _set_selected_segment_joint(self, value: str) -> None: + if self.selected_kind != "segments" or len(self.selected_rows) != 1: + return + row = self.selected_rows[0] + if self.mode_value == "Edit": + self._set_segment_joint(row, value) + return + self.preview_document.segments[row].joint = value + self._rebuild_preview_segments() + self._sync_selection() + + def _set_selected_taxel_segment(self, slot: int, value: str) -> None: + if self.selected_kind != "taxels" or len(self.selected_rows) != 1: + return + self._set_taxel_segment(self.selected_rows[0], slot, value) + + def _set_selected_taxel_weight(self, slot: int, value) -> None: + if self.selected_kind != "taxels" or len(self.selected_rows) != 1: + return + self._set_taxel_weight(self.selected_rows[0], slot, value) + + def _set_selection_position(self, axis: int, value: float) -> None: + values = self._selection_values() + center = np.asarray([values[row].pose[:3] for row in self.selected_rows]).mean( + axis=0 + ) + delta = np.eye(4) + delta[axis, 3] = float(value) - center[axis] + geo.apply_delta(values, self.selected_rows, delta) + self._sync_selected_list_values() + + def _set_selection_rotation(self, axis: int, value: float) -> None: + values = self._selection_values() + if len(self.selected_rows) == 1: + row = self.selected_rows[0] + position, rotation = geo.components(geo.matrix(values[row].pose)) + rotation = list(rotation) + rotation[axis] = float(value) + values[row].pose = geo.pose(geo.transform(position, rotation)) + else: + points = np.asarray([values[row].pose[:3] for row in self.selected_rows]) + center = points.mean(axis=0) + angle = np.deg2rad(float(value) - self.selection_rotation[axis]) + rotation = tf.rotation_matrix(angle, np.eye(3)[axis], center) + geo.apply_delta(values, self.selected_rows, rotation) + self.selection_rotation[axis] = float(value) + self._sync_selected_list_values() + + def _set_selection_scale(self, value: float) -> None: + values = self._selection_values() + factor = float(value) / self.selection_scale + geo.scale_positions(values, self.selected_rows, factor) + self.selection_scale = float(value) + self._sync_selected_list_values() + + def _sync_selected_list_values(self) -> None: + if self.mode_value == "Edit": + widgets = ( + self.segment_widgets + if self.selected_kind == "segments" + else self.taxel_widgets + ) + values = self._selection_values() + for row in self.selected_rows: + if self.selected_kind == "segments": + if row >= len(widgets): + continue + row_widgets = widgets[row] + else: + row_widgets = widgets.get(row) + if row_widgets is None: + continue + position, rotation = geo.components(geo.matrix(values[row].pose)) + position_widgets, rotation_widgets = ( + (row_widgets[1], row_widgets[2]) + if self.selected_kind == "segments" + else (row_widgets[0], row_widgets[1]) + ) + self._set_values(position_widgets, position) + self._set_values(rotation_widgets, rotation) + else: + self._sync_preview_values() + self.redraw = True + + def _delete_selection(self) -> None: + if self.mode_value == "Preview": + for row in sorted(self.selected_rows, reverse=True): + self.preview_document.segments.pop(row) + self._rebuild_preview_segments() + self._update_preview_error() + else: + self.document.delete(self.selected_kind, self.selected_rows) + self.selected_kind, self.selected_rows = None, [] + self._rebuild_edit_lists() + return + self.selected_kind, self.selected_rows = None, [] + self._sync_selection() + + def _draw_instances( + self, + name: str, + kind: str, + values, + colors, + opacity: float, + ) -> None: + matrices = [geo.matrix(value.pose) for value in values] + if not matrices: + if name in self.scene_handles: + self.scene_handles[name].visible = False + return + positions = np.asarray([value[:3, 3] for value in matrices]) + rotations = np.asarray([tf.quaternion_from_matrix(value) for value in matrices]) + scale = SEGMENT_CUBE_SIDE_M if kind == "segments" else TAXEL_SIDE_M + if name not in self.scene_handles: + mesh = ( + trimesh.creation.box(extents=(1.0, 1.0, 1.0)) + if kind == "segments" + else trimesh.Trimesh( + vertices=[ + [-0.5, -0.5, 0.0], + [0.5, -0.5, 0.0], + [0.5, 0.5, 0.0], + [-0.5, 0.5, 0.0], + ], + faces=[[0, 1, 2], [0, 2, 3]], + process=False, + ) + ) + self.scene_handles[name] = self.server.scene.add_batched_meshes_simple( + name, + mesh.vertices, + mesh.faces, + rotations, + positions, + batched_scales=np.full(len(values), scale), + batched_colors=np.asarray(colors, dtype=np.uint8), + opacity=opacity, + flat_shading=SEGMENT_FLAT_SHADING if kind == "segments" else False, + side="double", + cast_shadow=VISUALIZER_CAST_SHADOW, + receive_shadow=VISUALIZER_RECEIVE_SHADOW, + ) + handle = self.scene_handles[name] + handle.visible = True + handle.batched_positions = positions + handle.batched_wxyzs = rotations + handle.batched_scales = np.full(len(values), scale) + handle.batched_colors = np.asarray(colors, dtype=np.uint8) + handle.opacity = opacity + + def _hide_prefix(self, prefix: str) -> None: + for name, handle in self.scene_handles.items(): + if name.startswith(prefix): + handle.visible = False + + @staticmethod + def _side_highlight_rows( + segments: list[Segment], + taxels: list[Taxel], + selected_kind: str | None, + selected_rows: list[int], + ) -> tuple[set[int], set[int]]: + segment_rows: set[int] = set() + taxel_rows: set[int] = set() + if selected_kind == "segments": + selected_joints = { + segments[row].joint for row in selected_rows if 0 <= row < len(segments) + } + taxel_rows = { + row + for row, taxel in enumerate(taxels) + if any( + joint in selected_joints + and type(weight) in (int, float) + and weight > 0 + for joint, weight in zip(taxel.segments, taxel.weights) + ) + } + elif selected_kind == "taxels": + bound_joints = { + joint + for row in selected_rows + if 0 <= row < len(taxels) + for joint, weight in zip(taxels[row].segments, taxels[row].weights) + if joint is not None and type(weight) in (int, float) and weight > 0 + } + segment_rows = { + row + for row, segment in enumerate(segments) + if segment.joint in bound_joints + } + return segment_rows, taxel_rows + + def _draw_dataset( + self, + prefix: str, + segments: list[Segment], + taxels: list[Taxel], + opacity: float, + active: bool, + forces=None, + ) -> None: + side_segment_rows, side_taxel_rows = ( + self._side_highlight_rows( + segments, taxels, self.selected_kind, self.selected_rows + ) + if active + else (set(), set()) + ) + segment_colors = [] + for row in range(len(segments)): + color = SEGMENT_COLOR if active else INACTIVE_COLOR + if row in side_segment_rows: + color = SIDE_HIGHLIGHT_COLOR + if ( + self.selected_kind == "segments" + and row in self.selected_rows + and active + ): + color = SELECTED_COLOR + if self.binding_phase == "segments" and row in self.binding_segments: + color = SELECTED_COLOR + if self.binding_phase == "taxels" and active: + color = INACTIVE_COLOR + if ( + self.pick_mode is not None + and self.pick_mode != "segment_pose" + and active + ): + color = INACTIVE_COLOR + segment_colors.append(color) + taxel_colors = [TAXEL_COLOR] * len(taxels) + array = np.asarray(forces, dtype=float) if forces is not None else None + if ( + array is not None + and array.shape == (len(taxels), 3) + and np.isfinite(array).all() + ): + fraction = np.clip(np.linalg.norm(array, axis=1) / FORCE_MAX_N, 0.0, 1.0) + taxel_colors = [ + tuple( + int(low + value * (high - low)) + for low, high in zip(FORCE_MIN_COLOR, FORCE_MAX_COLOR) + ) + for value in fraction + ] + for row in range(len(taxels)): + if row in side_taxel_rows: + taxel_colors[row] = SIDE_HIGHLIGHT_COLOR + if self.binding_phase == "segments" and active: + taxel_colors[row] = INACTIVE_COLOR + if self.pick_mode is not None and active: + taxel_colors[row] = INACTIVE_COLOR + if prefix == "/edit": + binding = self.document if active else self._binding_document() + error_rows = { + issue["row"] + for issue in binding.validate() + if issue["kind"] == "taxels" and issue["row"] is not None + } + for row in error_rows: + if 0 <= row < len(taxel_colors): + taxel_colors[row] = ERROR_COLOR + # Selection is the final visual state so invalid selected taxels stay visible. + if active and self.pick_mode is None: + if self.selected_kind == "taxels": + for row in self.selected_rows: + if 0 <= row < len(taxel_colors): + taxel_colors[row] = SELECTED_COLOR + if self.binding_phase == "taxels": + for row in self.binding_taxels: + if 0 <= row < len(taxel_colors): + taxel_colors[row] = SELECTED_COLOR + self._draw_instances( + prefix + "/segments", "segments", segments, segment_colors, opacity + ) + self._draw_instances( + prefix + "/taxels", "taxels", taxels, taxel_colors, opacity + ) + + def _draw_selection_gizmo(self) -> None: + values = self._selection_values() + rows = [row for row in self.selected_rows if 0 <= row < len(values)] + if not rows: + self.selection_gizmo.visible = False + self.selection_gizmo_dragging = False + self.selection_gizmo_last_position = None + return + self.selection_gizmo.visible = True + if self.selection_gizmo_dragging: + return + self.selection_gizmo.position = np.asarray( + [values[row].pose[:3] for row in rows] + ).mean(axis=0) + self.selection_gizmo.wxyz = ( + tf.quaternion_from_matrix(geo.matrix(values[rows[0]].pose)) + if len(rows) == 1 + else np.array([1.0, 0.0, 0.0, 0.0]) + ) + + def _draw(self) -> None: + active_opacity = ACTIVE_DATASET_OPACITY + if self.model is not None: + q = ( + self.document.q + if self.mode_value == "Edit" + else self.preview_document.q + ) + self.model.draw( + self.server.scene, + q, + self.urdf_alignment, + self.urdf_scale, + highlighted=self.pick_mode in ("segment", "taxel_pose"), + opacity=URDF_OPACITY * active_opacity, + ) + if self.mode_value == "Edit": + self._hide_prefix("/preview") + self._draw_dataset( + "/edit", + self.document.segments, + self.document.taxels, + active_opacity, + True, + ) + else: + valid_rows, _, source_count = self._rendered_preview_taxels() + preview_taxels = [taxel for _, taxel in valid_rows] + raw_forces = ( + self.receiver.latest.get("forces") + if self.receiver + and self.receiver.active + and time.monotonic() + - self.receiver.latest.get("read_time", float("-inf")) + < 0.5 + else None + ) + forces = ( + [raw_forces[row] for row, _ in valid_rows] + if raw_forces is not None and len(raw_forces) == source_count + else None + ) + self._draw_dataset( + "/preview", + self.preview_document.segments, + preview_taxels, + active_opacity, + True, + forces, + ) + if self.show_binding.value: + binding = self._binding_document() + self._draw_dataset( + "/edit", + binding.segments, + binding.taxels, + BINDING_OVERLAY_OPACITY, + False, + ) + else: + self._hide_prefix("/edit") + saved = self.state.saved + if self.show_saved.value and saved is not None: + self._draw_dataset( + "/saved", + saved.segments, + saved.taxels, + SAVED_BINDING_OPACITY, + False, + ) + else: + self._hide_prefix("/saved") + self._draw_selection_gizmo() + + def run(self) -> None: + try: + while True: + try: + while True: + self.events.get_nowait()() + except queue.Empty: + pass + self._update_live() + if self.redraw or self.receiver is not None: + self._draw() + self.redraw = False + time.sleep(1 / 30) + except KeyboardInterrupt: + pass + finally: + self._disconnect() + self.server.stop() diff --git a/scripts/tactile_metadata/tactile_metadata_editor/viewport_axes.js b/scripts/tactile_metadata/tactile_metadata_editor/viewport_axes.js new file mode 100644 index 0000000000..ea144aa295 --- /dev/null +++ b/scripts/tactile_metadata/tactile_metadata_editor/viewport_axes.js @@ -0,0 +1,148 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +(() => { + "use strict"; + + const key = "__tactileMetadataViewportAxes"; + globalThis[key]?.stop(); + + const ns = "http://www.w3.org/2000/svg"; + const svg = document.createElementNS(ns, "svg"); + Object.assign(svg.style, { + position: "absolute", + top: "16px", + right: "16px", + width: "112px", + height: "112px", + zIndex: "10", + pointerEvents: "none", + overflow: "visible", + }); + svg.setAttribute("viewBox", "0 0 112 112"); + + const background = document.createElementNS(ns, "circle"); + background.setAttribute("cx", "56"); + background.setAttribute("cy", "56"); + background.setAttribute("r", "49"); + background.setAttribute("fill", "rgba(255,255,255,0.72)"); + background.setAttribute("stroke", "rgba(0,0,0,0.2)"); + svg.appendChild(background); + + const axes = [ + ["x", "#e53935", [1, 0, 0]], + ["y", "#43a047", [0, 1, 0]], + ["z", "#1e88e5", [0, 0, 1]], + ].map(([name, color, vector]) => { + const line = document.createElementNS(ns, "line"); + line.setAttribute("x1", "56"); + line.setAttribute("y1", "56"); + line.setAttribute("stroke", color); + line.setAttribute("stroke-width", "3"); + line.setAttribute("stroke-linecap", "round"); + svg.appendChild(line); + + const end = document.createElementNS(ns, "circle"); + end.setAttribute("r", "4"); + end.setAttribute("stroke", color); + end.setAttribute("stroke-width", "2"); + svg.appendChild(end); + + const label = document.createElementNS(ns, "text"); + label.textContent = name; + label.setAttribute("fill", color); + label.setAttribute("font-size", "15"); + label.setAttribute("font-family", "sans-serif"); + label.setAttribute("font-weight", "700"); + label.setAttribute("text-anchor", "middle"); + label.setAttribute("dominant-baseline", "central"); + svg.appendChild(label); + return { color, end, label, line, vector }; + }); + + const multiply = (a, b) => [ + a[3] * b[0] + a[0] * b[3] + a[1] * b[2] - a[2] * b[1], + a[3] * b[1] - a[0] * b[2] + a[1] * b[3] + a[2] * b[0], + a[3] * b[2] + a[0] * b[1] - a[1] * b[0] + a[2] * b[3], + a[3] * b[3] - a[0] * b[0] - a[1] * b[1] - a[2] * b[2], + ]; + + const rotate = (q, v) => { + const tx = 2 * (q[1] * v[2] - q[2] * v[1]); + const ty = 2 * (q[2] * v[0] - q[0] * v[2]); + const tz = 2 * (q[0] * v[1] - q[1] * v[0]); + return [ + v[0] + q[3] * tx + q[1] * tz - q[2] * ty, + v[1] + q[3] * ty + q[2] * tx - q[0] * tz, + v[2] + q[3] * tz + q[0] * ty - q[1] * tx, + ]; + }; + + const state = { + controls: null, + raf: 0, + stopped: false, + stop() { + this.stopped = true; + cancelAnimationFrame(this.raf); + this.controls?.removeEventListener("update", draw); + svg.remove(); + }, + }; + globalThis[key] = state; + + const draw = () => { + const mutable = globalThis.__viserMutable; + const camera = mutable?.camera; + const root = mutable?.nodeRefFromName?.[""]; + const parent = mutable?.canvas?.parentElement; + if (!camera || !root || !parent) return; + if (svg.parentElement !== parent) parent.appendChild(svg); + + const cameraQuaternion = camera.quaternion; + const rootQuaternion = root.quaternion; + const cameraFromWorld = [ + -cameraQuaternion.x, + -cameraQuaternion.y, + -cameraQuaternion.z, + cameraQuaternion.w, + ]; + const cameraFromVisualization = multiply(cameraFromWorld, [ + rootQuaternion.x, + rootQuaternion.y, + rootQuaternion.z, + rootQuaternion.w, + ]); + + for (const [index, axis] of axes.entries()) { + const direction = rotate(cameraFromVisualization, axis.vector); + const x = 56 + 36 * direction[0]; + const y = 56 - 36 * direction[1]; + const projected = Math.hypot(direction[0], direction[1]); + const fallback = [[10, -8], [-10, -8], [0, 12]][index]; + const labelOffset = + projected > 0.12 + ? [8 * direction[0] / projected, -8 * direction[1] / projected] + : fallback; + axis.line.setAttribute("x2", x.toFixed(2)); + axis.line.setAttribute("y2", y.toFixed(2)); + axis.end.setAttribute("cx", x.toFixed(2)); + axis.end.setAttribute("cy", y.toFixed(2)); + axis.end.setAttribute("fill", direction[2] >= 0 ? axis.color : "white"); + axis.label.setAttribute("x", (x + labelOffset[0]).toFixed(2)); + axis.label.setAttribute("y", (y + labelOffset[1]).toFixed(2)); + } + }; + + const frame = () => { + if (state.stopped) return; + const controls = globalThis.__viserMutable?.cameraControl; + if (controls !== state.controls) { + state.controls?.removeEventListener("update", draw); + controls?.addEventListener("update", draw); + state.controls = controls; + } + draw(); + state.raf = requestAnimationFrame(frame); + }; + frame(); +})(); diff --git a/scripts/tactile_metadata/test_document.py b/scripts/tactile_metadata/test_document.py new file mode 100644 index 0000000000..28f1c29f41 --- /dev/null +++ b/scripts/tactile_metadata/test_document.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import unittest + +from tactile_metadata_editor.document import Document, Segment, State, Taxel + + +class BindingTests(unittest.TestCase): + def test_segment_edits_update_the_preserved_copy(self): + taxel = Taxel( + segments=["WRIST", None, None, None], + stored_segments=["WRIST", "INDEX_TIP", "WRIST", "WRIST"], + ) + + taxel.set_segment(1, "MIDDLE_TIP") + self.assertEqual(taxel.segments[1], "MIDDLE_TIP") + self.assertEqual(taxel.stored_segments[1], "MIDDLE_TIP") + + taxel.set_segment(1, None) + self.assertIsNone(taxel.segments[1]) + self.assertIsNone(taxel.stored_segments[1]) + + def test_assign_replaces_the_preserved_bindings(self): + document = Document() + document.segments = [Segment("WRIST"), Segment("INDEX_TIP")] + document.taxels = [Taxel(stored_segments=["WRIST"] * 4)] + + document.assign([0], [1]) + + expected = ["INDEX_TIP", None, None, None] + self.assertEqual(document.taxels[0].segments, expected) + self.assertEqual(document.taxels[0].stored_segments, expected) + + +class StateTests(unittest.TestCase): + def test_replace_keeps_one_document_and_an_independent_saved_snapshot(self): + state = State() + first = Document() + first.segments = [Segment("WRIST")] + first.q = {"joint": 0.5} + state.replace(first, "first.json") + + first.segments[0].pose[0] = 0.01 + first.alignment[0][3] = 0.01 + self.assertEqual(state.saved.segments[0].pose[0], 0.0) + self.assertEqual(state.saved.alignment[0][3], 0.0) + self.assertEqual(state.saved.q, {}) + self.assertEqual(first.q, {"joint": 0.5}) + + second = Document() + self.assertEqual(second.alignment[0][3], 0.0) + state.replace(second, "second.json") + self.assertIs(state.document, second) + self.assertEqual(state.path, "second.json") + self.assertIsNot(state.saved, second) + self.assertEqual(state.saved.segments, []) + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/tactile_metadata/test_editor.py b/scripts/tactile_metadata/test_editor.py new file mode 100644 index 0000000000..8e6867801c --- /dev/null +++ b/scripts/tactile_metadata/test_editor.py @@ -0,0 +1,703 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import queue +import unittest +from pathlib import Path +from tempfile import TemporaryDirectory +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import numpy as np + +from tactile_metadata_editor import geometry as geo, storage +from tactile_metadata_editor.document import JOINTS, Segment, State, Taxel +from tactile_metadata_editor.ui import ERROR_COLOR, TAXEL_COLOR, Editor + + +class EditorTests(unittest.TestCase): + def test_apply_delta_transforms_duplicate_rows_once(self): + segments = [Segment("WRIST"), Segment("INDEX_TIP")] + delta = np.eye(4) + delta[0, 3] = 0.01 + + geo.apply_delta(segments, [0, 0], delta) + + np.testing.assert_allclose(segments[0].pose[:3], [0.01, 0.0, 0.0]) + np.testing.assert_allclose(segments[1].pose[:3], [0.0, 0.0, 0.0]) + + def make_editor(self): + editor = Editor.__new__(Editor) + editor.state = State() + editor.document.segments = [Segment("WRIST")] + editor.document.taxels = [ + Taxel(segments=["WRIST", None, None, None], weights=[1.0, 0.0, 0.0, 0.0]) + ] + editor.preview_document = editor._preview_from(editor.document) + editor.binding_source_live = SimpleNamespace(value=False) + editor.info = SimpleNamespace(value="") + editor.binding_phase, editor.pick_mode = None, None + editor.events = queue.SimpleQueue() + editor.server = SimpleNamespace( + gui=SimpleNamespace(_gui_input_handle_from_uuid={}) + ) + editor.server.gui.add_button = lambda *args, **kwargs: self.make_input(editor) + editor.edit_joint_widgets = {} + editor._rebuild_edit_lists = Mock() + editor._sync_joint_widget_values = Mock() + return editor + + def make_input(self, editor, value=None): + widget = SimpleNamespace( + value=value, disabled=False, _impl=SimpleNamespace(update_cb=[]) + ) + widget.on_update = widget._impl.update_cb.append + widget.on_click = widget._impl.update_cb.append + inputs = editor.server.gui._gui_input_handle_from_uuid + inputs[str(len(inputs))] = widget + return widget + + def make_procedure_editor(self): + editor = self.make_editor() + editor.mode_edit = SimpleNamespace(value=True) + editor.selected_kind, editor.selected_rows = "taxels", [0] + editor.binding_info = SimpleNamespace(content="") + editor.binding_start = self.make_input(editor) + editor.binding_next = self.make_input(editor) + editor.binding_complete = self.make_input(editor) + editor._button(editor.binding_start, editor._start_binding) + editor._button(editor.binding_next, editor._select_binding_segments) + editor._button(editor.binding_complete, editor._complete_binding) + for mode, name in ( + ("segment", "segment_click_button"), + ("taxel_pose", "taxel_capture_button"), + ("segment_pose", "segment_capture_button"), + ): + button = editor._action_button( + mode, lambda mode=mode: editor._toggle_pick_mode(mode) + ) + setattr(editor, name, button) + editor._sync_selection = Mock() + editor._sync_taxel_binding_widgets = Mock() + editor._document_error_html = Mock(return_value="") + return editor + + @staticmethod + def dispatch(widget, client=None): + async def emit(): + for callback in widget._impl.update_cb: + await callback(SimpleNamespace(client=client, target=widget)) + + asyncio.run(emit()) + + @staticmethod + def drain_events(editor): + while not editor.events.empty(): + editor.events.get_nowait()() + + def make_live_editor(self, frame): + editor = self.make_editor() + editor.receiver = SimpleNamespace(status="Live", active=True, latest=frame) + editor.live_status = SimpleNamespace(content="") + editor.live_document, editor.live_metadata_key = None, "" + editor.binding_source_live.disabled = True + editor.joint_enabled = SimpleNamespace(value=False) + editor.joint_mapping = SimpleNamespace(value="{}") + editor.preview_error = editor.info + editor._sync_preview_values = Mock() + editor._rebuild_preview_segments = Mock() + editor._update_preview_error = Mock() + editor.mode_edit = SimpleNamespace(value=False) + editor.selected_kind, editor.selected_rows = None, [] + editor._sync_selection = Mock() + return editor + + def test_malformed_live_input_reports_an_error_without_raising(self): + for frame, mapping, message in ( + ({"metadata": {}}, "{}", "Expected geometry_mapping"), + ({}, "[]", "Joint Mappings JSON must be an object"), + ({}, "null", "Joint Mappings JSON must be an object"), + ({}, "42", "Joint Mappings JSON must be an object"), + ): + with self.subTest(frame=frame, mapping=mapping): + editor = self.make_live_editor(frame) + editor.joint_enabled.value = True + editor.joint_mapping.value = mapping + + editor._update_live() + + self.assertIn("Live error: " + message, editor.info.value) + self.assertTrue(editor.redraw) + + def test_invalid_live_hand_retains_applicable_pose_and_recovers_status(self): + valid_hand = { + "positions": [[0.0, 0.0, 0.0] for _ in JOINTS], + "orientations": [[0.0, 0.0, 0.0, 1.0] for _ in JOINTS], + "is_valid": [True for _ in JOINTS], + } + invalid_wrist = { + **valid_hand, + "is_valid": [joint != "WRIST" for joint in JOINTS], + } + for hand in (None, invalid_wrist): + with self.subTest(hand=hand): + editor = self.make_live_editor({"hand": hand}) + retained = editor.preview_document.segments + retained[0].pose[0] = 0.02 + + editor._update_live() + + self.assertIs(editor.preview_document.segments, retained) + self.assertEqual( + editor.live_status.content, "Current Status: Data Invalid" + ) + editor._rebuild_preview_segments.assert_not_called() + editor._apply_preview_segments() + self.assertEqual(editor.document.segments, retained) + + editor.receiver.latest = {"hand": valid_hand} + editor._update_live() + + self.assertEqual(editor.live_status.content, "Current Status: Live") + self.assertEqual( + [segment.joint for segment in editor.preview_document.segments], + list(JOINTS), + ) + + def test_live_joint_changes_clear_only_preview_selection(self): + original = ["PALM", "WRIST", "THUMB_METACARPAL"] + for editing, joints in ( + (False, original[1:]), + (False, original), + (True, original[1:]), + ): + with self.subTest(editing=editing, joints=joints): + editor = self.make_live_editor({"hand": {}}) + editor.mode_edit.value = editing + editor.preview_document.segments = [ + Segment(joint) for joint in original + ] + editor.selected_kind, editor.selected_rows = "segments", [1] + with patch.object( + geo, + "hand_segments", + return_value=[Segment(joint) for joint in joints], + ): + editor._update_live() + + cleared = not editing and joints != original + self.assertEqual(editor.selected_kind, None if cleared else "segments") + self.assertEqual(editor.selected_rows, [] if cleared else [1]) + self.assertEqual(editor._sync_selection.call_count, int(cleared)) + + def test_reset_preview_refreshes_selection_without_retargeting_it(self): + for editing, joints in ( + (False, ["PALM", "WRIST"]), + (False, ["WRIST"]), + (True, ["PALM", "WRIST"]), + ): + with self.subTest(editing=editing, joints=joints): + editor = self.make_live_editor({}) + editor.mode_edit.value = editing + editor.preview_document.segments = [Segment(joint) for joint in joints] + editor.document.segments[0].pose[0] = 0.02 + editor.selected_kind, editor.selected_rows = "segments", [0] + + editor._restore_preview() + + cleared = not editing and joints != ["WRIST"] + self.assertEqual(editor.selected_kind, None if cleared else "segments") + self.assertEqual(editor.selected_rows, [] if cleared else [0]) + self.assertEqual(editor._sync_selection.call_count, int(not editing)) + self.assertEqual( + editor.preview_document.segments, editor.document.segments + ) + + def test_apply_all_preserves_rendered_taxels(self): + editor = self.make_editor() + editor.preview_document.segments[0].pose[0] = 0.01 + editor.preview_document.q = {"joint": 0.5} + rows, errors, _ = editor._rendered_preview_taxels() + self.assertFalse(errors) + expected = rows[0][1].pose + self.assertAlmostEqual(expected[0], 0.01) + + editor._apply_preview_all() + + np.testing.assert_allclose(editor.document.taxels[0].pose, expected) + self.assertEqual(editor.document.segments, editor.preview_document.segments) + self.assertEqual(editor.document.q, editor.preview_document.q) + + def test_joint_slider_initialization_matches_the_rendered_configuration(self): + for lower, upper, initial, expected in ( + (1.0, 2.0, None, 1.0), + (-2.0, -1.0, None, -1.0), + (1.0, 2.0, 4.0, 2.0), + ): + with self.subTest(lower=lower, upper=upper, initial=initial): + editor = self.make_editor() + editor.model = SimpleNamespace( + joints=[ + SimpleNamespace( + name="joint", + limit=SimpleNamespace(lower=lower, upper=upper), + ) + ] + ) + editor.server = SimpleNamespace(gui=Mock()) + editor._watch = Mock() + if initial is not None: + editor.document.q["joint"] = initial + + editor._add_joint_sliders(editor.document, {}) + + self.assertEqual(editor.document.q["joint"], expected) + self.assertEqual( + editor.server.gui.add_slider.call_args.args[-1], expected + ) + + def test_load_preserves_current_alignment_and_visible_poses(self): + editor = self.make_editor() + editor.document.alignment = geo.transform( + (0.12, 0.03, -0.04), (10.0, 20.0, 30.0) + ).tolist() + expected = editor.document.clone() + editor._sync_document = Mock() + with TemporaryDirectory() as directory: + path = Path(directory) / "metadata.json" + storage.save(editor.state, path) + editor.metadata_path = SimpleNamespace(value=str(path)) + + editor._load() + + for document in (editor.document, editor.state.saved, editor.preview_document): + np.testing.assert_allclose(document.alignment, expected.alignment) + np.testing.assert_allclose( + document.segments[0].pose, expected.segments[0].pose, atol=1e-12 + ) + for document in (editor.document, editor.state.saved): + np.testing.assert_allclose( + document.taxels[0].pose, expected.taxels[0].pose, atol=1e-12 + ) + editor._sync_document.assert_called_once() + + def test_error_html_escapes_external_text(self): + editor = self.make_editor() + payload = '' + editor.document.taxels[0].segments[0] = payload + validation = editor._document_error_html() + + editor.preview_error = editor.info + editor._rendered_preview_taxels = Mock(return_value=([], [payload], 0)) + editor._update_preview_error() + preview = editor.info.value + + widget = SimpleNamespace(content="", visible=False) + editor._set_error_widget(widget, [payload, payload]) + self.assertTrue(widget.visible) + self.assertEqual(widget.content.count("<img"), 1) + + with self.assertLogs("isaaccapture.tactile_metadata.editor", level="ERROR"): + editor._run_action(Mock(side_effect=ValueError(payload))) + for rendered in (validation, preview, widget.content, editor.info.value): + self.assertIn("<img src=x onerror="alert(1)">", rendered) + self.assertNotIn(payload, rendered) + + def test_list_pose_edits_refresh_selection_before_further_inspector_edits(self): + for editing, kind in ( + (True, "segments"), + (True, "taxels"), + (False, "segments"), + ): + for rows in ([0], [0, 1]): + with self.subTest(editing=editing, kind=kind, rows=rows): + editor = self.make_editor() + editor.mode_edit = SimpleNamespace(value=editing) + values = ( + editor._values(kind) + if editing + else editor.preview_document.segments + ) + values.append( + Segment("INDEX_TIP") if kind == "segments" else Taxel() + ) + editor.selected_kind, editor.selected_rows = kind, rows.copy() + editor.selection_folder = object() + for name in ( + "selection_summary", + "selection_names", + "selection_segment_editor", + "selection_taxel_editor", + "selection_segment_joint", + "selection_scale_widget", + "selection_delete", + ): + setattr(editor, name, self.make_input(editor)) + editor.selection_position_widgets = [ + self.make_input(editor, 0.0) for _ in range(3) + ] + editor.selection_rotation_widgets = [ + self.make_input(editor, 0.0) for _ in range(3) + ] + editor.selection_taxel_binding_widgets = [ + (self.make_input(editor), self.make_input(editor)) + for _ in range(4) + ] + editor._update_preview_error = Mock() + editor._sync_selected_list_values = Mock() + editor._sync_selection() + + if editing: + editor._set_pose_component(kind, 0, 0, 0.02) + editor._set_pose_rotation(kind, 0, 2, 30.0) + else: + editor._set_preview_pose(0, 0, 0.02, False) + editor._set_preview_pose(0, 2, 30.0, True) + + position = editor.selection_position_widgets[0].value + rotation = editor.selection_rotation_widgets[2].value + self.assertAlmostEqual(position, 0.02 / len(rows)) + self.assertAlmostEqual(rotation, 30.0 if len(rows) == 1 else 0.0) + editor._set_selection_position(0, position + 0.001) + editor._set_selection_rotation(2, rotation + 1.0) + center = np.mean([values[row].pose[:3] for row in rows], axis=0) + self.assertAlmostEqual(center[0], position + 0.001) + self.assertAlmostEqual( + geo.components(geo.matrix(values[0].pose))[1][2], 31.0 + ) + + def test_binding_overlay_follows_selected_source_and_its_validation(self): + editor = self.make_editor() + editor.live_document = editor.document.clone() + editor.live_document.segments[0].pose[0] = 0.02 + editor.live_document.taxels[0].pose[0] = 0.03 + editor.document.taxels[0].segments[0] = "MISSING" + editor.model, editor.receiver = None, None + editor.mode_edit = SimpleNamespace(value=False) + editor.show_binding = SimpleNamespace(value=True) + editor.show_saved = SimpleNamespace(value=False) + editor.selected_kind, editor.selected_rows = None, [] + editor.binding_phase, editor.binding_segments = None, [] + editor.pick_mode = None + editor._rendered_preview_taxels = Mock(return_value=([], [], 0)) + editor._draw_instances = Mock() + editor._hide_prefix = Mock() + editor._draw_selection_gizmo = Mock() + + for live in (False, True): + with self.subTest(live=live): + editor.binding_source_live.value = live + editor._draw_instances.reset_mock() + + editor._draw() + + calls = { + call.args[0]: call.args + for call in editor._draw_instances.call_args_list + } + source = editor.live_document if live else editor.document + self.assertIs(calls["/edit/segments"][2], source.segments) + self.assertIs(calls["/edit/taxels"][2], source.taxels) + self.assertEqual( + calls["/edit/taxels"][3], [TAXEL_COLOR if live else ERROR_COLOR] + ) + + def test_start_binding_reuses_only_selected_taxels(self): + for kind, rows in ( + ("taxels", [0]), + ("taxels", [0, 1]), + ("segments", [0]), + (None, []), + ): + with self.subTest(kind=kind, rows=rows): + editor = self.make_procedure_editor() + editor.document.taxels.append(Taxel()) + editor.selected_kind, editor.selected_rows = kind, rows.copy() + + editor._start_binding() + + expected = rows if kind == "taxels" else [] + self.assertEqual(editor.binding_taxels, expected) + self.assertEqual(editor.binding_segments, []) + self.assertEqual(editor.selected_rows, []) + self.assertIsNone(editor.selected_kind) + editor._sync_selection.assert_called_once() + self.assertEqual(editor.binding_phase, "taxels") + self.assertTrue(editor.binding_start.disabled) + self.assertEqual(editor.binding_next.disabled, not expected) + self.assertTrue(editor.binding_complete.disabled) + + def test_procedures_listen_to_passive_inputs_and_remove_the_listener(self): + for mode in ("binding", "segment", "taxel_pose", "segment_pose"): + with self.subTest(mode=mode): + editor = self.make_procedure_editor() + widget = self.make_input(editor) + if mode == "binding": + editor._start_binding() + else: + editor._toggle_pick_mode(mode) + self.assertIsNone(editor.selected_kind) + self.assertEqual(editor.selected_rows, []) + editor._sync_selection.assert_called_once() + + self.dispatch(widget) + widget.disabled = True + self.dispatch(widget, client=object()) + self.assertTrue(editor.events.empty()) + + widget.disabled = False + self.dispatch(widget, client=object()) + with self.assertLogs( + "isaaccapture.tactile_metadata.editor", level="WARNING" + ) as logged: + self.drain_events(editor) + self.assertEqual(len(logged.records), 1) + self.assertIsNone(editor.binding_phase) + self.assertIsNone(editor.pick_mode) + for handle in editor.server.gui._gui_input_handle_from_uuid.values(): + self.assertNotIn( + editor._procedure_interaction, handle._impl.update_cb + ) + self.dispatch(widget, client=object()) + self.assertTrue(editor.events.empty()) + + def test_interruption_runs_before_button_and_value_business_callbacks(self): + for button in (False, True): + with self.subTest(button=button): + editor = self.make_procedure_editor() + observed = [] + + def observe(*_, editor=editor, observed=observed): + observed.append(editor.binding_phase) + + if button: + widget = editor._action_button("Other action", observe) + else: + widget = self.make_input(editor, 2.0) + + async def flush_interruption(_, editor=editor): + self.drain_events(editor) + + widget.on_update(flush_interruption) + editor._watch(widget, observe) + editor._start_binding() + + with self.assertLogs( + "isaaccapture.tactile_metadata.editor", level="WARNING" + ): + self.dispatch(widget, client=object()) + self.drain_events(editor) + + self.assertEqual(observed, [None]) + + def test_binding_buttons_and_scene_picks_preserve_the_flow_until_completion(self): + editor = self.make_procedure_editor() + editor._start_binding() + self.dispatch(editor.binding_next, client=object()) + self.drain_events(editor) + self.assertEqual(editor.binding_phase, "segments") + + editor._handle_scene_event = Mock() + asyncio.run(editor._scene_click(SimpleNamespace())) + self.drain_events(editor) + editor._handle_scene_event.assert_called_once() + self.assertEqual(editor.binding_phase, "segments") + + editor.binding_segments = [0] + editor._set_binding_button("complete") + self.dispatch(editor.binding_complete, client=object()) + self.drain_events(editor) + self.assertIsNone(editor.binding_phase) + self.assertEqual(editor.selected_kind, "taxels") + self.assertEqual(editor.selected_rows, [0]) + for widget in editor.server.gui._gui_input_handle_from_uuid.values(): + self.assertNotIn(editor._procedure_interaction, widget._impl.update_cb) + + def test_capture_rectangles_do_not_select_objects_or_enable_the_gizmo(self): + for mode in ("segment", "taxel_pose", "segment_pose"): + with self.subTest(mode=mode): + editor = self.make_procedure_editor() + editor.selection_gizmo = SimpleNamespace(visible=True) + editor._box_rows = Mock(return_value=[0]) + editor._handle_pick = Mock() + editor._toggle_pick_mode(mode) + editor._sync_selection.reset_mock() + event = SimpleNamespace(modifier="shift") + + editor._handle_scene_event(event, True) + editor._draw_selection_gizmo() + + self.assertIsNone(editor.selected_kind) + self.assertEqual(editor.selected_rows, []) + self.assertEqual(editor.pick_mode, mode) + self.assertFalse(editor.selection_gizmo.visible) + editor._sync_selection.assert_not_called() + editor._box_rows.assert_not_called() + editor._handle_pick.assert_not_called() + + editor._handle_scene_event(event, False) + editor._handle_pick.assert_called_once_with(event) + + def test_capture_buttons_toggle_or_replace_the_active_procedure(self): + editor = self.make_procedure_editor() + editor._toggle_pick_mode("taxel_pose") + for button, expected in ( + (editor.segment_capture_button, "segment_pose"), + (editor.segment_capture_button, None), + ): + self.dispatch(button, client=object()) + with self.assertLogs( + "isaaccapture.tactile_metadata.editor", level="WARNING" + ) as logged: + self.drain_events(editor) + self.assertEqual(len(logged.records), 1) + self.assertEqual(editor.pick_mode, expected) + + def test_deleted_segment_remains_invalid_after_round_trip(self): + document = self.make_editor().document + document.segments.append(Segment("INDEX_TIP")) + document.taxels[0].segments[0] = "INDEX_TIP" + document.delete("segments", [1]) + + encoded = storage.metadata(document) + expected = [{"values": [0xFFFFFFFF, 0, 0, 0]}] + self.assertEqual(encoded["geometry_mapping"]["taxel_segment_indices"], expected) + + loaded = storage.parse(encoded) + self.assertTrue( + any(issue["field"] == "segment0" for issue in loaded.validate()) + ) + self.assertEqual( + storage.metadata(loaded)["geometry_mapping"]["taxel_segment_indices"], + expected, + ) + + def test_zero_weight_segment_indices_survive_round_trip(self): + joints = ["WRIST", "INDEX_TIP", "MIDDLE_TIP", "THUMB_TIP"] + for bindings, indices in ( + (["WRIST", None, None, None], [0, 0, 0, 0]), + (joints, [0, 1, 2, 3]), + ): + with self.subTest(bindings=bindings): + document = self.make_editor().document + document.segments = [Segment(joint) for joint in joints] + document.taxels[0].segments = bindings.copy() + encoded = storage.metadata(document) + + loaded = storage.parse(encoded) + + self.assertEqual(loaded.taxels[0].segments, ["WRIST", None, None, None]) + self.assertEqual(loaded.taxels[0].weights, [1.0, 0.0, 0.0, 0.0]) + self.assertFalse(loaded.validate()) + self.assertEqual( + storage.metadata(loaded)["geometry_mapping"][ + "taxel_segment_indices" + ], + [{"values": indices}], + ) + + def test_segment_edits_and_preview_apply_preserve_serialized_slots(self): + editor = self.make_editor() + joints = ["WRIST", "INDEX_TIP", "MIDDLE_TIP", "THUMB_TIP"] + editor.document.segments = [Segment(joint) for joint in joints] + editor.document.taxels[0].segments = joints.copy() + editor.state.document = storage.parse(storage.metadata(editor.document)) + editor.preview_document = editor._preview_from(editor.document) + editor._sync_taxel_binding_widgets = Mock() + editor._sync_selection = Mock() + + editor._set_taxel_segment(0, 1, "MIDDLE_TIP") + self.assertEqual(editor.document.taxels[0].stored_segments[1], "MIDDLE_TIP") + self.assertEqual( + storage.metadata(editor.document)["geometry_mapping"][ + "taxel_segment_indices" + ], + [{"values": [0, 2, 2, 3]}], + ) + + editor._set_taxel_segment(0, 1, "None") + source = editor.document.taxels[0].stored_segments + editor._apply_preview_taxels() + self.assertIsNot(editor.document.taxels[0].stored_segments, source) + self.assertEqual( + storage.metadata(editor.document)["geometry_mapping"][ + "taxel_segment_indices" + ], + [{"values": [0, 0, 2, 3]}], + ) + + def test_deleting_an_unused_loaded_segment_keeps_padding_valid(self): + document = self.make_editor().document + document.segments.append(Segment("INDEX_TIP")) + document.taxels[0].segments[1] = "INDEX_TIP" + loaded = storage.parse(storage.metadata(document)) + + loaded.delete("segments", [1]) + + self.assertFalse(loaded.validate()) + self.assertEqual( + storage.metadata(loaded)["geometry_mapping"]["taxel_segment_indices"], + [{"values": [0, 0, 0, 0]}], + ) + + def test_hand_side_change_preserves_document_and_snapshots(self): + editor = self.make_editor() + editor.state.replace(editor.document, "metadata.json") + document, saved = editor.document, editor.state.saved + preview = editor.preview_document + receiver = SimpleNamespace(close=Mock()) + editor.receiver = receiver + editor.connect_button = SimpleNamespace(label="Disconnect Live") + editor.live_status = SimpleNamespace(content="Live") + editor._update_preview_error = Mock() + editor.live_document, editor.live_metadata_key = document.clone(), "left" + editor.binding_source_editor = SimpleNamespace(value=False) + editor.binding_source_live = SimpleNamespace(value=True, disabled=False) + + editor._hand_side_changed("right") + + receiver.close.assert_called_once() + editor._update_preview_error.assert_called_once() + self.assertIsNone(editor.receiver) + self.assertIs(editor.document, document) + self.assertIs(editor.state.saved, saved) + self.assertIs(editor.preview_document, preview) + self.assertEqual(editor.state.path, "metadata.json") + self.assertIsNone(editor.live_document) + self.assertEqual(editor.live_metadata_key, "") + self.assertTrue(editor.binding_source_editor.value) + self.assertFalse(editor.binding_source_live.value) + self.assertTrue(editor.binding_source_live.disabled) + + def test_connect_discards_metadata_from_the_previous_stream(self): + editor = self.make_live_editor({}) + editor.receiver = None + editor.live_document, editor.live_metadata_key = ( + editor.document.clone(), + "previous", + ) + editor.binding_source_editor = SimpleNamespace(value=False) + editor.binding_source_live = SimpleNamespace(value=True, disabled=False) + editor.hand_side = SimpleNamespace(value="left") + editor.collection = SimpleNamespace(value="next_collection") + editor.connect_button = SimpleNamespace(label="Connect Live") + preview = editor.preview_document + + with patch("tactile_metadata_editor.ui.sdk.Receiver") as receiver: + editor._connect() + + receiver.assert_called_once_with( + {"side": "left", "collection": "next_collection", "joint_state": ""} + ) + self.assertIs(editor.receiver, receiver.return_value) + self.assertIsNone(editor.live_document) + self.assertEqual(editor.live_metadata_key, "") + self.assertTrue(editor.binding_source_editor.value) + self.assertFalse(editor.binding_source_live.value) + self.assertTrue(editor.binding_source_live.disabled) + self.assertIs(editor.preview_document, preview) + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/tactile_metadata/test_sdk.py b/scripts/tactile_metadata/test_sdk.py new file mode 100644 index 0000000000..631d5862f9 --- /dev/null +++ b/scripts/tactile_metadata/test_sdk.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import unittest +from unittest.mock import Mock, patch + +from tactile_metadata_editor.sdk import Receiver + + +class ReceiverTests(unittest.TestCase): + def test_receiver_is_inactive_until_the_first_frame(self): + context = Mock() + context.Pipe.return_value = Mock(), Mock() + with ( + patch( + "tactile_metadata_editor.sdk.multiprocessing.get_context", + return_value=context, + ), + patch("tactile_metadata_editor.sdk.threading.Thread"), + ): + receiver = Receiver({}) + + self.assertFalse(receiver.active) + self.assertEqual(receiver.status, "Connecting") + + def receive(_): + if not receiver.latest: + return b'{"frame": {"hand": null}}' + self.assertTrue(receiver.active) + self.assertEqual(receiver.status, "Live") + raise EOFError + + receiver.connection.recv_bytes.side_effect = receive + receiver._receive() + self.assertEqual(receiver.latest, {"hand": None}) + self.assertFalse(receiver.active) + + def test_pipe_close_preserves_child_errors(self): + for messages, expected in ( + ( + [b'{"status": "ImportError: SDK unavailable"}'], + "ImportError: SDK unavailable", + ), + ([b'{"frame": {}}'], "Disconnected"), + ([], "Disconnected"), + ): + with self.subTest(messages=messages): + receiver = Receiver.__new__(Receiver) + receiver.status, receiver.active = "Connecting", True + receiver.connection = Mock() + receiver.connection.recv_bytes.side_effect = [*messages, EOFError()] + + receiver._receive() + + self.assertEqual(receiver.status, expected) + self.assertFalse(receiver.active) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/core/python/requirements-ui.txt b/src/core/python/requirements-ui.txt index 5bb07526e4..a0c3a7521b 100644 --- a/src/core/python/requirements-ui.txt +++ b/src/core/python/requirements-ui.txt @@ -1,2 +1,9 @@ # Optional UI dependencies for parameter tuning imgui[glfw]>=2.0.0 +numpy>=1.23.0 +# Dockable panels and Trimesh 5 compatibility require Viser >=1.1.1. +# https://github.com/viser-project/viser/pull/769 +viser>=1.1.1 +yourdfpy +trimesh +rtree