Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
+174
@@ -0,0 +1,174 @@
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
# from body_model import MANO_JOINTS
|
||||
from geometry.camera import invert_camera, compose_cameras
|
||||
from geometry.rotation import batch_rodrigues
|
||||
from util.logger import Logger
|
||||
from util.tensor import detach_all
|
||||
|
||||
# J_BODY = len(MANO_JOINTS) - 1 # no root
|
||||
# assert J_BODY == 15
|
||||
|
||||
class Params(nn.Module):
|
||||
def __init__(self, batch_size):
|
||||
super().__init__()
|
||||
self.batch_size = batch_size
|
||||
self.param_names = set()
|
||||
|
||||
def set_param(self, name, val, requires_grad=False):
|
||||
print("SETTING PARAM", name, val.shape)
|
||||
with torch.no_grad():
|
||||
setattr(
|
||||
self, name, nn.Parameter(val.contiguous(), requires_grad=requires_grad)
|
||||
)
|
||||
self.param_names.add(name)
|
||||
|
||||
def get_param(self, name):
|
||||
if name not in self.param_names:
|
||||
raise ValueError(f"{name} not stored as opt param")
|
||||
return getattr(self, name)
|
||||
|
||||
def load_dict(self, param_dict):
|
||||
for name, val in param_dict.items():
|
||||
self.set_param(name, val, requires_grad=False)
|
||||
|
||||
def get_dict(self):
|
||||
return {name: self.get_param_item(name) for name in self.param_names}
|
||||
|
||||
def get_vars(self, names=None):
|
||||
if names is None:
|
||||
names = self.param_names
|
||||
return {name: self.get_param(name) for name in names}
|
||||
|
||||
def get_param_item(self, name):
|
||||
with torch.no_grad():
|
||||
param = self.get_param(name)
|
||||
return param.detach()
|
||||
|
||||
def _set_param_grad(self, name, val: bool):
|
||||
if name not in self.param_names:
|
||||
raise ValueError(f"{name} not stored as param")
|
||||
param = getattr(self, name)
|
||||
assert isinstance(param, torch.Tensor)
|
||||
param.requires_grad = val
|
||||
|
||||
def set_require_grads(self, names):
|
||||
"""
|
||||
set parameters in names to True, set all others to False
|
||||
"""
|
||||
for name in self.param_names:
|
||||
self._set_param_grad(name, False)
|
||||
|
||||
for name in names:
|
||||
self._set_param_grad(name, True)
|
||||
|
||||
Logger.log("Set parameter grads:")
|
||||
Logger.log(
|
||||
{name: getattr(self, name).requires_grad for name in self.param_names}
|
||||
)
|
||||
|
||||
|
||||
class CameraParams(Params):
|
||||
"""
|
||||
Parameter container with cameras
|
||||
"""
|
||||
|
||||
def set_cameras(
|
||||
self, cam_data, opt_scale=True, opt_cams=False, opt_focal=False, **kwargs
|
||||
):
|
||||
self.opt_scale = opt_scale
|
||||
self.opt_cams = opt_cams
|
||||
self.opt_focal = opt_focal
|
||||
|
||||
# (T, 3, 3), (T, 3)
|
||||
cam_R, cam_t = cam_data["cam_R"], cam_data["cam_t"]
|
||||
intrins = cam_data["intrins"] # (T, 4)
|
||||
|
||||
T = cam_R.shape[0]
|
||||
device = cam_R.device
|
||||
|
||||
# assume focal length and center are same for all timesteps
|
||||
# print(intrins.shape, cam_R.shape, cam_t.shape)
|
||||
# torch.Size([128, 4]) torch.Size([128, 3, 3]) torch.Size([128, 3])
|
||||
|
||||
self.cam_center = intrins[:, 2:] # (T, 2)
|
||||
cam_f = intrins[:, :2] # (T, 2)
|
||||
if self.opt_focal:
|
||||
self.set_param("cam_f", cam_f)
|
||||
else:
|
||||
self.cam_f = cam_f
|
||||
|
||||
self._cam_R = cam_R # (T, 3, 3)
|
||||
self._cam_t = cam_t # (T, 3)
|
||||
|
||||
world_scale = torch.ones(1, 1, device=device)
|
||||
if self.opt_scale:
|
||||
if "world_scale" in kwargs:
|
||||
world_scale = kwargs["world_scale"]
|
||||
self.set_param("world_scale", world_scale)
|
||||
else:
|
||||
self.world_scale = world_scale
|
||||
|
||||
if self.opt_cams:
|
||||
init_delta_R = torch.zeros(T, 3, device=device)
|
||||
init_delta_t = torch.zeros(T, 3, device=device)
|
||||
if "delta_cam_R" in kwargs:
|
||||
print("setting delta_cam_R from kwargs")
|
||||
init_delta_R = kwargs["delta_cam_R"]
|
||||
if "delta_cam_t" in kwargs:
|
||||
print("setting delta_cam_t from kwargs")
|
||||
init_delta_t = kwargs["delta_cam_t"]
|
||||
|
||||
self.set_param("delta_cam_R", init_delta_R)
|
||||
self.set_param("delta_cam_t", init_delta_t)
|
||||
|
||||
@property
|
||||
def intrins(self): # (4,)
|
||||
return torch.cat([self.cam_f[0], self.cam_center[0]], dim=-1).detach().cpu()
|
||||
|
||||
def get_extrinsics(self):
|
||||
"""
|
||||
returns (T, 3, 3), (T, 3)
|
||||
"""
|
||||
cam_R, cam_t = self._cam_R, self._cam_t
|
||||
print('get_extrinsics')
|
||||
if self.opt_cams:
|
||||
raise ValueError
|
||||
print('self.opt_cams:', self.opt_cams)
|
||||
# raise ValueError
|
||||
dR = batch_rodrigues(self.delta_cam_R)
|
||||
cam_R = torch.matmul(cam_R, dR)
|
||||
cam_t = cam_t + self.delta_cam_t
|
||||
return cam_R, cam_t * self.world_scale
|
||||
|
||||
def get_intrinsics(self):
|
||||
"""
|
||||
returns (T, 4)
|
||||
"""
|
||||
return torch.cat([self.cam_f, self.cam_center], axis=-1)
|
||||
|
||||
def get_cameras(self, idcs=None):
|
||||
"""
|
||||
returns cam_R (B, T, 3, 3) cam_t (B, T, 3), cam_f (T, 2), cam_center (T, 2)
|
||||
"""
|
||||
if idcs is None:
|
||||
idcs = np.arange(self._cam_R.shape[0])
|
||||
|
||||
# print('idcs:', idcs, self._cam_R.shape[0], self._cam_R.shape)
|
||||
cam_R, cam_t = self.get_extrinsics()
|
||||
cam_R, cam_t = cam_R[None, idcs], cam_t[None, idcs]
|
||||
cam_f, cam_center = self.cam_f[idcs], self.cam_center[idcs]
|
||||
|
||||
B = self.batch_size
|
||||
cam_R = cam_R.repeat(B, 1, 1, 1)
|
||||
cam_t = cam_t.repeat(B, 1, 1)
|
||||
# print(cam_R.shape, cam_t.shape, cam_f.shape, cam_center.shape)
|
||||
# print(B, self.batch_size)
|
||||
# raise ValueError
|
||||
# torch.Size([2, 128, 3, 3]) torch.Size([2, 128, 3]) torch.Size([128, 2]) torch.Size([128, 2])
|
||||
# 2 2
|
||||
return cam_R, cam_t, cam_f, cam_center
|
||||
Reference in New Issue
Block a user