Forward function. Args: *args: extra arguments for SMPL return_verts: whether to return vertices return_full_pose: whether to return full pose parameters **kwargs: extra arguments for SMPL Returns: output: contains output
(self,
*args,
return_verts: bool = True,
return_full_pose: bool = False,
**kwargs)
| 78 | self.body_part_segmentation = body_segmentation('smplx') |
| 79 | |
| 80 | def forward(self, |
| 81 | *args, |
| 82 | return_verts: bool = True, |
| 83 | return_full_pose: bool = False, |
| 84 | **kwargs) -> dict: |
| 85 | """Forward function. |
| 86 | |
| 87 | Args: |
| 88 | *args: extra arguments for SMPL |
| 89 | return_verts: whether to return vertices |
| 90 | return_full_pose: whether to return full pose parameters |
| 91 | **kwargs: extra arguments for SMPL |
| 92 | |
| 93 | Returns: |
| 94 | output: contains output parameters and attributes |
| 95 | """ |
| 96 | |
| 97 | kwargs['get_skin'] = True |
| 98 | smplx_output = super(SMPLX, self).forward(*args, **kwargs) |
| 99 | |
| 100 | if not hasattr(self, 'joints_regressor'): |
| 101 | joints = smplx_output.joints |
| 102 | else: |
| 103 | joints = vertices2joints(self.joints_regressor, |
| 104 | smplx_output.vertices) |
| 105 | |
| 106 | if hasattr(self, 'joints_regressor_extra'): |
| 107 | extra_joints = vertices2joints(self.joints_regressor_extra, |
| 108 | smplx_output.vertices) |
| 109 | joints = torch.cat([joints, extra_joints], dim=1) |
| 110 | |
| 111 | joints, joint_mask = convert_kps(joints, |
| 112 | src=self.keypoint_src, |
| 113 | dst=self.keypoint_dst, |
| 114 | approximate=self.keypoint_approximate) |
| 115 | if isinstance(joint_mask, np.ndarray): |
| 116 | joint_mask = torch.tensor(joint_mask, |
| 117 | dtype=torch.uint8, |
| 118 | device=joints.device) |
| 119 | |
| 120 | batch_size = joints.shape[0] |
| 121 | joint_mask = joint_mask.reshape(1, -1).expand(batch_size, -1) |
| 122 | |
| 123 | output = dict(global_orient=smplx_output.global_orient, |
| 124 | body_pose=smplx_output.body_pose, |
| 125 | joints=joints, |
| 126 | joint_mask=joint_mask, |
| 127 | keypoints=torch.cat([joints, joint_mask[:, :, None]], |
| 128 | dim=-1), |
| 129 | betas=smplx_output.betas) |
| 130 | |
| 131 | if return_verts: |
| 132 | output['vertices'] = smplx_output.vertices |
| 133 | if return_full_pose: |
| 134 | output['full_pose'] = smplx_output.full_pose |
| 135 | |
| 136 | return output |
| 137 |
no test coverage detected