| 1614 | |
| 1615 | |
| 1616 | class MANOLayer(MANO): |
| 1617 | def __init__(self, *args, **kwargs) -> None: |
| 1618 | """MANO as a layer model constructor.""" |
| 1619 | super(MANOLayer, self).__init__(create_global_orient=False, |
| 1620 | create_hand_pose=False, |
| 1621 | create_betas=False, |
| 1622 | create_transl=False, |
| 1623 | *args, |
| 1624 | **kwargs) |
| 1625 | |
| 1626 | def name(self) -> str: |
| 1627 | return 'MANO' |
| 1628 | |
| 1629 | def forward(self, |
| 1630 | betas: Optional[Tensor] = None, |
| 1631 | global_orient: Optional[Tensor] = None, |
| 1632 | hand_pose: Optional[Tensor] = None, |
| 1633 | transl: Optional[Tensor] = None, |
| 1634 | return_verts: bool = True, |
| 1635 | return_full_pose: bool = False, |
| 1636 | **kwargs) -> MANOOutput: |
| 1637 | """Forward pass for the MANO model.""" |
| 1638 | device, dtype = self.shapedirs.device, self.shapedirs.dtype |
| 1639 | if global_orient is None: |
| 1640 | batch_size = 1 |
| 1641 | global_orient = torch.zeros(3, device=device, dtype=dtype).view( |
| 1642 | 1, 1, 3).expand(batch_size, -1, -1).contiguous() |
| 1643 | else: |
| 1644 | batch_size = global_orient.shape[0] |
| 1645 | if hand_pose is None: |
| 1646 | hand_pose = torch.zeros(3, device=device, dtype=dtype).view( |
| 1647 | 1, 1, 3).expand(batch_size, 15, -1).contiguous() |
| 1648 | if betas is None: |
| 1649 | betas = torch.zeros([batch_size, self.num_betas], |
| 1650 | dtype=dtype, |
| 1651 | device=device) |
| 1652 | if transl is None: |
| 1653 | transl = torch.zeros([batch_size, 3], dtype=dtype, device=device) |
| 1654 | |
| 1655 | full_pose = torch.cat([global_orient, hand_pose], dim=1) |
| 1656 | vertices, joints = lbs(betas, |
| 1657 | full_pose, |
| 1658 | self.v_template, |
| 1659 | self.shapedirs, |
| 1660 | self.posedirs, |
| 1661 | self.J_regressor, |
| 1662 | self.parents, |
| 1663 | self.lbs_weights, |
| 1664 | pose2rot=True) |
| 1665 | |
| 1666 | if self.joint_mapper is not None: |
| 1667 | joints = self.joint_mapper(joints) |
| 1668 | |
| 1669 | if transl is not None: |
| 1670 | joints = joints + transl.unsqueeze(dim=1) |
| 1671 | vertices = vertices + transl.unsqueeze(dim=1) |
| 1672 | |
| 1673 | output = MANOOutput(vertices=vertices if return_verts else None, |