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,124 @@
+"""Run the fixed 128-frame HMP on overlapping windows of a full sequence."""
+import copy
+import json
+from pathlib import Path
+import numpy as np
+import torch
+from scipy.spatial.transform import Rotation
+
+
+def window_starts(length, size=128, overlap=16):
+    if length < 1 or not 0 <= overlap < size:
+        raise ValueError('Invalid sequence length or overlap')
+    starts = [0]
+    while starts[-1] + size < length:
+        starts.append(starts[-1] + size - overlap)
+    return starts
+
+
+def slice_window(data, start, length, size=128):
+    result = {}
+    for key, value in data.items():
+        if isinstance(value, (torch.Tensor, np.ndarray)) and value.ndim >= 2 and value.shape[1] == length:
+            clip = value[:, start:min(start + size, length)].clone() if isinstance(value, torch.Tensor) else value[:, start:min(start + size, length)].copy()
+            valid = clip.shape[1]
+            if valid < size:
+                if isinstance(clip, torch.Tensor):
+                    pad = clip[:, -1:].repeat_interleave(size-valid, dim=1)
+                    clip = torch.cat([clip, pad], dim=1)
+                else:
+                    clip = np.concatenate([clip, np.repeat(clip[:, -1:], size-valid, axis=1)], axis=1)
+                if key == 'vis_mask':
+                    clip[:, valid:] = 0
+                elif key == 'joints2d':
+                    clip[:, valid:, :, -1] = 0
+            result[key] = clip
+        else:
+            result[key] = copy.deepcopy(value)
+    return result
+
+
+def blend_windows(original, clips, size=128, overlap=16):
+    result = {k: np.array(v, copy=True) for k, v in original.items()}
+    B, T = result['trans'].shape[:2]
+    denom = np.zeros(T, dtype=np.float64)
+    trans = np.zeros((B, T, 3), dtype=np.float64)
+    rotation_keys = ['root_orient', 'pose_body']
+    if all('decode_root' in clip for _, clip in clips):
+        result['decode_root'] = np.zeros((B, T, 3), dtype=np.float32)
+        rotation_keys.append('decode_root')
+    accum, refs = {}, {}
+    for key in rotation_keys:
+        shape = result[key].shape
+        refs[key] = Rotation.from_rotvec(result[key].reshape(-1, 3)).as_quat().reshape(*shape[:-1], 4)
+        accum[key] = np.zeros_like(refs[key])
+    for start, clip in clips:
+        n = min(size, T-start)
+        w = np.ones(n, dtype=np.float64)
+        if start > 0:
+            m = min(overlap, n)
+            w[:m] *= np.arange(1, m+1)/(m+1)
+        if start + size < T:
+            m = min(overlap, n)
+            w[-m:] *= np.arange(m, 0, -1)/(m+1)
+        denom[start:start+n] += w
+        trans[:, start:start+n] += clip['trans'][:, :n] * w[None, :, None]
+        for key in rotation_keys:
+            a = clip[key][:, :n].reshape(result[key][:, start:start+n].shape)
+            q = Rotation.from_rotvec(a.reshape(-1, 3)).as_quat().reshape(*a.shape[:-1], 4)
+            reference = refs[key][:, start:start+n]
+            q *= np.where(np.sum(q*reference, axis=-1, keepdims=True) < 0, -1., 1.)
+            weights = w.reshape((1,n) + (1,)*(q.ndim-2))
+            accum[key][:, start:start+n] += q * weights
+    if np.any(denom <= 0):
+        raise ValueError('Uncovered frames in HMP windows')
+    result['trans'] = (trans/denom[None,:,None]).astype(np.float32)
+    for key, q in accum.items():
+        if np.any(np.linalg.norm(q, axis=-1) < 1e-8):
+            raise ValueError('Degenerate quaternion blend')
+        result[key] = Rotation.from_quat(q.reshape(-1,4)).as_rotvec().reshape(result[key].shape).astype(np.float32)
+    # HMP currently fixes shape. Reject unhandled per-window shape changes.
+    for _, clip in clips:
+        if not np.allclose(clip['betas'], result['betas'], atol=1e-5):
+            raise ValueError('Windowed HMP requires fixed shape (opt_betas=false)')
+    result['latent_pose'] = result['pose_body'].copy()
+    return result
+
+
+def fit_windowed(obs_data, res_dicts, opt, data_args, out_dir, device, fit_fn):
+    from body_model import MANO
+    original = {k: v.detach().cpu().numpy() if isinstance(v, torch.Tensor) else np.asarray(v) for k,v in res_dicts[0].items()}
+    T = original['trans'].shape[1]
+    starts = window_starts(T)
+    out = Path(out_dir)
+    out.mkdir(parents=True, exist_ok=True)
+    mano_cfg = {k.lower():v for k,v in dict(opt.MANO).items()}
+    hand_model = MANO(batch_size=128, pose2rot=True, **mano_cfg).to(device)
+    clips = []
+    manifest = {'frames':T, 'window_size':128, 'overlap':16, 'windows':[]}
+    for index, start in enumerate(starts):
+        clip_dir = out / 'windows' / f'{start:06d}'
+        path = clip_dir / f'{Path(opt.HMP.vid_path).stem}_000000_world_results.npz'
+        print(f'HMP window {index+1}/{len(starts)}: frames {start}:{min(start+128,T)}', flush=True)
+        if path.is_file():
+            with np.load(path) as cached:
+                compatible = np.allclose(cached['betas'], original['betas'], atol=1e-5)
+            if not compatible:
+                path.rename(path.with_suffix('.variable_shape_backup.npz'))
+        if not path.is_file():
+            obs = slice_window(obs_data, start, T)
+            res = slice_window(res_dicts[0], start, T)
+            fit_fn(obs, (res,), hand_model, opt, data_args, str(clip_dir), device)
+        with np.load(path) as saved:
+            clip = {k:saved[k].copy() for k in saved.files}
+        if clip['trans'].shape[1] != 128 or not all(np.isfinite(v).all() for v in clip.values() if v.dtype.kind in 'fc'):
+            raise ValueError(f'Invalid HMP window output: {path}')
+        clips.append((start,clip))
+        manifest['windows'].append({'start':start,'valid_frames':min(128,T-start),'path':str(path)})
+        (out/'window_manifest.json').write_text(json.dumps(manifest,indent=2)+'\n')
+    result = blend_windows(original,clips)
+    target = out/f'{Path(opt.HMP.vid_path).stem}_000000_world_results.npz'
+    temporary = target.with_suffix('.partial.npz')
+    np.savez(temporary,**result)
+    temporary.replace(target)
+    print(f'Saved full-length windowed HMP: {target}', flush=True)