Spaces:
Running
on
Zero
Running
on
Zero
| # Multi-HMR | |
| # Copyright (c) 2024-present NAVER Corp. | |
| # CC BY-NC-SA 4.0 license | |
| import torch | |
| from torch import nn | |
| from torch import nn | |
| import smplx | |
| import torch | |
| import numpy as np | |
| import pose_utils | |
| from pose_utils import inverse_perspective_projection, perspective_projection | |
| import roma | |
| import pickle | |
| import os | |
| from pose_utils.constants_service import SMPLX_DIR | |
| from pose_utils.rot6d import rotation_6d_to_matrix | |
| from smplx.lbs import vertices2joints | |
| class SMPL_Layer(nn.Module): | |
| """ | |
| Extension of the SMPL Layer with information about the camera for (inverse) projection the camera plane. | |
| """ | |
| def __init__( | |
| self, | |
| smpl_dir, | |
| type="smplx", | |
| gender="neutral", | |
| num_betas=10, | |
| kid=False, | |
| person_center=None, | |
| *args, | |
| **kwargs, | |
| ): | |
| super().__init__() | |
| # Args | |
| assert type == "smplx" | |
| self.type = type | |
| self.kid = kid | |
| self.num_betas = num_betas | |
| self.bm_x = smplx.create( | |
| smpl_dir, "smplx", gender=gender, use_pca=False, flat_hand_mean=True, num_betas=num_betas | |
| ) | |
| # Primary keypoint - root | |
| self.joint_names = eval(f"pose_utils.get_{self.type}_joint_names")() | |
| self.person_center = person_center | |
| self.person_center_idx = None | |
| if self.person_center is not None: | |
| self.person_center_idx = self.joint_names.index(self.person_center) | |
| def forward( | |
| self, | |
| pose, | |
| shape, | |
| loc, | |
| dist, | |
| transl, | |
| K, | |
| expression=None, # facial expression | |
| rot6d=False, | |
| j_regressor=None, | |
| ): | |
| """ | |
| Args: | |
| - pose: pose of the person in axis-angle - torch.Tensor [bs,24,3] | |
| - shape: torch.Tensor [bs,10] | |
| - loc: 2D location of the pelvis in pixel space - torch.Tensor [bs,2] | |
| - dist: distance of the pelvis from the camera in m - torch.Tensor [bs,1] | |
| Return: | |
| - dict containing a bunch of useful information about each person | |
| """ | |
| if loc is not None and dist is not None: | |
| assert pose.shape[0] == shape.shape[0] == loc.shape[0] == dist.shape[0] | |
| POSE_TYPE_LENGTH = 6 if rot6d else 3 | |
| if self.type == "smpl": | |
| assert len(pose.shape) == 3 and list(pose.shape[1:]) == [24, POSE_TYPE_LENGTH] | |
| elif self.type == "smplx": | |
| assert len(pose.shape) == 3 and list(pose.shape[1:]) == [ | |
| 53, | |
| POSE_TYPE_LENGTH, | |
| ] # taking root_orient, body_pose, lhand, rhan and jaw for the moment | |
| else: | |
| raise NameError | |
| assert len(shape.shape) == 2 and ( | |
| list(shape.shape[1:]) == [self.num_betas] or list(shape.shape[1:]) == [self.num_betas + 1] | |
| ) | |
| if loc is not None and dist is not None: | |
| assert len(loc.shape) == 2 and list(loc.shape[1:]) == [2] | |
| assert len(dist.shape) == 2 and list(dist.shape[1:]) == [1] | |
| bs = pose.shape[0] | |
| out = {} | |
| # No humans | |
| if bs == 0: | |
| return {} | |
| # Low dimensional parameters | |
| kwargs_pose = { | |
| "betas": shape, | |
| } | |
| kwargs_pose["global_orient"] = self.bm_x.global_orient.repeat(bs, 1) | |
| kwargs_pose["body_pose"] = pose[:, 1:22].flatten(1) | |
| kwargs_pose["left_hand_pose"] = pose[:, 22:37].flatten(1) | |
| kwargs_pose["right_hand_pose"] = pose[:, 37:52].flatten(1) | |
| kwargs_pose["jaw_pose"] = pose[:, 52:53].flatten(1) | |
| if expression is not None: | |
| kwargs_pose["expression"] = expression.flatten(1) # [bs,10] | |
| else: | |
| kwargs_pose["expression"] = self.bm_x.expression.repeat(bs, 1) | |
| # default - to be generalized | |
| kwargs_pose["leye_pose"] = self.bm_x.leye_pose.repeat(bs, 1) | |
| kwargs_pose["reye_pose"] = self.bm_x.reye_pose.repeat(bs, 1) | |
| # kwargs_pose['pose2rot'] = not rot6d | |
| # Forward using the parametric 3d model SMPL-X layer | |
| output = self.bm_x(pose2rot=not rot6d, **kwargs_pose) | |
| verts = output.vertices | |
| j3d = output.joints # 45 joints | |
| if rot6d: | |
| R = rotation_6d_to_matrix(pose[:, 0]) | |
| else: | |
| R = roma.rotvec_to_rotmat(pose[:, 0]) | |
| # Apply global orientation on 3D points | |
| pelvis = j3d[:, [0]] | |
| j3d = (R.unsqueeze(1) @ (j3d - pelvis).unsqueeze(-1)).squeeze(-1) | |
| # Apply global orientation on 3D points - bis | |
| verts = (R.unsqueeze(1) @ (verts - pelvis).unsqueeze(-1)).squeeze(-1) | |
| # Location of the person in 3D | |
| if transl is None: | |
| if K.dtype == torch.float16: | |
| # because of torch.inverse - not working with float16 at the moment | |
| transl = inverse_perspective_projection( | |
| loc.unsqueeze(1).float(), K.float(), dist.unsqueeze(1).float() | |
| )[:, 0] | |
| transl = transl.half() | |
| else: | |
| transl = inverse_perspective_projection(loc.unsqueeze(1), K, dist.unsqueeze(1))[:, 0] | |
| # Updating transl if we choose a certain person center | |
| transl_up = transl.clone() | |
| # Definition of the translation depend on the args: 1) vanilla SMPL - 2) computed from a given joint | |
| if self.person_center_idx is None: | |
| # Add pelvis to transl - standard way for SMPLX layer | |
| transl_up = transl_up + pelvis[:, 0] | |
| else: | |
| # Center around the joint because teh translation is computed from this joint | |
| person_center = j3d[:, [self.person_center_idx]] | |
| verts = verts - person_center | |
| j3d = j3d - person_center | |
| # Moving into the camera coordinate system | |
| j3d_cam = j3d + transl_up.unsqueeze(1) | |
| verts_cam = verts + transl_up.unsqueeze(1) | |
| # Projection in camera plane | |
| if j_regressor is not None: | |
| # for smplify | |
| j3d_cam = vertices2joints(j_regressor, verts_cam) | |
| j2d = perspective_projection(j3d_cam, K) | |
| v2d = perspective_projection(verts_cam, K) | |
| out.update( | |
| { | |
| "v3d": verts_cam, # in 3d camera space | |
| "j3d": j3d_cam, # in 3d camera space | |
| "j2d": j2d, | |
| "v2d": v2d, | |
| "transl": transl, # translation of the primary keypoint | |
| "transl_pelvis": transl.unsqueeze(1) - person_center - pelvis, # root=pelvis | |
| "j3d_world": output.joints, | |
| } | |
| ) | |
| return out | |
| def forward_local(self, pose, shape): | |
| N, J, L = pose.shape | |
| if N < 1: | |
| return None | |
| kwargs_pose = { | |
| "betas": shape, | |
| } | |
| if J == 53: | |
| kwargs_pose["global_orient"] = self.bm_x.global_orient.repeat(N, 1) | |
| kwargs_pose["body_pose"] = pose[:, 1:22].flatten(1) | |
| kwargs_pose["left_hand_pose"] = pose[:, 22:37].flatten(1) | |
| kwargs_pose["right_hand_pose"] = pose[:, 37:52].flatten(1) | |
| kwargs_pose["jaw_pose"] = pose[:, 52:53].flatten(1) | |
| elif J==55: | |
| kwargs_pose["global_orient"] = self.bm_x.global_orient.repeat(N, 1) | |
| kwargs_pose["body_pose"] = pose[:, 1:22].flatten(1) | |
| kwargs_pose["left_hand_pose"] = pose[:, 25:40].flatten(1) | |
| kwargs_pose["right_hand_pose"] = pose[:, 40:55].flatten(1) | |
| kwargs_pose["jaw_pose"] = pose[:, 22:23].flatten(1) | |
| else: | |
| raise ValueError(f"pose dim error, should be 53 or 55, but got {J}") | |
| kwargs_pose["expression"] = self.bm_x.expression.repeat(N, 1) | |
| # default - to be generalized | |
| kwargs_pose["leye_pose"] = self.bm_x.leye_pose.repeat(N, 1) | |
| kwargs_pose["reye_pose"] = self.bm_x.reye_pose.repeat(N, 1) | |
| output = self.bm_x(**kwargs_pose) | |
| return output | |
| def convert_standard_pose(self, poses): | |
| # pose: N, J, 3 | |
| n = poses.shape[0] | |
| poses = torch.cat( | |
| [ | |
| poses[:, :22], | |
| poses[:, 52:53], | |
| self.bm_x.leye_pose.repeat(n, 1, 1), | |
| self.bm_x.reye_pose.repeat(n, 1, 1), | |
| poses[:, 22:52], | |
| ], | |
| dim=1, | |
| ) | |
| return poses | |