Class JointMapper

src/modules/pixie/pixie.py:119–132

            )

        class JointMapper(nn.Module):
            def __init__(self, joint_maps=None):
                super(JointMapper, self).__init__()
                if joint_maps is None:
                    self.joint_maps = joint_maps
                else:
                    self.register_buffer('joint_maps',
                                         torch.tensor(joint_maps, dtype=torch.long))

            def forward(self, joints, **kwargs):
                if self.joint_maps is None:
                    return joints
                else:
                    return torch.index_select(joints, 1, self.joint_maps)
        # import smplx
        # smplx_model_path = r'C:\Users\lithiumice\code\smplify-x\models'
        # joint_mapper = JointMapper(