Method forward

detrsmpl/models/body_models/star.py:215–312

Forward pass for the STAR model. Args: global_orient: torch.tensor, optional, shape Bx3 Global orientation (rotation) of the body. If given, ignore the member variable and use it as the global rotation of the body. Useful if someon

(self,
                global_orient: Optional[torch.Tensor] = None,
                body_pose: Optional[torch.Tensor] = None,
                betas: Optional[torch.Tensor] = None,
                transl: Optional[torch.Tensor] = None,
                return_verts: bool = True,
                return_full_pose: bool = True)

        self.R = None

    def forward(self,
                global_orient: Optional[torch.Tensor] = None,
                body_pose: Optional[torch.Tensor] = None,
                betas: Optional[torch.Tensor] = None,
                transl: Optional[torch.Tensor] = None,
                return_verts: bool = True,
                return_full_pose: bool = True) -> torch.Tensor:
        """Forward pass for the STAR model.

        Args:
            global_orient: torch.tensor, optional, shape Bx3
                Global orientation (rotation) of the body. If given, ignore the
                member variable and use it as the global rotation of the body.
                Useful if someone wishes to predicts this with an external
                model. (default=None)
            body_pose: torch.Tensor, shape Bx(J*3)
                Pose parameters for the STAR model. It should be a tensor that
                contains joint rotations in axis-angle format. If given, ignore
                the member variable and use it as the body parameters.
                (default=None)
            betas: torch.Tensor, shape Bx10
                Shape parameters for the STAR model. If given, ignore the
                member variable and use it as shape parameters. (default=None)
            transl: torch.Tensor, shape Bx3
                Translation vector for the STAR model. If given, ignore the
                member variable and use it as the translation of the body.
                (default=None)
        Returns:
            output: Contains output parameters and attributes corresponding
            to other body models.
        """
        global_orient = (global_orient
                         if global_orient is not None else self.global_orient)
        body_pose = body_pose if body_pose is not None else self.body_pose
        betas = betas if betas is not None else self.betas
        apply_transl = transl is not None or hasattr(self, 'transl')
        if transl is None and hasattr(self, 'transl'):
            transl = self.transl

        batch_size = body_pose.shape[0]
        v_template = self.v_template[None, :]
        shapedirs = self.shapedirs.view(-1, self.num_betas)[None, :].expand(
            batch_size, -1, -1)
        beta = betas[:, :, None]
        v_shaped = torch.matmul(shapedirs, beta).view(-1, 6890, 3) + v_template
        J = torch.einsum('bik,ji->bjk', [v_shaped, self.J_regressor])

        pose_quat = self.normalize_quaternion(body_pose.view(-1, 3)).view(
            batch_size, -1)
        pose_feat = torch.cat((pose_quat[:, 4:], beta[:, 1]), 1)

        R = aa_to_rotmat(body_pose.view(-1, 3)).view(batch_size, 24, 3, 3)
        R = R.view(batch_size, 24, 3, 3)

        posedirs = self.posedirs[None, :].expand(batch_size, -1, -1)
        v_posed = v_shaped + torch.matmul(
            posedirs, pose_feat[:, :, None]).view(-1, 6890, 3)