dhb-xr 0.2.1__py3-none-any.whl

dhb_xr/utils/__init__.py

@@ -0,0 +1,24 @@
+"""Utilities: I/O, resampling, XDOF loader."""
+
+from dhb_xr.utils.resampling import resample_trajectory
+
+__all__ = [
+    "resample_trajectory",
+    "load_trajectory",
+    "save_trajectory",
+    "load_xdof_episode",
+    "list_xdof_episodes",
+    "se3_poses_to_positions_quaternions",
+]
+
+try:
+    from dhb_xr.utils.io import load_trajectory, save_trajectory
+except ImportError:
+    load_trajectory = None
+    save_trajectory = None
+
+from dhb_xr.utils.xdof_loader import (
+    load_xdof_episode,
+    list_xdof_episodes,
+    se3_poses_to_positions_quaternions,
+)

dhb_xr/utils/io.py

@@ -0,0 +1,59 @@
+"""I/O for trajectories (load/save)."""
+
+import numpy as np
+from typing import Dict, Any, Optional
+from pathlib import Path
+
+
+def load_trajectory(path: str) -> Dict[str, np.ndarray]:
+    """Load trajectory from .npz or .txt. Returns dict with 'positions' (N,3), 'quaternions' (N,4) wxyz."""
+    path = Path(path)
+    if path.suffix == ".npz":
+        data = np.load(path)
+        return {
+            "positions": data["positions"],
+            "quaternions": data["quaternions"],
+        }
+    if path.suffix in (".txt", ".csv"):
+        arr = np.loadtxt(path)
+        if arr.ndim == 1:
+            arr = arr.reshape(1, -1)
+        # Assume columns: x, y, z, qw, qx, qy, qz or timestamp, x, y, z, qx, qy, qz, qw
+        if arr.shape[1] >= 7:
+            positions = arr[:, :3]
+            quat = arr[:, 3:7]
+            if quat.shape[1] == 4:
+                # If scalar-last (xyzw), convert to wxyz
+                if np.all(np.abs(quat[:, 0]) < 1.1):
+                    quat = np.column_stack([quat[:, 3], quat[:, 0], quat[:, 1], quat[:, 2]])
+                quaternions = quat
+            else:
+                quaternions = np.tile(np.array([1.0, 0, 0, 0]), (len(positions), 1))
+        else:
+            positions = arr[:, :3]
+            quaternions = np.tile(np.array([1.0, 0, 0, 0]), (len(positions), 1))
+        return {"positions": positions, "quaternions": quaternions}
+    raise ValueError(f"Unsupported format: {path.suffix}")
+
+
+def save_trajectory(
+    path: str,
+    positions: np.ndarray,
+    quaternions: np.ndarray,
+    timestamps: Optional[np.ndarray] = None,
+) -> None:
+    """Save trajectory to .npz or .txt."""
+    path = Path(path)
+    if path.suffix == ".npz":
+        out = {"positions": np.asarray(positions), "quaternions": np.asarray(quaternions)}
+        if timestamps is not None:
+            out["timestamps"] = np.asarray(timestamps)
+        np.savez(path, **out)
+        return
+    if path.suffix in (".txt", ".csv"):
+        arr = np.hstack([positions, quaternions])
+        if timestamps is not None:
+            arr = np.column_stack([timestamps, arr])
+        np.savetxt(path, arr)
+        return
+    raise ValueError(f"Unsupported format: {path.suffix}")

dhb_xr/utils/resampling.py

@@ -0,0 +1,66 @@
+"""Trajectory resampling and smoothing."""
+
+import numpy as np
+from typing import Tuple, Optional
+from scipy.interpolate import CubicSpline
+
+from dhb_xr.core import geometry as geom
+
+
+def resample_trajectory(
+    positions: np.ndarray,
+    quaternions: np.ndarray,
+    num_points: int,
+    smoothing: bool = False,
+    window_length: int = 10,
+    polyorder: int = 3,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Resample trajectory to num_points. Optionally apply Savitzky-Golay smoothing.
+    positions (N, 3), quaternions (N, 4) wxyz.
+    Returns (positions_resampled, quaternions_resampled) (num_points, 3), (num_points, 4).
+    """
+    from scipy.signal import savgol_filter
+
+    n = positions.shape[0]
+    time_orig = np.linspace(0, n - 1, n)
+    time_new = np.linspace(0, n - 1, num_points)
+    rvec = np.array([geom.quat_to_axis_angle(quaternions[i]) for i in range(n)])
+    pos_resample = np.zeros((num_points, 3))
+    rvec_resample = np.zeros((num_points, 3))
+    for i in range(3):
+        cs_p = CubicSpline(time_orig, positions[:, i])
+        cs_r = CubicSpline(time_orig, rvec[:, i])
+        pos_resample[:, i] = cs_p(time_new)
+        rvec_resample[:, i] = cs_r(time_new)
+        if smoothing:
+            pos_resample[:, i] = savgol_filter(
+                pos_resample[:, i], window_length=window_length, polyorder=polyorder
+            )
+            rvec_resample[:, i] = savgol_filter(
+                rvec_resample[:, i], window_length=window_length, polyorder=polyorder
+            )
+    quat_resample = np.array([geom.axis_angle_to_quat(rvec_resample[i]) for i in range(num_points)])
+    return pos_resample, quat_resample
+
+
+def resample_and_smooth(
+    pos_data: np.ndarray,
+    quat_data: np.ndarray,
+    length: int,
+    smoothing: bool = False,
+    window_length: int = 10,
+    polyorder: int = 3,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """
+    Resample to length points; return both original-length resampled and padded for encoder.
+    Returns (pos_orig, quat_orig, rvec_orig, pos_padded, quat_padded, rvec_padded).
+    """
+    pos_orig, quat_orig = resample_trajectory(
+        pos_data, quat_data, length, smoothing, window_length, polyorder
+    )
+    rvec_orig = np.array([geom.quat_to_axis_angle(quat_orig[i]) for i in range(length)])
+    pos_padded = np.vstack([pos_orig, np.tile(pos_orig[-1], (3, 1))])
+    quat_padded = np.vstack([quat_orig, np.tile(quat_orig[-1], (3, 1))])
+    rvec_padded = np.vstack([rvec_orig, np.tile(rvec_orig[-1], (3, 1))])
+    return pos_orig, quat_orig, rvec_orig, pos_padded, quat_padded, rvec_padded

dhb_xr/utils/xdof_loader.py

@@ -0,0 +1,89 @@
+"""
+Load UMI-XDOF trajectory data from numpy files (no xdof_sdk dependency).
+
+Episodes are directories containing:
+  - action-{arm}-hand_in_quest_world_frame.npy  (N, 4, 4) SE(3) poses
+  - metadata.json
+"""
+
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+import numpy as np
+
+from dhb_xr.core import geometry as geom
+
+
+def se3_poses_to_positions_quaternions(poses: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Convert (N, 4, 4) SE(3) pose matrices to positions (N, 3) and quaternions (N, 4) wxyz.
+
+    Args:
+        poses: Array of shape (N, 4, 4) homogeneous transformation matrices.
+
+    Returns:
+        positions: (N, 3) XYZ positions.
+        quaternions: (N, 4) quaternions in wxyz (scalar-first) convention.
+    """
+    poses = np.asarray(poses)
+    if poses.ndim == 2:
+        poses = poses[np.newaxis, ...]
+    positions = poses[:, :3, 3].copy()
+    quaternions = np.array([geom.rot_to_quat(poses[i, :3, :3]) for i in range(len(poses))])
+    return positions, quaternions
+
+
+def load_xdof_episode(
+    episode_path: Path,
+    arm: str = "left",
+) -> Optional[Tuple[np.ndarray, np.ndarray]]:
+    """
+    Load one UMI-XDOF episode and return (positions, quaternions) for the given arm.
+
+    Args:
+        episode_path: Path to episode directory (contains action-*-hand_in_quest_world_frame.npy).
+        arm: "left" or "right".
+
+    Returns:
+        (positions, quaternions) with shapes (N, 3) and (N, 4) wxyz, or None if load fails.
+    """
+    episode_path = Path(episode_path)
+    fname = f"action-{arm.lower()}-hand_in_quest_world_frame.npy"
+    npy_path = episode_path / fname
+    if not npy_path.exists():
+        return None
+    try:
+        poses = np.load(npy_path)
+    except Exception:
+        return None
+    if poses.ndim != 3 or poses.shape[1] != 4 or poses.shape[2] != 4:
+        return None
+    return se3_poses_to_positions_quaternions(poses)
+
+
+def list_xdof_episodes(data_root: Path) -> List[Tuple[str, Path]]:
+    """
+    List all (task, episode_path) pairs under data_root.
+
+    Expects structure: data_root / task_name / episode_xxx / metadata.json (or .npy files).
+
+    Returns:
+        List of (task_name, episode_dir_path).
+    """
+    data_root = Path(data_root)
+    if not data_root.is_dir():
+        return []
+    out = []
+    for task_dir in sorted(data_root.iterdir()):
+        if not task_dir.is_dir():
+            continue
+        task_name = task_dir.name
+        for ep_dir in sorted(task_dir.iterdir()):
+            if not ep_dir.is_dir() or not ep_dir.name.startswith("episode_"):
+                continue
+            # Consider valid if it has at least one hand pose file
+            if (ep_dir / "action-left-hand_in_quest_world_frame.npy").exists() or (
+                ep_dir / "action-right-hand_in_quest_world_frame.npy"
+            ).exists():
+                out.append((task_name, ep_dir))
+    return out

dhb_xr/visualization/__init__.py

@@ -0,0 +1,5 @@
+"""Visualization: SE3 trajectories and invariant plots."""
+
+from dhb_xr.visualization.plot import plot_se3_trajectory, plot_invariants
+
+__all__ = ["plot_se3_trajectory", "plot_invariants"]

dhb_xr/visualization/plot.py

@@ -0,0 +1,242 @@
+"""Plot SE(3) trajectories and invariant sequences with cubes and coordinate frames."""
+
+import numpy as np
+from typing import Optional, Dict, Any, List, Union
+
+try:
+    import matplotlib.pyplot as plt
+    from mpl_toolkits.mplot3d import Axes3D
+    from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+    HAS_MPL = True
+except ImportError:
+    HAS_MPL = False
+
+# Axis colors: x=red, y=green, z=blue
+AXIS_COLORS = ["r", "g", "b"]
+
+
+def draw_box(
+    ax,
+    center: np.ndarray,
+    rotation_matrix: np.ndarray,
+    size: tuple = (0.05, 0.05, 0.05),
+    color: str = "blue",
+    alpha: float = 0.3,
+) -> None:
+    """
+    Draw a 3D oriented box (cuboid) at the given center with the given rotation.
+
+    Args:
+        ax: Matplotlib 3D axes
+        center: (3,) position of box center
+        rotation_matrix: (3, 3) rotation matrix
+        size: (width, height, depth) of box
+        color: Face color
+        alpha: Transparency
+    """
+    # Cuboid vertices before rotation (unit cube centered at origin)
+    vertices = np.array([
+        [0.5, 0.5, -0.5], [-0.5, 0.5, -0.5], [-0.5, -0.5, -0.5], [0.5, -0.5, -0.5],  # Bottom
+        [0.5, 0.5, 0.5], [-0.5, 0.5, 0.5], [-0.5, -0.5, 0.5], [0.5, -0.5, 0.5]  # Top
+    ], dtype=np.float64)
+
+    # Scale and rotate vertices
+    vertices = vertices * np.array(size)
+    vertices = np.dot(vertices, rotation_matrix.T)
+    vertices += center
+
+    # Define the indices of vertices for each face of the cuboid
+    faces = [
+        [vertices[0], vertices[1], vertices[2], vertices[3]],
+        [vertices[4], vertices[5], vertices[6], vertices[7]],
+        [vertices[0], vertices[3], vertices[7], vertices[4]],
+        [vertices[1], vertices[2], vertices[6], vertices[5]],
+        [vertices[0], vertices[1], vertices[5], vertices[4]],
+        [vertices[2], vertices[3], vertices[7], vertices[6]]
+    ]
+
+    # Create a Poly3DCollection object
+    face_collection = Poly3DCollection(faces, facecolors=color, linewidths=1, edgecolors='k', alpha=alpha)
+    ax.add_collection3d(face_collection)
+
+
+def draw_frame(
+    ax,
+    position: np.ndarray,
+    rotation_matrix: np.ndarray,
+    length: float = 0.05,
+    linewidth: float = 2.0,
+) -> None:
+    """
+    Draw a coordinate frame (x=red, y=green, z=blue arrows) at the given pose.
+
+    Args:
+        ax: Matplotlib 3D axes
+        position: (3,) position of frame origin
+        rotation_matrix: (3, 3) rotation matrix
+        length: Length of arrows
+        linewidth: Line width
+    """
+    for j, c in enumerate(AXIS_COLORS):
+        ax.quiver(
+            position[0], position[1], position[2],
+            rotation_matrix[0, j], rotation_matrix[1, j], rotation_matrix[2, j],
+            color=c, length=length, normalize=True, linewidth=linewidth
+        )
+
+
+def plot_se3_trajectory(
+    positions: np.ndarray,
+    quaternions: Optional[np.ndarray] = None,
+    ax: Optional[Any] = None,
+    title: str = "SE(3) trajectory",
+    show_orientation: bool = False,
+    vis_type: str = "arrow",
+    box_size: tuple = (0.03, 0.03, 0.03),
+    color: str = "b",
+    alpha: float = 0.7,
+    num_frames: int = 8,
+    label: Optional[str] = None,
+) -> Any:
+    """
+    Plot 3D position trajectory with optional orientation visualization.
+
+    Args:
+        positions: (N, 3) position array
+        quaternions: (N, 4) quaternion array (wxyz), optional
+        ax: Matplotlib 3D axes, created if None
+        title: Plot title
+        show_orientation: If True, show orientation frames or cubes
+        vis_type: "arrow" for coordinate frames, "cube" for oriented boxes
+        box_size: Size of cubes (if vis_type="cube")
+        color: Trajectory line color
+        alpha: Transparency for cubes
+        num_frames: Number of orientation samples to show along trajectory
+        label: Legend label for trajectory
+
+    Returns:
+        Matplotlib axes
+    """
+    if not HAS_MPL:
+        raise ImportError("matplotlib required for plot_se3_trajectory")
+    positions = np.asarray(positions)
+    if ax is None:
+        fig = plt.figure()
+        ax = fig.add_subplot(111, projection="3d")
+    ax.plot(positions[:, 0], positions[:, 1], positions[:, 2], "-", color=color, label=label or "position", alpha=0.8)
+    ax.set_xlabel("x")
+    ax.set_ylabel("y")
+    ax.set_zlabel("z")
+    ax.set_title(title)
+
+    if show_orientation and quaternions is not None:
+        from dhb_xr.core import geometry as geom
+        n = len(positions)
+        indices = np.linspace(0, n - 1, num_frames, dtype=int)
+        for i in indices:
+            R = geom.quat_to_rot(quaternions[i])
+            p = positions[i]
+            if vis_type == "cube":
+                draw_box(ax, p, R, size=box_size, color=color, alpha=alpha)
+            else:  # arrow
+                draw_frame(ax, p, R, length=0.05)
+    return ax
+
+
+def plot_se3_trajectories(
+    trajectories: Union[List[Dict], Dict[str, Dict]],
+    ax: Optional[Any] = None,
+    show_orientation: bool = True,
+    vis_type: str = "cube",
+    box_size_scale: float = 0.05,
+    num_frames: int = 6,
+    title: str = "SE(3) Trajectories",
+    show_legend: bool = True,
+) -> Any:
+    """
+    Plot multiple SE(3) trajectories with orientation visualization.
+
+    Args:
+        trajectories: List of dicts with 'positions' and 'quaternions', or dict of name -> trajectory
+        ax: Matplotlib 3D axes
+        show_orientation: Show orientation cubes/frames
+        vis_type: "cube" or "arrow"
+        box_size_scale: Size of cubes
+        num_frames: Number of orientation samples per trajectory
+        title: Plot title
+        show_legend: Show legend
+
+    Returns:
+        Matplotlib axes
+    """
+    if not HAS_MPL:
+        raise ImportError("matplotlib required for plot_se3_trajectories")
+
+    # Normalize input to dict
+    if isinstance(trajectories, list):
+        trajectories = {f"traj_{i}": t for i, t in enumerate(trajectories)}
+
+    if ax is None:
+        fig = plt.figure(figsize=(10, 8))
+        ax = fig.add_subplot(111, projection="3d")
+
+    colors = ["blue", "green", "red", "orange", "purple", "cyan", "magenta", "brown"]
+    box_size = (box_size_scale, box_size_scale, box_size_scale)
+
+    for idx, (name, traj) in enumerate(trajectories.items()):
+        positions = np.asarray(traj["positions"])
+        quaternions = np.asarray(traj.get("quaternions"))
+        color = colors[idx % len(colors)]
+
+        ax.plot(positions[:, 0], positions[:, 1], positions[:, 2], color=color, label=name, linewidth=2)
+
+        if show_orientation and quaternions is not None:
+            from dhb_xr.core import geometry as geom
+            n = len(positions)
+            indices = np.linspace(0, n - 1, num_frames, dtype=int)
+            for i in indices:
+                R = geom.quat_to_rot(quaternions[i])
+                p = positions[i]
+                if vis_type == "cube":
+                    draw_box(ax, p, R, size=box_size, color=color, alpha=0.5)
+                else:
+                    draw_frame(ax, p, R, length=0.05)
+
+    ax.set_xlabel("x")
+    ax.set_ylabel("y")
+    ax.set_zlabel("z")
+    ax.set_title(title)
+    if show_legend:
+        ax.legend()
+
+    return ax
+
+
+def plot_invariants(
+    linear_inv: np.ndarray,
+    angular_inv: np.ndarray,
+    ax: Optional[Any] = None,
+    title: str = "DHB invariants",
+) -> Any:
+    """Plot linear and angular invariant time series."""
+    if not HAS_MPL:
+        raise ImportError("matplotlib required for plot_invariants")
+    n = linear_inv.shape[0]
+    k_lin, k_ang = linear_inv.shape[1], angular_inv.shape[1]
+    if ax is None:
+        fig, axes = plt.subplots(2, 1, sharex=True)
+        ax_lin, ax_ang = axes[0], axes[1]
+    else:
+        ax_lin, ax_ang = ax
+    t = np.arange(n)
+    for j in range(k_lin):
+        ax_lin.plot(t, linear_inv[:, j], label=f"lin_{j}")
+    ax_lin.set_ylabel("linear")
+    ax_lin.legend(loc="right", fontsize=8)
+    ax_lin.set_title(title)
+    for j in range(k_ang):
+        ax_ang.plot(t, angular_inv[:, j], label=f"ang_{j}")
+    ax_ang.set_ylabel("angular")
+    ax_ang.legend(loc="right", fontsize=8)
+    ax_ang.set_xlabel("step")
+    return (ax_lin, ax_ang)