Spaces:
Runtime error
Runtime error
# -*- coding: utf-8 -*- | |
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is | |
# holder of all proprietary rights on this computer program. | |
# You can only use this computer program if you have closed | |
# a license agreement with MPG or you get the right to use the computer | |
# program from someone who is authorized to grant you that right. | |
# Any use of the computer program without a valid license is prohibited and | |
# liable to prosecution. | |
# | |
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung | |
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute | |
# for Intelligent Systems. All rights reserved. | |
# | |
# Contact: [email protected] | |
from __future__ import absolute_import | |
from __future__ import print_function | |
from __future__ import division | |
import numpy as np | |
import torch | |
import torch.nn as nn | |
from .utils import to_tensor | |
class VertexJointSelector(nn.Module): | |
def __init__(self, vertex_ids=None, | |
use_hands=True, | |
use_feet_keypoints=True, **kwargs): | |
super(VertexJointSelector, self).__init__() | |
extra_joints_idxs = [] | |
face_keyp_idxs = np.array([ | |
vertex_ids['nose'], | |
vertex_ids['reye'], | |
vertex_ids['leye'], | |
vertex_ids['rear'], | |
vertex_ids['lear']], dtype=np.int64) | |
extra_joints_idxs = np.concatenate([extra_joints_idxs, | |
face_keyp_idxs]) | |
if use_feet_keypoints: | |
feet_keyp_idxs = np.array([vertex_ids['LBigToe'], | |
vertex_ids['LSmallToe'], | |
vertex_ids['LHeel'], | |
vertex_ids['RBigToe'], | |
vertex_ids['RSmallToe'], | |
vertex_ids['RHeel']], dtype=np.int32) | |
extra_joints_idxs = np.concatenate( | |
[extra_joints_idxs, feet_keyp_idxs]) | |
if use_hands: | |
self.tip_names = ['thumb', 'index', 'middle', 'ring', 'pinky'] | |
tips_idxs = [] | |
for hand_id in ['l', 'r']: | |
for tip_name in self.tip_names: | |
tips_idxs.append(vertex_ids[hand_id + tip_name]) | |
extra_joints_idxs = np.concatenate( | |
[extra_joints_idxs, tips_idxs]) | |
self.register_buffer('extra_joints_idxs', | |
to_tensor(extra_joints_idxs, dtype=torch.long)) | |
def forward(self, vertices, joints): | |
extra_joints = torch.index_select(vertices, 1, self.extra_joints_idxs) | |
joints = torch.cat([joints, extra_joints], dim=1) | |
return joints | |