(
motion_tensor: torch.Tensor,
smplx_model,
device: torch.device,
noise_std: float = 0.0,
)
| 51 | |
| 52 | |
| 53 | def run_smplx( |
| 54 | motion_tensor: torch.Tensor, |
| 55 | smplx_model, |
| 56 | device: torch.device, |
| 57 | noise_std: float = 0.0, |
| 58 | ) -> dict: |
| 59 | smpl_data, recovered_joints = motion_rep_to_SMPL( |
| 60 | motion_tensor, |
| 61 | recover_from_velocity=True, |
| 62 | equal_length=False, |
| 63 | ) |
| 64 | if noise_std > 0.0: |
| 65 | for key, value in smpl_data.items(): |
| 66 | if torch.is_tensor(value): |
| 67 | smpl_data[key] = value + torch.randn_like(value) * noise_std |
| 68 | recovered_joints = recovered_joints + torch.randn_like(recovered_joints) * noise_std |
| 69 | smpl_params = {k: v.to(device) for k, v in smpl_data.items()} |
| 70 | recovered_joints = recovered_joints.to(device) |
| 71 | betas = torch.zeros( |
| 72 | (motion_tensor.shape[0], smplx_model.num_betas), |
| 73 | dtype=smpl_params["body_pose"].dtype, |
| 74 | device=device, |
| 75 | ) |
| 76 | with torch.no_grad(): |
| 77 | smpl_output = smplx_model(**smpl_params, betas=betas) |
| 78 | |
| 79 | convert_matrix = COORD_CONVERSION.to(device) |
| 80 | vertices = torch.einsum("ij,tvj->tvi", convert_matrix, smpl_output.vertices) |
| 81 | joints = torch.einsum("ij,tvj->tvi", convert_matrix, smpl_output.joints[:, :22]) |
| 82 | recovered_joints = torch.einsum("ij,tvj->tvi", convert_matrix, recovered_joints) |
| 83 | smpl_params["transl"] = torch.einsum("ij,tj->ti", convert_matrix, smpl_params["transl"]) |
| 84 | |
| 85 | vertices = vertices.detach().cpu().numpy().astype("float32") |
| 86 | joints = joints.detach().cpu().numpy().astype("float32") |
| 87 | recovered_np = recovered_joints.detach().cpu().numpy().astype("float32") |
| 88 | pose = torch.cat( |
| 89 | [smpl_params["global_orient"].cpu(), smpl_params["body_pose"].cpu()], dim=1 |
| 90 | ) |
| 91 | pose = pose.view(pose.shape[0], -1, 3).numpy().astype("float32") |
| 92 | smpl_param_np = {k: v.detach().cpu().numpy().astype("float32") for k, v in smpl_params.items()} |
| 93 | |
| 94 | return { |
| 95 | "pose": pose, |
| 96 | "joints": joints, |
| 97 | "vertices": vertices, |
| 98 | "smpl_params": smpl_param_np, |
| 99 | "recovered_joints": recovered_np, |
| 100 | } |
| 101 | |
| 102 | |
| 103 | def render_video_from_vertices( |
no test coverage detected