Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-09-15 14:13:49 +08:00
commit 7e4ef6f98b
7002 changed files with 1452249 additions and 0 deletions
@@ -0,0 +1,96 @@
+"""
+Ray Direction Skeleton Encoder.
+
+2D skeleton -> inverse intrinsics -> ray direction -> high-frequency sinusoidal PE -> flatten -> MLP -> embed_dim.
+Outputs one skeleton token per frame.
+
+Input: landmarks (B, T, 21, 2) + crop_intrinsics (B, T, 4)
+Output: (B, T, embed_dim)
+"""
+
+import torch
+import torch.nn as nn
+
+
+class RayDirectionSkeletonEncoder(nn.Module):
+    """
+    Encodes a 2D skeleton into one token per frame via ray direction.
+
+    Pipeline:
+      1. un-normalize landmarks -> pixel coordinates (in crop coordinate system)
+      2. ray = [(u-cx)/fx, (v-cy)/fy, 1] -> normalize -> (21, 3)
+      3. High-frequency sinusoidal PE: [sin(2^L * ray), cos(2^L * ray)] L=0..num_freqs-1
+      4. flatten: 21 joints x (3 x 2 x num_freqs) per joint
+      5. MLP -> embed_dim
+    """
+
+    def __init__(
+        self,
+        n_joints: int = 21,
+        num_freqs: int = 6,
+        embed_dim: int = 512,
+        crop_h: int = 256,
+        crop_w: int = 256,
+        mlp_hidden: int = 512,
+    ):
+        super().__init__()
+        self.n_joints = n_joints
+        self.num_freqs = num_freqs
+        self.crop_h = crop_h
+        self.crop_w = crop_w
+
+        # Per-joint dim after PE: 3 (ray) x 2 (sin+cos) x num_freqs
+        pe_dim = 3 * 2 * num_freqs  # = 36
+        flat_dim = n_joints * pe_dim  # = 756
+
+        self.mlp = nn.Sequential(
+            nn.Linear(flat_dim, mlp_hidden),
+            nn.SiLU(),
+            nn.Linear(mlp_hidden, embed_dim),
+        )
+
+    def forward(
+        self,
+        landmarks: torch.Tensor,
+        crop_intrinsics: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Args:
+            landmarks: (B, T, 21, 2) normalized [0, 1]
+            crop_intrinsics: (B, T, 4) = [fx, fy, cx, cy] in crop coords
+        Returns:
+            (B, T, embed_dim)
+        """
+        B, T, J, _ = landmarks.shape
+        fx = crop_intrinsics[..., 0:1]  # (B, T, 1)
+        fy = crop_intrinsics[..., 1:2]
+        cx = crop_intrinsics[..., 2:3]
+        cy = crop_intrinsics[..., 3:4]
+
+        # un-normalize -> crop pixel coordinates
+        u = landmarks[..., 0:1] * self.crop_w  # (B, T, 21, 1)
+        v = landmarks[..., 1:2] * self.crop_h
+
+        # ray direction: [(u-cx)/fx, (v-cy)/fy, 1]
+        ray_x = (u - cx.unsqueeze(-2)) / (fx.unsqueeze(-2) + 1e-8)
+        ray_y = (v - cy.unsqueeze(-2)) / (fy.unsqueeze(-2) + 1e-8)
+        ray_z = torch.ones_like(ray_x)
+        ray = torch.cat([ray_x, ray_y, ray_z], dim=-1)  # (B, T, 21, 3)
+
+        # normalize ray direction
+        ray = ray / (ray.norm(dim=-1, keepdim=True) + 1e-8)
+
+        # High-frequency sinusoidal PE
+        pe_parts = []
+        for L in range(self.num_freqs):
+            freq = 2.0 ** L
+            pe_parts.append(torch.sin(freq * ray))
+            pe_parts.append(torch.cos(freq * ray))
+        # pe: (B, T, 21, 3 * 2 * num_freqs) = (B, T, 21, 36)
+        pe = torch.cat(pe_parts, dim=-1)
+
+        # flatten all joints: (B, T, 21 * 36) = (B, T, 756)
+        flat = pe.reshape(B, T, -1)
+
+        # MLP -> embed_dim
+        return self.mlp(flat)