def _fk()

in empose/bodymodels/smpl.py [0:0]

• 21 lines of code
• 6 McCabe index (conditional complexity)

    def _fk(self, poses_body, betas, poses_root=None, trans=None, normalize_root=False):
        """
        Evaluate the SMPL body model, i.e. compute joint and vertex positions given pose and shape parameters.
        :param poses_body: A tensor of shape (N, 21*3) specifying the body pose.
        :param betas: A tensor of shape (N, N_BETAS) specifying the SMPL parameters. If N_BETAS > self.num_betas, the
          excessive parameters are shaved off.
        :param poses_root: Root orientation as a tensor of shape (N, 3) or None.
        :param trans: Root translation as a tensor of shape (N, 3) or None.
        :param normalize_root: Normalize the root into the coordinate system defined by the first frame (i.e.
          root orientation at frame 0 is the identity and the root position is at the origin).
        :return: Vertex and joint positions of the posed SMPL mesh.
        """
        assert poses_body.shape[1] >= C.N_JOINTS * 3

        batch_size = poses_body.shape[0]
        device = poses_body.device

        # The body model expects hand poses, so supply a dummy value.
        poses_hands = torch.zeros([batch_size, C.N_JOINTS_HAND * 3 * 2]).to(dtype=poses_body.dtype, device=device)

        if poses_root is None:
            poses_root = torch.zeros([batch_size, 3]).to(dtype=poses_body.dtype, device=device)

        if trans is None:
            trans = torch.zeros([batch_size, 3]).to(dtype=poses_body.dtype, device=device)

        # Broadcast shape parameters.
        if len(betas.shape) == 1 or betas.shape[0] == 1:
            betas = betas.repeat(poses_body.shape[0], 1)
        betas = betas[:, :self.num_betas]

        if normalize_root:
            # Make everything relative to the first root orientation.
            root_ori = aa2rot(poses_root)
            first_root_ori = torch.inverse(root_ori[0:1])
            root_ori = torch.matmul(first_root_ori, root_ori)
            poses_root = rot2aa(root_ori)
            trans = torch.matmul(first_root_ori.unsqueeze(0), trans.unsqueeze(-1)).squeeze()
            trans = trans - trans[0:1]

        body = self.bm(root_orient=poses_root, pose_body=poses_body, betas=betas, pose_hand=poses_hands, trans=trans)
        return body.v, body.Jtr