346 lines
14 KiB
Python
346 lines
14 KiB
Python
"""
|
||
HandFlow Rectified Flow Denoiser — all-frame direct regression.
|
||
|
||
Pose (48D axis-angle) and Trans (3D) are both regressed per-frame directly,
|
||
without velocity/cumsum. Beta (10D) serves as a global token.
|
||
|
||
Flow Target layout (flat packing per window, managed by mano_slice):
|
||
x_1 = [beta(10) | pose_0(48) | ... | pose_{T-1}(48) | trans_0(3) | ... | trans_{T-1}(3)]
|
||
flat_dim = 10 + T × 51
|
||
|
||
z-stream: 1 beta token + T pose tokens + T trans tokens = 1+2T tokens
|
||
c-stream: 2 × T tokens (1 image + 1 skeleton per frame, with cmask)
|
||
|
||
Condition pipeline:
|
||
HaMeR backbone (frozen) → FrameCompressor (8-layer cross-attention pooling) → image tokens
|
||
2D skeleton → ray direction → PE → MLP → skeleton tokens
|
||
image + skeleton → cmask → condition stream
|
||
"""
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
from omegaconf import DictConfig
|
||
|
||
from model.flow_matching.dualstream_transformer import (
|
||
FlowMatchingTransformer,
|
||
MLPEmbedder,
|
||
)
|
||
from model.flow_matching.helpers.model_wrapper import ModelWrapper
|
||
from model.flow_matching.helpers.scheduler import FluxTimeSampler
|
||
from model.feature_extractors.hamer_extractor import HaMeRBackbone
|
||
from model.feature_extractors.image_refiner import FrameCompressor
|
||
from model.feature_extractors.skeleton_encoder import RayDirectionSkeletonEncoder
|
||
from model.feature_extractors.condition_builder import ConditionBuilder
|
||
from model.pose_embedding import RotaryEmbedding
|
||
|
||
|
||
class ManoSlice(nn.Module):
|
||
"""Manage slice positions and normalization of each component in the flow target flat vector.
|
||
|
||
Flat layout: [beta(10) | poses(T,48) | trans(T,3)]
|
||
flat_dim = 10 + T * 51
|
||
"""
|
||
|
||
def __init__(self, beta_dim: int, pose_dim: int,
|
||
trans_dim: int, window_size: int, stats_path: str | None = None):
|
||
super().__init__()
|
||
self.beta_dim = beta_dim # 10
|
||
self.pose_dim = pose_dim # 48
|
||
self.trans_dim = trans_dim # 3
|
||
self.window_size = window_size
|
||
|
||
self.flat_dim = beta_dim + window_size * (pose_dim + trans_dim)
|
||
|
||
# Build slices
|
||
s = 0
|
||
self.beta = slice(s, s + beta_dim) # [0:10]
|
||
s += beta_dim
|
||
self.poses = slice(s, s + window_size * pose_dim) # [10:10+48T]
|
||
s += window_size * pose_dim
|
||
self.trans = slice(s, s + window_size * trans_dim) # [10+48T:10+51T]
|
||
|
||
# Normalization statistics
|
||
flat_mean = torch.zeros(self.flat_dim)
|
||
flat_std = torch.ones(self.flat_dim)
|
||
if stats_path is not None:
|
||
import numpy as np
|
||
stats = np.load(stats_path)
|
||
beta_mean = torch.from_numpy(stats["beta_mean"]) # (10,)
|
||
beta_std = torch.from_numpy(stats["beta_std"]) # (10,)
|
||
pose_mean = torch.from_numpy(stats["pose_mean"]) # (48,)
|
||
pose_std = torch.from_numpy(stats["pose_std"]) # (48,)
|
||
trans_mean = torch.from_numpy(stats["trans_mean"]) # (3,)
|
||
trans_std = torch.from_numpy(stats["trans_std"]) # (3,)
|
||
flat_mean = torch.cat([
|
||
beta_mean,
|
||
pose_mean.repeat(window_size),
|
||
trans_mean.repeat(window_size),
|
||
])
|
||
flat_std = torch.cat([
|
||
beta_std,
|
||
pose_std.repeat(window_size),
|
||
trans_std.repeat(window_size),
|
||
])
|
||
self.register_buffer("flat_mean", flat_mean)
|
||
self.register_buffer("flat_std", flat_std)
|
||
|
||
# per-component statistics for auxiliary loss denormalization
|
||
self.register_buffer("beta_mean_", flat_mean[self.beta])
|
||
self.register_buffer("beta_std_", flat_std[self.beta])
|
||
|
||
if stats_path is not None:
|
||
self.register_buffer("pose_mean_", pose_mean)
|
||
self.register_buffer("pose_std_", pose_std)
|
||
self.register_buffer("trans_mean_", trans_mean)
|
||
self.register_buffer("trans_std_", trans_std)
|
||
else:
|
||
self.register_buffer("pose_mean_", torch.zeros(pose_dim))
|
||
self.register_buffer("pose_std_", torch.ones(pose_dim))
|
||
self.register_buffer("trans_mean_", torch.zeros(trans_dim))
|
||
self.register_buffer("trans_std_", torch.ones(trans_dim))
|
||
|
||
def normalize(self, x_flat: torch.Tensor) -> torch.Tensor:
|
||
"""Normalize the flat vector to ~N(0,1)."""
|
||
return (x_flat - self.flat_mean) / self.flat_std
|
||
|
||
def denormalize(self, x_flat: torch.Tensor) -> torch.Tensor:
|
||
"""Denormalize the flat vector back to original scale."""
|
||
return x_flat * self.flat_std + self.flat_mean
|
||
|
||
def denormalize_components(
|
||
self, beta: torch.Tensor, poses: torch.Tensor, trans: torch.Tensor,
|
||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||
"""Denormalize each component back to original scale."""
|
||
beta_orig = beta * self.beta_std_ + self.beta_mean_
|
||
poses_orig = poses * self.pose_std_ + self.pose_mean_ # (48,) broadcast
|
||
trans_orig = trans * self.trans_std_ + self.trans_mean_ # (3,) broadcast
|
||
return beta_orig, poses_orig, trans_orig
|
||
|
||
def unpack(self, x_flat: torch.Tensor):
|
||
"""Unpack flat → (beta, poses, trans)"""
|
||
B = x_flat.shape[0]
|
||
beta = x_flat[:, self.beta] # (B, 10)
|
||
poses = x_flat[:, self.poses].reshape(B, self.window_size, self.pose_dim) # (B, T, 48)
|
||
trans = x_flat[:, self.trans].reshape(B, self.window_size, self.trans_dim) # (B, T, 3)
|
||
return beta, poses, trans
|
||
|
||
def pack(self, beta: torch.Tensor, poses: torch.Tensor,
|
||
trans: torch.Tensor) -> torch.Tensor:
|
||
"""Pack → flat"""
|
||
B = beta.shape[0]
|
||
return torch.cat([
|
||
beta.reshape(B, -1),
|
||
poses.reshape(B, -1),
|
||
trans.reshape(B, -1),
|
||
], dim=-1)
|
||
|
||
|
||
class HandPoseDenoiser(nn.Module):
|
||
"""
|
||
Rectified Flow Denoiser — all-frame direct regression.
|
||
"""
|
||
|
||
def __init__(self, config: DictConfig):
|
||
super().__init__()
|
||
fm_cfg = config.model.fm
|
||
cond_cfg = fm_cfg.condition
|
||
comp_cfg = fm_cfg.frame_compressor
|
||
tf_cfg = fm_cfg.transformer
|
||
|
||
self.d_model = tf_cfg.hidden_size
|
||
self.beta_dim = fm_cfg.beta_dim # 10
|
||
self.pose_dim = fm_cfg.pose_dim # 48
|
||
self.trans_dim = fm_cfg.trans_dim # 3
|
||
self.window_size = config.dataset.window_size
|
||
|
||
# mano_slice (includes normalization statistics)
|
||
stats_path = fm_cfg.get("normalization_stats", None)
|
||
self.mano_slice = ManoSlice(
|
||
beta_dim=self.beta_dim,
|
||
pose_dim=self.pose_dim,
|
||
trans_dim=self.trans_dim,
|
||
window_size=self.window_size,
|
||
stats_path=stats_path,
|
||
)
|
||
|
||
# z-stream input projections: each component → d_model
|
||
self.beta_proj_in = nn.Linear(self.beta_dim, self.d_model)
|
||
self.pose_proj_in = nn.Linear(self.pose_dim, self.d_model)
|
||
self.trans_proj_in = nn.Linear(self.trans_dim, self.d_model)
|
||
|
||
# z-stream output projections: d_model → native dims
|
||
self.beta_proj_out = nn.Linear(self.d_model, self.beta_dim)
|
||
self.pose_proj_out = nn.Linear(self.d_model, self.pose_dim)
|
||
self.trans_proj_out = nn.Linear(self.d_model, self.trans_dim)
|
||
|
||
# HaMeR backbone (frozen)
|
||
hamer_ckpt = config.model.get("hamer", {}).get("checkpoint", None)
|
||
self.hamer_backbone = HaMeRBackbone(checkpoint_path=hamer_ckpt)
|
||
|
||
# Frame Compressor (directly compresses HaMeR patch tokens)
|
||
self.frame_compressor = FrameCompressor(
|
||
d_model=self.d_model,
|
||
input_dim=cond_cfg.hamer_embed_dim,
|
||
num_heads=tf_cfg.num_heads,
|
||
num_layers=comp_cfg.num_layers,
|
||
)
|
||
|
||
# Skeleton Encoder
|
||
self.skeleton_encoder = RayDirectionSkeletonEncoder(
|
||
n_joints=cond_cfg.n_joints,
|
||
num_freqs=cond_cfg.skeleton_pe_freqs,
|
||
embed_dim=self.d_model,
|
||
)
|
||
|
||
# Condition Builder (cmask; mask_ratio only used for optional regularization)
|
||
self.condition_builder = ConditionBuilder(
|
||
d_model=self.d_model,
|
||
mask_ratio=cond_cfg.get("cmask_ratio", 0.0),
|
||
)
|
||
|
||
# RoPE for Denoiser
|
||
self.rope = RotaryEmbedding(dim=self.d_model // tf_cfg.num_heads)
|
||
|
||
# Flow Matching Transformer
|
||
self.transformer = FlowMatchingTransformer(
|
||
in_channels=self.d_model,
|
||
out_channels=self.d_model,
|
||
use_context_in=True,
|
||
context_in_dim=self.d_model,
|
||
use_txt_in=False,
|
||
vec_in_dim=0,
|
||
use_pre_text_attn=False,
|
||
config=tf_cfg,
|
||
)
|
||
|
||
def _extract_image_features(self, images: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
Image feature extraction: HaMeR backbone → Frame Compressor.
|
||
images: (B, T, 3, 256, 256) uint8
|
||
Returns: (B, T, d_model) — 1 image token per frame
|
||
"""
|
||
B, T = images.shape[:2]
|
||
imgs = images.reshape(B * T, 3, images.shape[-2], images.shape[-1]).float() / 255.0
|
||
with torch.no_grad():
|
||
patches = self.hamer_backbone(imgs)
|
||
patches = patches.reshape(B, T, 192, -1)
|
||
compressed = self.frame_compressor(patches)
|
||
return compressed
|
||
|
||
def _build_condition(self, batch: dict) -> tuple[torch.Tensor, torch.Tensor]:
|
||
"""Build the condition stream and return (condition, m)."""
|
||
if "image_tokens" in batch:
|
||
image_tokens = batch["image_tokens"]
|
||
else:
|
||
image_tokens = self._extract_image_features(batch["images"])
|
||
skeleton_tokens = self.skeleton_encoder(
|
||
batch["hamer_landmarks"],
|
||
batch["crop_intrinsics"],
|
||
)
|
||
condition, m = self.condition_builder(
|
||
image_tokens,
|
||
skeleton_tokens,
|
||
batch["hamer_confidence"],
|
||
)
|
||
return condition, m
|
||
|
||
def _build_z_positions(self, B: int, T: int, device: torch.device) -> torch.Tensor:
|
||
"""z-stream RoPE positions: [0, 1, 1, 2, 2, ..., T, T]"""
|
||
pos = torch.zeros(1 + 2 * T, dtype=torch.long, device=device)
|
||
for t in range(T):
|
||
pos[1 + t] = t + 1 # pose token
|
||
pos[1 + T + t] = t + 1 # trans token
|
||
return pos
|
||
|
||
def _build_cond_positions(self, B: int, T: int, device: torch.device) -> torch.Tensor:
|
||
"""c-stream RoPE positions: the 2 tokens per frame share position t"""
|
||
pos = torch.zeros(2 * T, dtype=torch.long, device=device)
|
||
for t in range(T):
|
||
pos[2 * t] = t
|
||
pos[2 * t + 1] = t
|
||
return pos
|
||
|
||
def forward(self, x_t_flat: torch.Tensor, t: torch.Tensor, batch: dict):
|
||
"""
|
||
Args:
|
||
x_t_flat: (B, flat_dim) — noisy flow target
|
||
t: (B,) — flow timestep
|
||
batch: dict with all condition data
|
||
Returns:
|
||
v_pred_flat: (B, flat_dim) — predicted flow velocity
|
||
"""
|
||
B, T = x_t_flat.shape[0], self.window_size
|
||
device = x_t_flat.device
|
||
|
||
# 1. Unpack flat → components
|
||
beta_t, poses_t, trans_t = self.mano_slice.unpack(x_t_flat)
|
||
|
||
# 2. Project to d_model → z-stream tokens
|
||
beta_tok = self.beta_proj_in(beta_t).unsqueeze(1) # (B, 1, D)
|
||
pose_tok = self.pose_proj_in(poses_t) # (B, T, D)
|
||
trans_tok = self.trans_proj_in(trans_t) # (B, T, D)
|
||
z_stream = torch.cat([beta_tok, pose_tok, trans_tok], dim=1) # (B, 1+2T, D)
|
||
|
||
# 3. Build condition → c-stream
|
||
condition, _ = self._build_condition(batch) # (B, 2*T, D)
|
||
|
||
# 4. RoPE positions
|
||
z_positions = self._build_z_positions(B, T, device)
|
||
cond_positions = self._build_cond_positions(B, T, device)
|
||
|
||
# 5. Transformer forward
|
||
pc_out, _ = self.transformer(
|
||
pc=z_stream,
|
||
dino=condition,
|
||
timesteps=t,
|
||
y=None,
|
||
txt_tokens=None,
|
||
z_positions=z_positions,
|
||
cond_positions=cond_positions,
|
||
rope_module=self.rope,
|
||
)
|
||
|
||
# 6. Project back to native dims
|
||
v_beta = self.beta_proj_out(pc_out[:, 0:1]) # (B, 1, 10)
|
||
v_pose = self.pose_proj_out(pc_out[:, 1:1 + T]) # (B, T, 48)
|
||
v_trans = self.trans_proj_out(pc_out[:, 1 + T:]) # (B, T, 3)
|
||
|
||
# 7. Pack flat
|
||
v_pred_flat = self.mano_slice.pack(
|
||
v_beta.squeeze(1), v_pose, v_trans
|
||
)
|
||
return v_pred_flat
|
||
|
||
@torch.no_grad()
|
||
def build_flow_target(self, batch: dict) -> torch.Tensor:
|
||
"""Build the GT flow target x_1 (normalized space) from the batch."""
|
||
beta = batch["mano_betas"] # (B, 10)
|
||
poses = batch["mano_params"] # (B, T, 48)
|
||
trans = batch["mano_trans"] # (B, T, 3)
|
||
flat = self.mano_slice.pack(beta, poses, trans)
|
||
return self.mano_slice.normalize(flat)
|
||
|
||
@torch.no_grad()
|
||
def reconstruct_params(self, x_1_flat: torch.Tensor):
|
||
"""
|
||
Reconstruct MANO parameters (original scale) from the predicted x_1 (normalized space).
|
||
|
||
All values are used directly, without cumsum.
|
||
|
||
Returns:
|
||
params: (B, T, 51) = cat[pose(48), trans(3)]
|
||
beta: (B, 10)
|
||
"""
|
||
x_1_orig = self.mano_slice.denormalize(x_1_flat)
|
||
beta, poses, trans = self.mano_slice.unpack(x_1_orig)
|
||
params = torch.cat([poses, trans], dim=-1) # (B, T, 51)
|
||
return params, beta
|
||
|
||
def get_model_wrapper(self):
|
||
"""Return a ModelWrapper for the ODE solver."""
|
||
return ModelWrapper(self)
|
||
|
||
def get_time_sampler(self):
|
||
"""Return the time sampler."""
|
||
return FluxTimeSampler()
|