Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
+727
@@ -0,0 +1,727 @@
|
||||
import os
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from body_model import MANO_JOINTS, KEYPT_VERTS, smpl_to_openpose
|
||||
from geometry.rotation import batch_rodrigues, rotation_matrix_to_angle_axis
|
||||
from util.logger import Logger
|
||||
from util.tensor import (
|
||||
detach_all,
|
||||
scatter_intervals,
|
||||
select_intervals,
|
||||
get_scatter_mask,
|
||||
)
|
||||
|
||||
from .base_scene import BaseSceneModel
|
||||
from .helpers import compute_world2prior, estimate_floor_planes
|
||||
|
||||
|
||||
J_BODY = len(MANO_JOINTS) - 1 # no root
|
||||
CONTACT_ORDERING = [
|
||||
"hips",
|
||||
"leftLeg",
|
||||
"rightLeg",
|
||||
"leftFoot",
|
||||
"rightFoot",
|
||||
"leftToeBase",
|
||||
"rightToeBase",
|
||||
"leftHand",
|
||||
"rightHand",
|
||||
]
|
||||
CONTACT_INDS = [MANO_JOINTS[jname] for jname in CONTACT_ORDERING]
|
||||
CONTACT_THRESH = 0.5
|
||||
|
||||
|
||||
class MovingSceneModel(BaseSceneModel):
|
||||
"""
|
||||
Scene model of moving people in a shared global reference frame
|
||||
|
||||
Parameters:
|
||||
batch_size: number of sequences to optimize
|
||||
seq_len: length of the sequences
|
||||
pose_prior: VPoser model
|
||||
motion_prior: humor model
|
||||
init_motion_prior: dict of GMM params to use for prior on initial motion state
|
||||
shared_floor: (default True) if true, sequences are in the same reference frame
|
||||
fit_gender: (optional) gender of SMPL model
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
batch_size,
|
||||
seq_len,
|
||||
body_model,
|
||||
pose_prior,
|
||||
motion_prior,
|
||||
init_motion_prior=None,
|
||||
fit_gender="neutral",
|
||||
use_init=False,
|
||||
opt_cams=False,
|
||||
opt_scale=True,
|
||||
cam_graph=False,
|
||||
est_floor=True,
|
||||
floor_type="shared",
|
||||
async_tracks=True,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(
|
||||
batch_size,
|
||||
seq_len,
|
||||
body_model,
|
||||
pose_prior,
|
||||
fit_gender=fit_gender,
|
||||
use_init=use_init,
|
||||
opt_cams=opt_cams,
|
||||
opt_scale=opt_scale,
|
||||
cam_graph=cam_graph,
|
||||
)
|
||||
assert motion_prior is not None
|
||||
assert motion_prior.model_data_config in [
|
||||
"mano+joints",
|
||||
"mano+joints+contacts",
|
||||
], "Only mano+joints motion prior configuration is supported!"
|
||||
|
||||
self.motion_prior = motion_prior
|
||||
self.init_motion_prior = init_motion_prior
|
||||
|
||||
# need latent dynamics sequence as well
|
||||
self.latent_motion_dim = self.motion_prior.latent_size
|
||||
self.cond_prior = self.motion_prior.use_conditional_prior
|
||||
assert self.cond_prior
|
||||
|
||||
# the frame chosen to use for the initial state (first frame by default)
|
||||
self.async_tracks = async_tracks
|
||||
self.register_buffer("track_start", torch.zeros(self.batch_size))
|
||||
self.register_buffer("track_end", torch.ones(self.batch_size) * self.seq_len)
|
||||
|
||||
self.shared_floor = floor_type == "shared"
|
||||
self.group_floor = floor_type == "group"
|
||||
self.floor_type = floor_type
|
||||
print("FLOOR TYPE", floor_type)
|
||||
self.est_floor = est_floor
|
||||
|
||||
@property
|
||||
def is_motion_active(self):
|
||||
return hasattr(self.params, "latent_motion")
|
||||
|
||||
def initialize(self, obs_data, cam_data, param_dict, data_fps):
|
||||
"""
|
||||
we need to also optimize for floor and world scale
|
||||
"""
|
||||
Logger.log("Initializing moving scene model with observed data")
|
||||
|
||||
self.params.set_cameras(
|
||||
cam_data,
|
||||
opt_scale=self.opt_scale,
|
||||
opt_cams=self.opt_cams,
|
||||
opt_focal=self.opt_cams,
|
||||
**param_dict,
|
||||
)
|
||||
# self.init_floor(obs_data, param_dict)
|
||||
self.init_first_state(obs_data, param_dict, data_fps)
|
||||
|
||||
# def init_floor(self, obs_data, param_dict):
|
||||
# if self.est_floor or self.group_floor:
|
||||
# with torch.no_grad():
|
||||
# smpl_preds = self.pred_mano(
|
||||
# param_dict["trans"],
|
||||
# param_dict["root_orient"],
|
||||
# self.latent2pose(param_dict["latent_pose"]),
|
||||
# param_dict["betas"],
|
||||
# )
|
||||
# floor_plane, floor_idcs = estimate_floor_planes(
|
||||
# smpl_preds["joints3d"].detach(),
|
||||
# obs_data["vis_mask"] > 0,
|
||||
# group=self.group_floor,
|
||||
# flatten=self.shared_floor,
|
||||
# )
|
||||
# if self.group_floor and not self.est_floor:
|
||||
# # don't use the estimated floor as initial plane
|
||||
# floor_plane = obs_data["floor_plane"][: len(floor_plane)]
|
||||
# else: # fixed shared or separate floors
|
||||
# num_floors = 1 if self.shared_floor else self.batch_size
|
||||
# floor_plane = obs_data["floor_plane"][:num_floors]
|
||||
# if self.shared_floor:
|
||||
# floor_idcs = torch.zeros(
|
||||
# self.batch_size, dtype=torch.long, device=floor_plane.device
|
||||
# )
|
||||
# else:
|
||||
# floor_idcs = torch.arange(
|
||||
# self.batch_size, dtype=torch.long, device=floor_plane.device
|
||||
# )
|
||||
|
||||
# Logger.log(f"ESTIMATED FLOORS: {str(floor_plane.detach().cpu())}")
|
||||
# Logger.log(f"FLOOR IDCS: {str(floor_idcs.detach().cpu())}")
|
||||
# self.params.set_param("floor_plane", floor_plane.float().detach())
|
||||
# self.params.set_param(
|
||||
# "floor_idcs", floor_idcs.long().detach(), requires_grad=False
|
||||
# )
|
||||
|
||||
def init_first_state(self, obs_data, param_dict, data_fps):
|
||||
"""
|
||||
initialize the latent motion and first state of each track
|
||||
using per-frame trajectories of trans, rot, latent_pose, betas
|
||||
and observed data
|
||||
"""
|
||||
B, T = self.batch_size, self.seq_len
|
||||
param_dict = detach_all(param_dict)
|
||||
|
||||
# select the valid segments of each track; pad the shorter tracks with final element
|
||||
if self.async_tracks and "track_interval" in obs_data:
|
||||
print("asynchronous tracks")
|
||||
interval = obs_data["track_interval"] # (B, 2)
|
||||
start, end = interval[:, 0], interval[:, 1]
|
||||
self.track_start, self.track_end = start, end
|
||||
param_dict = select_dict_segments(
|
||||
param_dict, start, end, names=["trans", "root_orient", "latent_pose"]
|
||||
)
|
||||
|
||||
trans = param_dict["trans"] # (B, T, 3)
|
||||
root_orient = param_dict["root_orient"] # (B, T, 3)
|
||||
latent_pose = param_dict["latent_pose"] # (B, T, D)
|
||||
betas = param_dict["betas"] # (B, b)
|
||||
|
||||
Logger.log(f"SEL TRACKS {trans.shape}, {root_orient.shape}")
|
||||
|
||||
# save each track's first appearance
|
||||
self.params.set_param("trans", trans[:, :1])
|
||||
Logger.log(f"INITIAL TRANS {trans[:, :1].detach().cpu()}")
|
||||
self.params.set_param("root_orient", root_orient[:, :1])
|
||||
self.params.set_param("latent_pose", latent_pose[:, :1])
|
||||
self.params.set_param("betas", betas)
|
||||
|
||||
# pass the current pose estimates through the motion prior
|
||||
self.data_fps = data_fps
|
||||
body_pose = self.latent2pose(latent_pose)
|
||||
init_latent = self.infer_latent_motion(
|
||||
trans, root_orient, body_pose, betas, data_fps
|
||||
).detach() # (B, T-1, D)
|
||||
self.params.set_param("latent_motion", init_latent)
|
||||
|
||||
# estimate velocities in prior frame and save initial state
|
||||
trans_vel, joints_vel, root_orient_vel = self.estimate_prior_velocities(
|
||||
trans, root_orient, body_pose, betas, data_fps
|
||||
)
|
||||
self.params.set_param("trans_vel", trans_vel[:, :1].detach())
|
||||
self.params.set_param("joints_vel", joints_vel[:, :1].detach())
|
||||
self.params.set_param("root_orient_vel", root_orient_vel[:, :1].detach())
|
||||
|
||||
def get_optim_result(self, num_steps=-1):
|
||||
res = super().get_optim_result()
|
||||
if not self.is_motion_active:
|
||||
return res
|
||||
|
||||
num_steps = self.seq_len if num_steps < 0 else num_steps
|
||||
prior_res, world_res = self.rollout_latent_motion(
|
||||
self.params.latent_motion[:, : num_steps - 1]
|
||||
)
|
||||
world_res = self.synchronize_preds(world_res, num_steps)
|
||||
res["world"].update(world_res)
|
||||
|
||||
prior_res = self.synchronize_preds(prior_res, num_steps)
|
||||
res["prior"] = prior_res
|
||||
return res
|
||||
|
||||
def estimate_prior_velocities(self, trans, root_orient, body_pose, betas, data_fps):
|
||||
"""
|
||||
compute velocities of trajectory in prior frame (not world frame)
|
||||
Transforms the first of each track into the prior frame
|
||||
:param trans (B, T, 3)
|
||||
:param root_orient (B, T, 3)
|
||||
:param body_pose (B, T, Dp)
|
||||
:param betas (B, Db)
|
||||
"""
|
||||
with torch.no_grad():
|
||||
self.update_world2prior(trans, root_orient, body_pose, betas)
|
||||
trans, root_orient = self.apply_world2prior(
|
||||
trans, root_orient, body_pose, betas
|
||||
)
|
||||
mano_results = self.pred_mano(trans, root_orient, body_pose, betas)
|
||||
return estimate_velocities(
|
||||
trans, root_orient, mano_results["joints3d"], data_fps
|
||||
)
|
||||
|
||||
def update_world2prior(self, trans, root_orient, body_pose, betas):
|
||||
B = trans.shape[0]
|
||||
mano_data = self.pred_mano(
|
||||
trans[:, :1], root_orient[:, :1], body_pose[:, :1], betas
|
||||
)
|
||||
floor_plane = self.params.floor_plane[self.params.floor_idcs]
|
||||
R, t, height = compute_world2prior(
|
||||
floor_plane,
|
||||
trans[:, 0],
|
||||
root_orient[:, 0],
|
||||
mano_data["joints3d"][:, 0, 0],
|
||||
)
|
||||
self.world2prior_R = R # (B, 3, 3)
|
||||
self.world2prior_t = t # (B, 3)
|
||||
self.world2prior_root_height = height # (B)
|
||||
|
||||
def apply_world2prior(self, trans, root_orient, body_pose, betas, inverse=False):
|
||||
"""
|
||||
Applies the world2prior tranformation to trans, root_orient
|
||||
If T=1, this function assumes they are at key_frame_idx,
|
||||
which we need to compute the offset from the origin
|
||||
:param trans (B, T, 3)
|
||||
:param root_orient (B, T, 3)
|
||||
:param body_pose (B, T, J, 3)
|
||||
:param betas (B, b)
|
||||
:param inverse (bool) optional, default False
|
||||
"""
|
||||
B, T, _ = trans.size()
|
||||
# (B, 3, 3), (B, 3), (B)
|
||||
R, t, root_height = (
|
||||
self.world2prior_R,
|
||||
self.world2prior_t,
|
||||
self.world2prior_root_height,
|
||||
)
|
||||
R_time = R.unsqueeze(1).expand((B, T, 3, 3))
|
||||
t_time = t.unsqueeze(1).expand((B, T, 3))
|
||||
root_orient_mat = batch_rodrigues(root_orient.reshape(-1, 3)).reshape(
|
||||
B, T, 3, 3
|
||||
)
|
||||
|
||||
if inverse:
|
||||
R_time = R_time.transpose(-1, -2)
|
||||
|
||||
root_orient_mat = torch.matmul(R_time, root_orient_mat)
|
||||
root_orient = rotation_matrix_to_angle_axis(
|
||||
root_orient_mat.reshape(B * T, 3, 3)
|
||||
).reshape(B, T, 3)
|
||||
|
||||
if inverse:
|
||||
# transform so first frame is at origin
|
||||
trans = trans - trans[:, :1, :]
|
||||
|
||||
# rotates to world frame
|
||||
trans = torch.matmul(R_time, trans[..., None])[..., 0]
|
||||
# translate to world frame
|
||||
trans = trans - t_time
|
||||
|
||||
return trans, root_orient
|
||||
|
||||
# first transform so the trans of key frame is at origin
|
||||
trans = trans + t_time
|
||||
# then rotate to canonical frame
|
||||
trans = torch.matmul(R_time, trans[..., None])[..., 0]
|
||||
# compute the root height after transforming
|
||||
cur_smpl_data = self.pred_mano(trans, root_orient, body_pose, betas)
|
||||
cur_root_height = cur_smpl_data["joints3d"][:, 0, 0, 2:3]
|
||||
# then apply floor offset so the root joint is at root_height
|
||||
height_diff = root_height - cur_root_height
|
||||
trans_offset = torch.cat(
|
||||
[torch.zeros((B, 2)).to(height_diff), height_diff], axis=1
|
||||
)
|
||||
trans = trans + trans_offset.reshape((B, 1, 3))
|
||||
|
||||
return trans, root_orient
|
||||
|
||||
def convert_prior_rot_inputs(self, root_orient, body_pose):
|
||||
# convert rots
|
||||
# body pose and root orient are both in aa
|
||||
B, T = root_orient.shape[:2]
|
||||
if (
|
||||
self.motion_prior.in_rot_rep == "mat"
|
||||
or self.motion_prior.in_rot_rep == "6d"
|
||||
):
|
||||
root_orient_in = batch_rodrigues(root_orient.reshape(-1, 3)).reshape(
|
||||
(B, T, 9)
|
||||
)
|
||||
body_pose_in = batch_rodrigues(body_pose.reshape(-1, 3)).reshape(
|
||||
(B, T, J_BODY * 9)
|
||||
)
|
||||
if self.motion_prior.in_rot_rep == "6d":
|
||||
root_orient_in = root_orient[:, :, :6]
|
||||
body_pose_in = body_pose.reshape((B, T, J_BODY, 9))[:, :, :, :6].reshape(
|
||||
(B, T, J_BODY * 6)
|
||||
)
|
||||
return root_orient_in, body_pose_in
|
||||
|
||||
def convert_prior_rot_outputs(self, out_dict):
|
||||
keys = ["root_orient", "pose_body"]
|
||||
for k in keys:
|
||||
out = out_dict[k]
|
||||
B, T = out.shape[:2]
|
||||
out_dict[k] = rotation_matrix_to_angle_axis(
|
||||
out.reshape((-1, 3, 3))
|
||||
).reshape((B, T, -1))
|
||||
return out_dict
|
||||
|
||||
def infer_latent_motion(
|
||||
self, trans, root_orient, body_pose, betas, data_fps, full_forward_pass=False
|
||||
):
|
||||
"""
|
||||
By default, gets a sequence of z's from the current SMPL optim params.
|
||||
|
||||
If full_forward_pass is true, also samples from the posterior and feeds
|
||||
through the motion prior decoder to get all terms needed to calculate the ELBO.
|
||||
"""
|
||||
B, T, _ = trans.size()
|
||||
h = 1.0 / data_fps
|
||||
|
||||
# need to first transform into canonical coordinate frame
|
||||
self.update_world2prior(trans, root_orient, body_pose, betas)
|
||||
trans, root_orient = self.apply_world2prior(
|
||||
trans, root_orient, body_pose, betas
|
||||
)
|
||||
|
||||
mano_results = self.pred_mano(trans, root_orient, body_pose, betas)
|
||||
joints = mano_results["joints3d"] # (B, T, len(MANO_JOINTS), 3)
|
||||
trans_vel, joints_vel, root_orient_vel = estimate_velocities(
|
||||
trans, root_orient, joints, data_fps
|
||||
)
|
||||
|
||||
root_orient_in, body_pose_in = self.convert_prior_rot_inputs(
|
||||
root_orient, body_pose
|
||||
)
|
||||
joints_in = joints.reshape((B, T, -1))
|
||||
joints_vel_in = joints_vel.reshape((B, T, -1))
|
||||
|
||||
seq_dict = {
|
||||
"trans": trans,
|
||||
"trans_vel": trans_vel,
|
||||
"root_orient": root_orient_in,
|
||||
"root_orient_vel": root_orient_vel,
|
||||
"pose_body": body_pose_in,
|
||||
"joints": joints_in,
|
||||
"joints_vel": joints_vel_in,
|
||||
}
|
||||
|
||||
infer_results = self.motion_prior.infer_global_seq(
|
||||
seq_dict, full_forward_pass=full_forward_pass
|
||||
)
|
||||
if full_forward_pass:
|
||||
# return both the given motion and the one from the forward pass
|
||||
# make sure rotations are matrix
|
||||
# NOTE: assumes seq_dict is same thing we want to compute loss on
|
||||
# Need to change if multiple future steps.
|
||||
if self.motion_prior.in_rot_rep != "mat":
|
||||
seq_dict["trans"] = batch_rodrigues(root_orient.reshape(-1, 3)).reshape(
|
||||
(B, T, 9)
|
||||
)
|
||||
seq_dict["pose_body"] = batch_rodrigues(
|
||||
body_pose.reshape(-1, 3)
|
||||
).reshape((B, T, J_BODY * 9))
|
||||
# do not need initial step anymore since output will be T-1
|
||||
for k, v in seq_dict.items():
|
||||
seq_dict[k] = v[:, 1:]
|
||||
for k in infer_results.keys():
|
||||
if k != "posterior_distrib" and k != "prior_distrib":
|
||||
infer_results[k] = infer_results[k][
|
||||
:, :, 0
|
||||
] # only want first output step
|
||||
infer_results = (seq_dict, infer_results)
|
||||
else:
|
||||
prior_z, posterior_z = infer_results
|
||||
infer_results = posterior_z[0] # mean of the approximate posterior
|
||||
|
||||
return infer_results
|
||||
|
||||
def rollout_mano_steps(self, num_steps=-1):
|
||||
# rollout motion given initial state
|
||||
num_steps = self.seq_len if num_steps < 0 else num_steps
|
||||
latent_motion = self.params.latent_motion[:, : num_steps - 1]
|
||||
|
||||
# roll out tracks from their first appearances
|
||||
# NOTE rolled out tracks are unsynced in time
|
||||
res, world_res = self.rollout_latent_motion(
|
||||
latent_motion, return_prior=self.cond_prior
|
||||
)
|
||||
|
||||
# get the smpl predictions for the unsynced tracks
|
||||
preds = self.pred_mano(
|
||||
res["trans"], res["root_orient"], res["pose_body"], res["betas"]
|
||||
)
|
||||
world_preds = self.pred_mano(
|
||||
world_res["trans"],
|
||||
world_res["root_orient"],
|
||||
world_res["pose_body"],
|
||||
world_res["betas"],
|
||||
)
|
||||
|
||||
# pass along relevant results
|
||||
preds["joints3d_rollout"] = res["joints"]
|
||||
preds["latent_pose"] = self.pose2latent(res["pose_body"])
|
||||
preds["joints3d_init"] = preds["joints3d"][:, :1]
|
||||
|
||||
if "contacts" in res:
|
||||
preds["contacts"] = res["contacts"]
|
||||
if "contacts_conf" in res:
|
||||
preds["contacts_conf"] = res["contacts_conf"]
|
||||
|
||||
# synchronize the tracklets to the same timesteps
|
||||
preds = self.synchronize_preds(preds, num_steps)
|
||||
world_preds = self.synchronize_preds(world_preds, num_steps)
|
||||
|
||||
# pass along unsynchronized motion latents
|
||||
preds["latent_motion"] = latent_motion
|
||||
if "cond_prior" in res:
|
||||
preds["cond_prior"] = res["cond_prior"]
|
||||
|
||||
return preds, world_preds
|
||||
|
||||
def rollout_latent_motion(
|
||||
self,
|
||||
latent_motion,
|
||||
return_prior=False,
|
||||
return_vel=False,
|
||||
num_steps=-1,
|
||||
canonicalize_input=False,
|
||||
):
|
||||
"""
|
||||
From the stored initial state, rolls out a sequence from the latent_motion vector
|
||||
NOTE: the stored initial states are not at the same time step, but we return all
|
||||
predicted sequences starting from time 0.
|
||||
|
||||
If latent_motion is None, samples num_steps into the future sequence from the prior.
|
||||
using the mean of the prior rather than random samples.
|
||||
|
||||
If canonicalize_input is True, transform the initial state into the local canonical
|
||||
frame before roll out
|
||||
"""
|
||||
# get the first frame state
|
||||
# NOTE: first states do not occur at same time step
|
||||
trans, root_orient, betas = (
|
||||
self.params.trans,
|
||||
self.params.root_orient,
|
||||
self.params.betas,
|
||||
)
|
||||
trans_vel, root_orient_vel, joints_vel = (
|
||||
self.params.trans_vel,
|
||||
self.params.root_orient_vel,
|
||||
self.params.joints_vel,
|
||||
)
|
||||
body_pose = self.latent2pose(self.params.latent_pose)
|
||||
|
||||
B = trans.size(0)
|
||||
is_sampling = latent_motion is None
|
||||
Tm1 = num_steps if latent_motion is None else latent_motion.size(1)
|
||||
if is_sampling and Tm1 <= 0:
|
||||
Logger.log("num_steps must be positive to sample!")
|
||||
exit()
|
||||
|
||||
# need to first transform initial state into canonical coordinate frame
|
||||
self.update_world2prior(trans, root_orient, body_pose, betas)
|
||||
trans, root_orient = self.apply_world2prior(
|
||||
trans, root_orient, body_pose, betas
|
||||
)
|
||||
|
||||
smpl_results = self.pred_mano(trans, root_orient, body_pose, betas)
|
||||
joints = smpl_results["joints3d"] # (B, T, len(MANO_JOINTS), 3)
|
||||
|
||||
# update to correct rotations for input
|
||||
root_orient_in, body_pose_in = self.convert_prior_rot_inputs(
|
||||
root_orient, body_pose
|
||||
)
|
||||
joints_in = joints.reshape((B, 1, -1))
|
||||
joints_vel_in = joints_vel.reshape((B, 1, -1))
|
||||
|
||||
rollout_in_dict = {
|
||||
"trans": trans,
|
||||
"root_orient": root_orient_in,
|
||||
"pose_body": body_pose_in,
|
||||
"joints": joints_in,
|
||||
"trans_vel": trans_vel,
|
||||
"root_orient_vel": root_orient_vel,
|
||||
"joints_vel": joints_vel_in,
|
||||
}
|
||||
|
||||
roll_output = self.motion_prior.roll_out(
|
||||
None,
|
||||
rollout_in_dict,
|
||||
Tm1,
|
||||
z_seq=latent_motion,
|
||||
return_prior=return_prior,
|
||||
return_z=is_sampling,
|
||||
canonicalize_input=canonicalize_input,
|
||||
gender=[self.fit_gender] * B,
|
||||
betas=betas.reshape((B, 1, -1)),
|
||||
)
|
||||
|
||||
pred_dict = prior_out = None
|
||||
if return_prior:
|
||||
pred_dict, prior_out = roll_output
|
||||
else:
|
||||
pred_dict = roll_output
|
||||
|
||||
pred_dict = self.convert_prior_rot_outputs(pred_dict)
|
||||
|
||||
# concat with initial state
|
||||
trans_out = torch.cat([trans, pred_dict["trans"]], dim=1)
|
||||
root_orient_out = torch.cat([root_orient, pred_dict["root_orient"]], dim=1)
|
||||
body_pose_out = torch.cat([body_pose, pred_dict["pose_body"]], dim=1)
|
||||
joints_out = torch.cat(
|
||||
[joints, pred_dict["joints"].reshape((B, Tm1, -1, 3))], dim=1
|
||||
)
|
||||
out_dict = {
|
||||
"trans": trans_out,
|
||||
"root_orient": root_orient_out,
|
||||
"pose_body": body_pose_out,
|
||||
"joints": joints_out,
|
||||
"betas": betas,
|
||||
}
|
||||
if return_vel:
|
||||
out_dict["trans_vel"] = torch.cat(
|
||||
[trans_vel, pred_dict["trans_vel"]], dim=1
|
||||
)
|
||||
out_dict["root_orient_vel"] = torch.cat(
|
||||
[root_orient_vel, pred_dict["root_orient_vel"]], dim=1
|
||||
)
|
||||
out_dict["joints_vel"] = torch.cat(
|
||||
[joints_vel, pred_dict["joints_vel"].reshape((B, Tm1, -1, 3))],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
if return_prior: # return the mean and var of distribution stacked
|
||||
pm, pv = prior_out
|
||||
out_dict["cond_prior"] = torch.stack([pm, pv], dim=-1)
|
||||
|
||||
if self.motion_prior.model_data_config == "smpl+joints+contacts":
|
||||
pred_contacts = pred_dict["contacts"]
|
||||
# get binary classification
|
||||
contact_conf = torch.sigmoid(pred_contacts)
|
||||
pred_contacts = (contact_conf > CONTACT_THRESH).to(torch.float)
|
||||
# expand to full body
|
||||
full_contact_conf = torch.zeros((B, Tm1, len(MANO_JOINTS))).to(contact_conf)
|
||||
full_contact_conf[:, :, CONTACT_INDS] = (
|
||||
full_contact_conf[:, :, CONTACT_INDS] + contact_conf
|
||||
)
|
||||
full_contacts = torch.zeros((B, Tm1, len(MANO_JOINTS))).to(pred_contacts)
|
||||
full_contacts[:, :, CONTACT_INDS] = (
|
||||
full_contacts[:, :, CONTACT_INDS] + pred_contacts
|
||||
)
|
||||
# repeat first entry for t0
|
||||
full_contact_conf = torch.cat(
|
||||
[full_contact_conf[:, 0:1], full_contact_conf], dim=1
|
||||
)
|
||||
full_contacts = torch.cat([full_contacts[:, 0:1], full_contacts], dim=1)
|
||||
out_dict["contacts_conf"] = full_contact_conf
|
||||
out_dict["contacts"] = full_contacts
|
||||
|
||||
if is_sampling:
|
||||
out_dict["z"] = pred_dict["z"]
|
||||
|
||||
cam_dict = {
|
||||
"trans": out_dict["trans"],
|
||||
"root_orient": out_dict["root_orient"],
|
||||
"pose_body": out_dict["pose_body"],
|
||||
"betas": out_dict["betas"],
|
||||
}
|
||||
|
||||
# also must return trans and root orient in camera frame
|
||||
cam_dict["trans"], cam_dict["root_orient"] = self.apply_world2prior(
|
||||
out_dict["trans"],
|
||||
out_dict["root_orient"],
|
||||
out_dict["pose_body"],
|
||||
out_dict["betas"],
|
||||
inverse=True,
|
||||
)
|
||||
return out_dict, cam_dict
|
||||
|
||||
def synchronize_preds(self, pred_dict, seg_len):
|
||||
"""
|
||||
synchronize predictions in time, scatter predictions into [0, seq_len)
|
||||
"""
|
||||
if not self.async_tracks:
|
||||
# predictions are already synchronized
|
||||
return pred_dict, None
|
||||
|
||||
# return time-synchronized predictions
|
||||
start = self.track_start
|
||||
end = torch.clip(start + seg_len, max=self.track_end)
|
||||
T = end.max()
|
||||
pred_dict = scatter_dict_segments(pred_dict, start, end, T)
|
||||
pred_mask = get_scatter_mask(start, end, T)
|
||||
pred_dict["track_mask"] = pred_mask
|
||||
return pred_dict
|
||||
|
||||
|
||||
def scatter_dict_segments(data_dict, start, end, T=None, names=None):
|
||||
"""
|
||||
the tracks as they are stored are synchronized by time step
|
||||
Uses the start and end to synchronize output predictions
|
||||
"""
|
||||
sync_dict = data_dict.copy()
|
||||
min_len = (end - start).min()
|
||||
if names is None:
|
||||
names = data_dict.keys()
|
||||
for name in names:
|
||||
val = data_dict[name]
|
||||
if not isinstance(val, torch.Tensor) or val.ndim < 3 or val.shape[1] < min_len:
|
||||
continue
|
||||
sync_dict[name] = scatter_intervals(val, start, end, T)
|
||||
return sync_dict
|
||||
|
||||
|
||||
def select_dict_segments(data_dict, start, end, names=None):
|
||||
out_data_dict = data_dict.copy()
|
||||
min_len = (end - start).min()
|
||||
if names is None:
|
||||
names = data_dict.keys()
|
||||
for name in names:
|
||||
val = data_dict[name]
|
||||
if not isinstance(val, torch.Tensor) or val.ndim < 3 or val.shape[1] < min_len:
|
||||
continue
|
||||
# only select for time-series observations
|
||||
out_data_dict[name] = select_intervals(val, start, end)
|
||||
return out_data_dict
|
||||
|
||||
|
||||
def estimate_velocities(trans, root_orient, joints3d, data_fps):
|
||||
"""
|
||||
Estimates velocity inputs to the motion prior.
|
||||
- trans (B, T, 3) root translation
|
||||
- root_orient (B, T, 3) aa root orientation
|
||||
- joints3d (B, T, len(MANO_JOINTS), 3) joints3d of SMPL prediction
|
||||
"""
|
||||
B, T, _ = trans.size()
|
||||
h = 1.0 / data_fps
|
||||
trans_vel = estimate_linear_velocity(trans, h)
|
||||
joints_vel = estimate_linear_velocity(joints3d, h)
|
||||
root_orient_mat = batch_rodrigues(root_orient.reshape((-1, 3))).reshape(
|
||||
(B, T, 3, 3)
|
||||
)
|
||||
root_orient_vel = estimate_angular_velocity(root_orient_mat, h)
|
||||
return trans_vel, joints_vel, root_orient_vel
|
||||
|
||||
|
||||
def estimate_linear_velocity(data_seq, h):
|
||||
"""
|
||||
Given some batched data sequences of T timesteps in the shape (B, T, ...), estimates
|
||||
the velocity for the middle T-2 steps using a second order central difference scheme.
|
||||
The first and last frames are with forward and backward first-order
|
||||
differences, respectively
|
||||
- h : step size
|
||||
"""
|
||||
# first steps is forward diff (t+1 - t) / h
|
||||
init_vel = (data_seq[:, 1:2] - data_seq[:, :1]) / h
|
||||
# middle steps are second order (t+1 - t-1) / 2h
|
||||
middle_vel = (data_seq[:, 2:] - data_seq[:, 0:-2]) / (2 * h)
|
||||
# last step is backward diff (t - t-1) / h
|
||||
final_vel = (data_seq[:, -1:] - data_seq[:, -2:-1]) / h
|
||||
|
||||
vel_seq = torch.cat([init_vel, middle_vel, final_vel], dim=1)
|
||||
return vel_seq
|
||||
|
||||
|
||||
def estimate_angular_velocity(rot_seq, h):
|
||||
"""
|
||||
Given a batch of sequences of T rotation matrices, estimates angular velocity at T-2 steps.
|
||||
Input sequence should be of shape (B, T, ..., 3, 3)
|
||||
"""
|
||||
# see https://en.wikipedia.org/wiki/Angular_velocity#Calculation_from_the_orientation_matrix
|
||||
dRdt = estimate_linear_velocity(rot_seq, h)
|
||||
R = rot_seq
|
||||
RT = R.transpose(-1, -2)
|
||||
# compute skew-symmetric angular velocity tensor
|
||||
w_mat = torch.matmul(dRdt, RT)
|
||||
# pull out angular velocity vector
|
||||
# average symmetric entries
|
||||
w_x = (-w_mat[..., 1, 2] + w_mat[..., 2, 1]) / 2.0
|
||||
w_y = (w_mat[..., 0, 2] - w_mat[..., 2, 0]) / 2.0
|
||||
w_z = (-w_mat[..., 0, 1] + w_mat[..., 1, 0]) / 2.0
|
||||
w = torch.stack([w_x, w_y, w_z], axis=-1)
|
||||
return w
|
||||
Reference in New Issue
Block a user