Method forward

detrsmpl/models/body_models/smplx.py:319–375

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)

        self.body_part_segmentation = body_segmentation('smplx')

    def forward(self,
                *args,
                return_verts: bool = True,
                return_full_pose: bool = False,
                **kwargs) -> dict:
        """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 parameters and attributes
        """

        kwargs['get_skin'] = True
        smplx_output = super(SMPLXLayer, self).forward(*args, **kwargs)

        if not hasattr(self, 'joints_regressor'):
            joints = smplx_output.joints
        else:
            joints = vertices2joints(self.joints_regressor,
                                     smplx_output.vertices)

        if hasattr(self, 'joints_regressor_extra'):
            extra_joints = vertices2joints(self.joints_regressor_extra,
                                           smplx_output.vertices)
            joints = torch.cat([joints, extra_joints], dim=1)

        joints, joint_mask = convert_kps(joints,
                                         src=self.keypoint_src,
                                         dst=self.keypoint_dst,
                                         approximate=self.keypoint_approximate)
        if isinstance(joint_mask, np.ndarray):
            joint_mask = torch.tensor(joint_mask,
                                      dtype=torch.uint8,
                                      device=joints.device)

        batch_size = joints.shape[0]
        joint_mask = joint_mask.reshape(1, -1).expand(batch_size, -1)

        output = dict(global_orient=smplx_output.global_orient,
                      body_pose=smplx_output.body_pose,
                      joints=joints,
                      joint_mask=joint_mask,
                      keypoints=torch.cat([joints, joint_mask[:, :, None]],
                                          dim=-1),
                      betas=smplx_output.betas)

        if return_verts:
            output['vertices'] = smplx_output.vertices
        if return_full_pose:
            output['full_pose'] = smplx_output.full_pose

        return output