Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""HandFlow feature extractors package."""
|
||||
@@ -0,0 +1,67 @@
|
||||
"""
|
||||
Condition Builder (with cmask mechanism).
|
||||
|
||||
Concatenates image token + skeleton token into the condition stream.
|
||||
cmask: learnable mask token; both training and inference use HaMeR confidence as a continuous mask.
|
||||
|
||||
Inputs:
|
||||
- image_tokens: (B, T, d_model) — from Image Refiner + Compression
|
||||
- skeleton_tokens: (B, T, d_model) — from Skeleton Encoder
|
||||
- hamer_confidence: (B, T) — HaMeR detection confidence
|
||||
|
||||
Output: condition (B, 2*T, d_model)
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class ConditionBuilder(nn.Module):
|
||||
def __init__(self, d_model: int, mask_ratio: float = 0.2):
|
||||
super().__init__()
|
||||
self.d_model = d_model
|
||||
self.mask_ratio = mask_ratio
|
||||
|
||||
# Learnable mask token — initialized with truncated normal
|
||||
self.c_mask = nn.Parameter(torch.empty(1, 1, d_model))
|
||||
nn.init.trunc_normal_(self.c_mask, std=0.02)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
image_tokens: torch.Tensor,
|
||||
skeleton_tokens: torch.Tensor,
|
||||
hamer_confidence: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Args:
|
||||
image_tokens: (B, T, d_model)
|
||||
skeleton_tokens: (B, T, d_model)
|
||||
hamer_confidence: (B, T) — used as m_t during inference
|
||||
|
||||
Returns:
|
||||
condition: (B, 2*T, d_model)
|
||||
m: (B, T) — actual mask values used (for logging)
|
||||
"""
|
||||
B, T, D = image_tokens.shape
|
||||
|
||||
# concat: (B, T, 2, d_model) -> (B, 2*T, d_model)
|
||||
tokens = torch.stack([image_tokens, skeleton_tokens], dim=2) # (B, T, 2, D)
|
||||
tokens = tokens.reshape(B, 2 * T, D)
|
||||
|
||||
# Use HaMeR confidence as a continuous mask (consistent between training and inference)
|
||||
m = hamer_confidence.float() # (B, T)
|
||||
|
||||
# Optional: randomly dropout some frames during training as regularization
|
||||
if self.training and self.mask_ratio > 0:
|
||||
keep = (torch.rand(B, T, device=tokens.device) >= self.mask_ratio).float()
|
||||
m = m * keep
|
||||
|
||||
# Expand m to match 2*T tokens (the two tokens per frame share the same m)
|
||||
# m: (B, T) -> (B, 2*T)
|
||||
m_expanded = m.unsqueeze(-1).expand(-1, -1, 2).reshape(B, 2 * T) # (B, 2*T)
|
||||
m_expanded = m_expanded.unsqueeze(-1) # (B, 2*T, 1)
|
||||
|
||||
# cmask: c'_t = m_t * c_t + (1 - m_t) * c_mask
|
||||
condition = m_expanded * tokens + (1 - m_expanded) * self.c_mask
|
||||
|
||||
return condition, m
|
||||
@@ -0,0 +1,110 @@
|
||||
"""
|
||||
HaMeR ViT-Huge backbone wrapper (frozen).
|
||||
|
||||
Provides image token extraction for the 2D condition stream.
|
||||
Input: (B, 3, 256, 256) float32 [0,1]
|
||||
Output: (B, 192, 1280) — 16x12 patches x 1280D
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
# Add HaMeR repository to the import path
|
||||
_HAMER_ROOT = Path(__file__).resolve().parents[2] / "third_party" / "hamer"
|
||||
if _HAMER_ROOT.exists() and str(_HAMER_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(_HAMER_ROOT))
|
||||
|
||||
|
||||
class HaMeRBackbone(nn.Module):
|
||||
"""
|
||||
Frozen HaMeR ViT-Huge backbone.
|
||||
|
||||
Internal pipeline:
|
||||
1. ImageNet normalization (mean/std)
|
||||
2. Crop 32px from each horizontal side: (B,3,256,256) -> (B,3,256,192)
|
||||
3. Patch embedding: 192/16 x 256/16 = 12x16 = 192 tokens x 1280D
|
||||
4. 32-layer ViT-H Transformer
|
||||
5. Reshape output -> (B, 192, 1280)
|
||||
"""
|
||||
|
||||
# ImageNet normalization parameters
|
||||
_IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406])
|
||||
_IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225])
|
||||
|
||||
def __init__(self, checkpoint_path: Optional[str] = None):
|
||||
super().__init__()
|
||||
# Import the ViT backbone directly, bypassing hamer/models/__init__.py (avoids smplx dependency)
|
||||
import importlib.util
|
||||
vit_path = _HAMER_ROOT / "hamer" / "models" / "backbones" / "vit.py"
|
||||
spec = importlib.util.spec_from_file_location("hamer_vit", str(vit_path))
|
||||
vit_mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(vit_mod)
|
||||
ViT = vit_mod.ViT
|
||||
|
||||
self.backbone = ViT(
|
||||
img_size=(256, 192),
|
||||
patch_size=16,
|
||||
embed_dim=1280,
|
||||
depth=32,
|
||||
num_heads=16,
|
||||
ratio=1,
|
||||
use_checkpoint=False,
|
||||
mlp_ratio=4,
|
||||
qkv_bias=True,
|
||||
drop_path_rate=0.55,
|
||||
)
|
||||
self.backbone.eval()
|
||||
|
||||
# Load pretrained weights
|
||||
if checkpoint_path is not None:
|
||||
ckpt = torch.load(checkpoint_path, map_location="cpu")
|
||||
if "state_dict" in ckpt:
|
||||
state_dict = {
|
||||
k.replace("backbone.", ""): v
|
||||
for k, v in ckpt["state_dict"].items()
|
||||
if k.startswith("backbone.")
|
||||
}
|
||||
else:
|
||||
state_dict = ckpt
|
||||
self.backbone.load_state_dict(state_dict, strict=False)
|
||||
|
||||
# Freeze all parameters
|
||||
self.backbone.requires_grad_(False)
|
||||
for p in self.backbone.parameters():
|
||||
p.requires_grad = False
|
||||
|
||||
# Register normalization buffers
|
||||
self.register_buffer("_mean", self._IMAGENET_MEAN.view(1, 3, 1, 1))
|
||||
self.register_buffer("_std", self._IMAGENET_STD.view(1, 3, 1, 1))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
x: (B, 3, 256, 256) float32 in range [0, 1]
|
||||
Returns:
|
||||
tokens: (B, 192, 1280) patch tokens
|
||||
"""
|
||||
# ImageNet normalization
|
||||
x = (x - self._mean) / self._std
|
||||
|
||||
# HaMeR internal crop: 32px from each side (256 -> 192)
|
||||
x = x[:, :, :, 32:-32]
|
||||
|
||||
# Backbone forward
|
||||
feat_map = self.backbone.forward_features(x) # (B, 1280, Hp, Wp)
|
||||
B, C, Hp, Wp = feat_map.shape
|
||||
|
||||
# Flatten into tokens
|
||||
tokens = feat_map.flatten(2).transpose(1, 2) # (B, Hp*Wp, 1280)
|
||||
|
||||
return tokens
|
||||
|
||||
def train(self, mode: bool = True):
|
||||
"""Always stay in eval mode."""
|
||||
super().train(False)
|
||||
self.backbone.eval()
|
||||
return self
|
||||
@@ -0,0 +1,282 @@
|
||||
"""
|
||||
Image Spatio-Temporal Refiner + Frame Compression.
|
||||
|
||||
Step 1: 8+8 alternating spatio-temporal attention over patch tokens of all frames
|
||||
Step 2: 4 layers of cross-attention pooling compress 192 tokens per frame -> 1 token
|
||||
|
||||
Input: HaMeR backbone output (B, T, 192, 1280)
|
||||
Output: (B, T, d_model) — one condensed image condition token per frame
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import rearrange
|
||||
from model import attention
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# 2D RoPE (consistent with the legacy TemporalRefiner)
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class RotaryEmbedding1D(nn.Module):
|
||||
"""1D RoPE: used for temporal positional encoding."""
|
||||
|
||||
def __init__(self, dim: int, base: int = 10000):
|
||||
super().__init__()
|
||||
assert dim % 2 == 0
|
||||
self.dim = dim
|
||||
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
|
||||
def forward(self, x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
x: (B, H, L, D), positions: (L,)
|
||||
"""
|
||||
freqs = torch.einsum("i,j->ij", positions.float(), self.inv_freq.to(positions))
|
||||
sin = freqs.sin().unsqueeze(0).unsqueeze(0)
|
||||
cos = freqs.cos().unsqueeze(0).unsqueeze(0)
|
||||
x1, x2 = x[..., 0::2], x[..., 1::2]
|
||||
return torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).flatten(-2)
|
||||
|
||||
|
||||
class RotaryEmbedding2D(nn.Module):
|
||||
"""2D RoPE: split head_dim in half and apply 1D RoPE with y/x coordinates respectively."""
|
||||
|
||||
def __init__(self, dim: int, base: int = 10000):
|
||||
super().__init__()
|
||||
assert dim % 2 == 0
|
||||
self.dim = dim
|
||||
self.dim_per_axis = dim // 2
|
||||
inv_freq = 1.0 / (base ** (torch.arange(0, self.dim_per_axis, 2).float() / self.dim_per_axis))
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
|
||||
def forward(self, x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
x: (B, H, L, D), positions: (L, 2) (y, x)
|
||||
"""
|
||||
x_y, x_x = x.chunk(2, dim=-1)
|
||||
x_y = self._apply_1d(x_y, positions[:, 0])
|
||||
x_x = self._apply_1d(x_x, positions[:, 1])
|
||||
return torch.cat([x_y, x_x], dim=-1)
|
||||
|
||||
def _apply_1d(self, x: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
|
||||
freqs = torch.einsum("i,j->ij", pos.float(), self.inv_freq.to(pos))
|
||||
sin = freqs.sin().unsqueeze(0).unsqueeze(0)
|
||||
cos = freqs.cos().unsqueeze(0).unsqueeze(0)
|
||||
x1, x2 = x[..., 0::2], x[..., 1::2]
|
||||
return torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).flatten(-2)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Building blocks
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, dim: int):
|
||||
super().__init__()
|
||||
self.scale = nn.Parameter(torch.ones(dim))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x_dtype = x.dtype
|
||||
x = x.float()
|
||||
rrms = torch.rsqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + 1e-6)
|
||||
return (x * rrms).to(dtype=x_dtype) * self.scale
|
||||
|
||||
|
||||
class RefinerBlock(nn.Module):
|
||||
"""Pre-LN Transformer block with optional 2D RoPE."""
|
||||
|
||||
def __init__(self, d_model: int, num_heads: int, mlp_ratio: float = 4.0, dropout: float = 0.1):
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = d_model // num_heads
|
||||
self.norm1 = nn.LayerNorm(d_model)
|
||||
self.qkv = nn.Linear(d_model, d_model * 3)
|
||||
self.q_norm = RMSNorm(self.head_dim)
|
||||
self.k_norm = RMSNorm(self.head_dim)
|
||||
self.proj = nn.Linear(d_model, d_model)
|
||||
self.norm2 = nn.LayerNorm(d_model)
|
||||
mlp_hidden = int(d_model * mlp_ratio)
|
||||
self.mlp = nn.Sequential(nn.Linear(d_model, mlp_hidden), nn.GELU(), nn.Linear(mlp_hidden, d_model))
|
||||
self.drop = nn.Dropout(dropout)
|
||||
|
||||
def forward(self, x: torch.Tensor, rope: Optional[RotaryEmbedding2D] = None, positions: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||||
residual = x
|
||||
x = self.norm1(x)
|
||||
qkv = self.qkv(x)
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
|
||||
q = self.q_norm(q)
|
||||
k = self.k_norm(k)
|
||||
if rope is not None and positions is not None:
|
||||
q = rope(q, positions)
|
||||
k = rope(k, positions)
|
||||
x = attention(q, k, v)
|
||||
x = self.proj(x)
|
||||
x = residual + self.drop(x)
|
||||
x = x + self.drop(self.mlp(self.norm2(x)))
|
||||
return x
|
||||
|
||||
|
||||
class CrossAttentionPoolBlock(nn.Module):
|
||||
"""Cross-attention pooling: 1 learnable query token cross-attends K patches."""
|
||||
|
||||
def __init__(self, d_model: int, num_heads: int, mlp_ratio: float = 4.0, dropout: float = 0.1):
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = d_model // num_heads
|
||||
self.norm_q = nn.LayerNorm(d_model)
|
||||
self.norm_kv = nn.LayerNorm(d_model)
|
||||
self.q_proj = nn.Linear(d_model, d_model)
|
||||
self.kv_proj = nn.Linear(d_model, d_model * 2)
|
||||
self.q_norm = RMSNorm(self.head_dim)
|
||||
self.k_norm = RMSNorm(self.head_dim)
|
||||
self.out_proj = nn.Linear(d_model, d_model)
|
||||
self.norm_ff = nn.LayerNorm(d_model)
|
||||
mlp_hidden = int(d_model * mlp_ratio)
|
||||
self.ff = nn.Sequential(nn.Linear(d_model, mlp_hidden), nn.GELU(), nn.Linear(mlp_hidden, d_model))
|
||||
self.drop = nn.Dropout(dropout)
|
||||
|
||||
def forward(self, query: torch.Tensor, patches: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
query: (B, 1, D), patches: (B, N, D)
|
||||
Returns: (B, 1, D)
|
||||
"""
|
||||
residual = query
|
||||
q = self.q_proj(self.norm_q(query))
|
||||
kv = self.kv_proj(self.norm_kv(patches))
|
||||
k, v = rearrange(kv, "B N (K H D) -> K B H N D", K=2, H=self.num_heads)
|
||||
q = rearrange(q, "B L (H D) -> B H L D", H=self.num_heads)
|
||||
q = self.q_norm(q)
|
||||
k = self.k_norm(k)
|
||||
out = attention(q, k, v) # (B, 1, H*D) = (B, 1, D)
|
||||
out = self.out_proj(out)
|
||||
query = residual + self.drop(out)
|
||||
query = query + self.drop(self.ff(self.norm_ff(query)))
|
||||
return query
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Spatio-Temporal Refiner
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class SpatioTemporalRefiner(nn.Module):
|
||||
"""
|
||||
8+8 alternating spatio-temporal attention.
|
||||
|
||||
Even layers: intra-frame spatial attention (B*T, 192, dim)
|
||||
Odd layers: cross-frame same-position attention (B, T, dim) per spatial position
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
d_model: int,
|
||||
hamer_dim: int = 1280,
|
||||
spatial_layers: int = 8,
|
||||
temporal_layers: int = 8,
|
||||
num_heads: int = 8,
|
||||
mlp_ratio: float = 4.0,
|
||||
dropout: float = 0.1,
|
||||
n_patches_h: int = 16,
|
||||
n_patches_w: int = 12,
|
||||
):
|
||||
super().__init__()
|
||||
self.d_model = d_model
|
||||
self.n_patches = n_patches_h * n_patches_w
|
||||
self.depth = spatial_layers + temporal_layers
|
||||
self.spatial_layers = spatial_layers
|
||||
self.temporal_layers = temporal_layers
|
||||
|
||||
self.input_proj = nn.Linear(hamer_dim, d_model)
|
||||
|
||||
# 2D spatial positions
|
||||
y = torch.arange(n_patches_h)
|
||||
x = torch.arange(n_patches_w)
|
||||
self.register_buffer("spatial_positions", torch.cartesian_prod(y, x))
|
||||
|
||||
self.rope = RotaryEmbedding2D(dim=d_model // num_heads)
|
||||
self.temporal_rope = RotaryEmbedding1D(dim=d_model // num_heads)
|
||||
self.blocks = nn.ModuleList([
|
||||
RefinerBlock(d_model, num_heads, mlp_ratio, dropout)
|
||||
for _ in range(self.depth)
|
||||
])
|
||||
self.use_checkpoint = False
|
||||
|
||||
def gradient_checkpointing_enable(self):
|
||||
self.use_checkpoint = True
|
||||
|
||||
def forward(self, hamer_features: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
hamer_features: (B, T, 192, 1280)
|
||||
Returns:
|
||||
(B, T, 192, d_model)
|
||||
"""
|
||||
B, T, N, _ = hamer_features.shape
|
||||
x = self.input_proj(hamer_features)
|
||||
|
||||
spatial_pos = self.spatial_positions # (N, 2)
|
||||
temporal_pos = torch.arange(T, device=x.device) # (T,)
|
||||
temporal_pos_expanded = temporal_pos.unsqueeze(1).expand(T, N).reshape(T * N) # (T*N,)
|
||||
|
||||
for i, blk in enumerate(self.blocks):
|
||||
if i % 2 == 0:
|
||||
# Intra-frame spatial attention (2D RoPE: y, x)
|
||||
x = x.reshape(B * T, N, self.d_model)
|
||||
x = blk(x, rope=self.rope, positions=spatial_pos)
|
||||
x = x.reshape(B, T, N, self.d_model)
|
||||
else:
|
||||
# Cross-frame global attention (1D temporal RoPE: frame index t)
|
||||
x = x.reshape(B, T * N, self.d_model)
|
||||
x = blk(x, rope=self.temporal_rope, positions=temporal_pos_expanded)
|
||||
x = x.reshape(B, T, N, self.d_model)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Frame Compressor (Cross-Attention Pooling)
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class FrameCompressor(nn.Module):
|
||||
"""
|
||||
Compress N patch tokens of each frame into 1 condensed token.
|
||||
Uses a learnable query token with multi-layer cross-attention.
|
||||
Includes input_proj: supports arbitrary input dims (e.g. HaMeR 1280D -> d_model).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
d_model: int,
|
||||
input_dim: int | None = None,
|
||||
num_heads: int = 8,
|
||||
num_layers: int = 8,
|
||||
mlp_ratio: float = 4.0,
|
||||
dropout: float = 0.1,
|
||||
):
|
||||
super().__init__()
|
||||
self.input_proj = nn.Linear(input_dim, d_model) if input_dim else nn.Identity()
|
||||
self.query = nn.Parameter(torch.randn(1, 1, d_model))
|
||||
nn.init.trunc_normal_(self.query, std=0.02)
|
||||
self.layers = nn.ModuleList([
|
||||
CrossAttentionPoolBlock(d_model, num_heads, mlp_ratio, dropout)
|
||||
for _ in range(num_layers)
|
||||
])
|
||||
|
||||
def forward(self, patches: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
patches: (B, T, N, input_dim or d_model)
|
||||
Returns: (B, T, d_model)
|
||||
"""
|
||||
B, T, N, _ = patches.shape
|
||||
patches = self.input_proj(patches) # (B, T, N, d_model)
|
||||
patches = patches.reshape(B * T, N, -1) # (B*T, N, D)
|
||||
q = self.query.expand(B * T, -1, -1) # (B*T, 1, D)
|
||||
for layer in self.layers:
|
||||
q = layer(q, patches)
|
||||
return q.reshape(B, T, -1)
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user