125 lines
5.9 KiB
Python
125 lines
5.9 KiB
Python
"""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)
|