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