lib/models/layers/utils.py

import torch
from lib.utils import transforms



def rollout_global_motion(root_r, root_v, init_trans=None):
    b, f = root_v.shape[:2]
    root = transforms.rotation_6d_to_matrix(root_r[:])
    vel_world = (root[:, :-1] @ root_v.unsqueeze(-1)).squeeze(-1)
    trans = torch.cumsum(vel_world, dim=1)
    
    if init_trans is not None: trans = trans + init_trans
    return root[:, 1:], trans

def compute_camera_motion(output, root_c_d6d, root_w, trans, pred_cam):
    root_c = transforms.rotation_6d_to_matrix(root_c_d6d)  # Root orient in cam coord
    cam_R = root_c @ root_w.mT
    pelvis_cam = output.full_cam.view_as(pred_cam)
    pelvis_world = (cam_R.mT @ pelvis_cam.unsqueeze(-1)).squeeze(-1)
    cam_T_world = pelvis_world - trans
    cam_T = (cam_R @ cam_T_world.unsqueeze(-1)).squeeze(-1)
    
    return cam_R, cam_T

def compute_camera_pose(root_c_d6d, root_w):
    root_c = transforms.rotation_6d_to_matrix(root_c_d6d)  # Root orient in cam coord
    cam_R = root_c @ root_w.mT
    return cam_R


def reset_root_velocity(smpl, output, stationary, pred_ori, pred_vel, thr=0.7):
    b, f = pred_vel.shape[:2]
    
    stationary_mask = (stationary.clone().detach() > thr).unsqueeze(-1).float()
    poses_root = transforms.rotation_6d_to_matrix(pred_ori.clone().detach())
    vel_world = (poses_root[:, 1:] @ pred_vel.clone().detach().unsqueeze(-1)).squeeze(-1)
    
    output = smpl.get_output(body_pose=output.body_pose.clone().detach(),
                             global_orient=poses_root[:, 1:].reshape(-1, 1, 3, 3),
                             betas=output.betas.clone().detach(),
                             pose2rot=False)
    feet = output.feet.reshape(b, f, 4, 3)
    feet_vel = feet[:, 1:] - feet[:, :-1] + vel_world[:, 1:].unsqueeze(-2)
    feet_vel = torch.cat((torch.zeros_like(feet_vel[:, :1]), feet_vel), dim=1)
    
    stationary_vel = feet_vel * stationary_mask
    del_vel = stationary_vel.sum(dim=2) / ((stationary_vel != 0).sum(dim=2) + 1e-4)
    vel_world_update = vel_world - del_vel
    
    vel_root = (poses_root[:, 1:].mT @ vel_world_update.unsqueeze(-1)).squeeze(-1)
    
    return vel_root