def __init__(self, config):
super(FlameDecoder, self).__init__()
print("Initializing a Flame decoder")
with open(config.flame_model_path, 'rb') as f:
self.flame_model = Struct(**pickle.load(f, encoding='latin1'))
self.dtype = torch.float32
self.batch_size = config.batch_size
self.faces = self.flame_model.f
self.register_buffer(
'faces_tensor',
to_tensor(to_np(self.faces, dtype=np.int64), dtype=torch.long))
# Eyeball and neck rotation
default_eyball_pose = torch.zeros((self.batch_size, 6),
dtype=self.dtype,
requires_grad=False)
self.register_parameter(
'eye_pose', nn.Parameter(default_eyball_pose, requires_grad=False))
# Fixing 3D translation since we use translation in the image plane
#self.use_3D_translation = config.use_3D_translation
# The vertices of the template model
self.register_buffer(
'v_template',
to_tensor(to_np(self.flame_model.v_template), dtype=self.dtype))
# The shape components
shapedirs = self.flame_model.shapedirs
# The shape components
self.register_buffer('shapedirs',
to_tensor(to_np(shapedirs), dtype=self.dtype))
j_regressor = to_tensor(to_np(self.flame_model.J_regressor),
dtype=self.dtype)
self.register_buffer('J_regressor', j_regressor)
# Pose blend shape basis
num_pose_basis = self.flame_model.posedirs.shape[-1]
posedirs = np.reshape(self.flame_model.posedirs,
[-1, num_pose_basis]).T
self.register_buffer('posedirs',
to_tensor(to_np(posedirs), dtype=self.dtype))
# indices of parents for each joints
parents = to_tensor(to_np(self.flame_model.kintree_table[0])).long()
parents[0] = -1
self.register_buffer('parents', parents)
self.register_buffer(
'lbs_weights',
to_tensor(to_np(self.flame_model.weights), dtype=self.dtype))
def __init__(self,
model_path,
data_struct=None,
create_betas=True,
betas=None,
create_global_orient=True,
global_orient=None,
create_body_pose=True,
body_pose=None,
create_transl=True,
transl=None,
dtype=torch.float32,
batch_size=1,
joint_mapper=None,
gender='neutral',
vertex_ids=None,
**kwargs):
''' SMPL model constructor
Parameters
----------
model_path: str
The path to the folder or to the file where the model
parameters are stored
data_struct: Strct
A struct object. If given, then the parameters of the model are
read from the object. Otherwise, the model tries to read the
parameters from the given `model_path`. (default = None)
create_global_orient: bool, optional
Flag for creating a member variable for the global orientation
of the body. (default = True)
global_orient: torch.tensor, optional, Bx3
The default value for the global orientation variable.
(default = None)
create_body_pose: bool, optional
Flag for creating a member variable for the pose of the body.
(default = True)
body_pose: torch.tensor, optional, Bx(Body Joints * 3)
The default value for the body pose variable.
(default = None)
create_betas: bool, optional
Flag for creating a member variable for the shape space
(default = True).
betas: torch.tensor, optional, Bx10
The default value for the shape member variable.
(default = None)
create_transl: bool, optional
Flag for creating a member variable for the translation
of the body. (default = True)
transl: torch.tensor, optional, Bx3
The default value for the transl variable.
(default = None)
dtype: torch.dtype, optional
The data type for the created variables
batch_size: int, optional
The batch size used for creating the member variables
joint_mapper: object, optional
An object that re-maps the joints. Useful if one wants to
re-order the SMPL joints to some other convention (e.g. MSCOCO)
(default = None)
gender: str, optional
Which gender to load
vertex_ids: dict, optional
A dictionary containing the indices of the extra vertices that
will be selected
'''
self.gender = gender
if data_struct is None:
if osp.isdir(model_path):
model_fn = 'SMPL_{}.{ext}'.format(gender.upper(), ext='pkl')
smpl_path = os.path.join(model_path, model_fn)
else:
smpl_path = model_path
assert osp.exists(smpl_path), 'Path {} does not exist!'.format(
smpl_path)
with open(smpl_path, 'rb') as smpl_file:
data_struct = Struct(
**pickle.load(smpl_file, encoding='latin1'))
super(SMPL, self).__init__()
self.batch_size = batch_size
if vertex_ids is None:
# SMPL and SMPL-H share the same topology, so any extra joints can
# be drawn from the same place
vertex_ids = VERTEX_IDS['smplh']
self.dtype = dtype
self.joint_mapper = joint_mapper
self.vertex_joint_selector = VertexJointSelector(vertex_ids=vertex_ids,
**kwargs)
self.faces = data_struct.f
self.register_buffer(
'faces_tensor',
to_tensor(to_np(self.faces, dtype=np.int64), dtype=torch.long))
if create_betas:
if betas is None:
default_betas = torch.zeros([batch_size, self.NUM_BETAS],
dtype=dtype)
else:
if 'torch.Tensor' in str(type(betas)):
default_betas = betas.clone().detach()
else:
default_betas = torch.tensor(betas, dtype=dtype)
self.register_parameter(
'betas', nn.Parameter(default_betas, requires_grad=True))
# The tensor that contains the global rotation of the model
# It is separated from the pose of the joints in case we wish to
# optimize only over one of them
if create_global_orient:
if global_orient is None:
default_global_orient = torch.zeros([batch_size, 3],
dtype=dtype)
else:
if 'torch.Tensor' in str(type(global_orient)):
default_global_orient = global_orient.clone().detach()
else:
default_global_orient = torch.tensor(global_orient,
dtype=dtype)
global_orient = nn.Parameter(default_global_orient,
requires_grad=True)
self.register_parameter('global_orient', global_orient)
if create_body_pose:
if body_pose is None:
default_body_pose = torch.zeros(
[batch_size, self.NUM_BODY_JOINTS * 3], dtype=dtype)
else:
if 'torch.Tensor' in str(type(body_pose)):
default_body_pose = body_pose.clone().detach()
else:
default_body_pose = torch.tensor(body_pose, dtype=dtype)
self.register_parameter(
'body_pose', nn.Parameter(default_body_pose,
requires_grad=True))
if create_transl:
if transl is None:
default_transl = torch.zeros([batch_size, 3],
dtype=dtype,
requires_grad=True)
else:
default_transl = torch.tensor(transl, dtype=dtype)
self.register_parameter(
'transl', nn.Parameter(default_transl, requires_grad=True))
# The vertices of the template model
self.register_buffer(
'v_template', to_tensor(to_np(data_struct.v_template),
dtype=dtype))
# The shape components
shapedirs = data_struct.shapedirs
# The shape components
self.register_buffer('shapedirs',
to_tensor(to_np(shapedirs), dtype=dtype))
j_regressor = to_tensor(to_np(data_struct.J_regressor), dtype=dtype)
self.register_buffer('J_regressor', j_regressor)
# Pose blend shape basis: 6890 x 3 x 207, reshaped to 6890*3 x 207
num_pose_basis = data_struct.posedirs.shape[-1]
# 207 x 20670
posedirs = np.reshape(data_struct.posedirs, [-1, num_pose_basis]).T
self.register_buffer('posedirs', to_tensor(to_np(posedirs),
dtype=dtype))
# indices of parents for each joints
parents = to_tensor(to_np(data_struct.kintree_table[0])).long()
parents[0] = -1
self.register_buffer('parents', parents)
self.register_buffer(
'lbs_weights', to_tensor(to_np(data_struct.weights), dtype=dtype))
def __init__(self, config):
super(FLAME, self).__init__()
print("creating the FLAME Decoder")
with open(config.flame_model_path, 'rb') as f:
self.flame_model = Struct(**pickle.load(f, encoding='latin1'))
self.NECK_IDX = 1
self.batch_size = config.batch_size
self.dtype = torch.float32
self.use_face_contour = config.use_face_contour
self.faces = self.flame_model.f
self.register_buffer('faces_tensor',
to_tensor(to_np(self.faces, dtype=np.int64),
dtype=torch.long))
# Fixing remaining Shape betas
# There are total 300 shape parameters to control FLAME; But one can use the first few parameters to express
# the shape. For example 100 shape parameters are used for RingNet project
default_shape = torch.zeros([self.batch_size, 300-config.shape_params],
dtype=self.dtype, requires_grad=False)
self.register_parameter('shape_betas', nn.Parameter(default_shape,
requires_grad=False))
# Fixing remaining expression betas
# There are total 100 shape expression parameters to control FLAME; But one can use the first few parameters to express
# the expression. For example 50 expression parameters are used for RingNet project
default_exp = torch.zeros([self.batch_size, 100 - config.expression_params],
dtype=self.dtype, requires_grad=False)
self.register_parameter('expression_betas', nn.Parameter(default_exp,
requires_grad=False))
# Eyeball and neck rotation
default_eyball_pose = torch.zeros([self.batch_size, 6],
dtype=self.dtype, requires_grad=False)
self.register_parameter('eye_pose', nn.Parameter(default_eyball_pose,
requires_grad=False))
default_neck_pose = torch.zeros([self.batch_size, 3],
dtype=self.dtype, requires_grad=False)
self.register_parameter('neck_pose', nn.Parameter(default_neck_pose,
requires_grad=False))
# Fixing 3D translation since we use translation in the image plane
self.use_3D_translation = config.use_3D_translation
default_transl = torch.zeros([self.batch_size, 3],
dtype=self.dtype, requires_grad=False)
self.register_parameter(
'transl',
nn.Parameter(default_transl, requires_grad=False))
# The vertices of the template model
self.register_buffer('v_template',
to_tensor(to_np(self.flame_model.v_template),
dtype=self.dtype))
# The shape components
shapedirs = self.flame_model.shapedirs
# The shape components
self.register_buffer(
'shapedirs',
to_tensor(to_np(shapedirs), dtype=self.dtype))
j_regressor = to_tensor(to_np(
self.flame_model.J_regressor), dtype=self.dtype)
self.register_buffer('J_regressor', j_regressor)
# Pose blend shape basis
num_pose_basis = self.flame_model.posedirs.shape[-1]
posedirs = np.reshape(self.flame_model.posedirs, [-1, num_pose_basis]).T
self.register_buffer('posedirs',
to_tensor(to_np(posedirs), dtype=self.dtype))
# indices of parents for each joints
parents = to_tensor(to_np(self.flame_model.kintree_table[0])).long()
parents[0] = -1
self.register_buffer('parents', parents)
self.register_buffer('lbs_weights',
to_tensor(to_np(self.flame_model.weights), dtype=self.dtype))
# Static and Dynamic Landmark embeddings for FLAME
with open(config.static_landmark_embedding_path, 'rb') as f:
static_embeddings = Struct(**pickle.load(f, encoding='latin1'))
lmk_faces_idx = (static_embeddings.lmk_face_idx).astype(np.int64)
self.register_buffer('lmk_faces_idx',
torch.tensor(lmk_faces_idx, dtype=torch.long))
lmk_bary_coords = static_embeddings.lmk_b_coords
self.register_buffer('lmk_bary_coords',
torch.tensor(lmk_bary_coords, dtype=self.dtype))
if self.use_face_contour:
conture_embeddings = np.load(config.dynamic_landmark_embedding_path,
allow_pickle=True, encoding='latin1')
conture_embeddings = conture_embeddings[()]
dynamic_lmk_faces_idx = np.array(conture_embeddings['lmk_face_idx']).astype(np.int64)
dynamic_lmk_faces_idx = torch.tensor(
dynamic_lmk_faces_idx,
dtype=torch.long)
self.register_buffer('dynamic_lmk_faces_idx',
dynamic_lmk_faces_idx)
dynamic_lmk_bary_coords = conture_embeddings['lmk_b_coords']
dynamic_lmk_bary_coords = torch.tensor(
dynamic_lmk_bary_coords, dtype=self.dtype)
self.register_buffer('dynamic_lmk_bary_coords',
dynamic_lmk_bary_coords)
neck_kin_chain = []
curr_idx = torch.tensor(self.NECK_IDX, dtype=torch.long)
while curr_idx != -1:
neck_kin_chain.append(curr_idx)
curr_idx = self.parents[curr_idx]
self.register_buffer('neck_kin_chain',
torch.stack(neck_kin_chain))
def init_flame_buffers(self, config):
# The vertices of the template model
self.register_buffer(
'v_template',
to_tensor(to_np(self.flame_model.v_template), dtype=self.dtype))
self.register_buffer(
'faces_tensor',
to_tensor(to_np(self.faces, dtype=np.int64), dtype=torch.long))
# The shape components
shapedirs = self.flame_model.shapedirs
# The shape components
self.register_buffer('shapedirs',
to_tensor(to_np(shapedirs), dtype=self.dtype))
j_regressor = to_tensor(to_np(self.flame_model.J_regressor),
dtype=self.dtype)
self.register_buffer('J_regressor', j_regressor)
# Pose blend shape basis
num_pose_basis = self.flame_model.posedirs.shape[-1]
posedirs = np.reshape(self.flame_model.posedirs,
[-1, num_pose_basis]).T
self.register_buffer('posedirs',
to_tensor(to_np(posedirs), dtype=self.dtype))
# indices of parents for each joints
parents = to_tensor(to_np(self.flame_model.kintree_table[0])).long()
parents[0] = -1
self.register_buffer('parents', parents)
self.register_buffer(
'lbs_weights',
to_tensor(to_np(self.flame_model.weights), dtype=self.dtype))
# Static and Dynamic Landmark embeddings for FLAME
with open(config.static_landmark_embedding_path, 'rb') as f:
static_embeddings = Struct(**pickle.load(f, encoding='latin1'))
lmk_faces_idx = (static_embeddings.lmk_face_idx).astype(np.int64)
self.register_buffer('lmk_faces_idx',
torch.tensor(lmk_faces_idx, dtype=torch.long))
lmk_bary_coords = static_embeddings.lmk_b_coords
self.register_buffer('lmk_bary_coords',
torch.tensor(lmk_bary_coords, dtype=self.dtype))
if self.use_face_contour:
conture_embeddings = np.load(
config.dynamic_landmark_embedding_path,
allow_pickle=True,
encoding='latin1')
conture_embeddings = conture_embeddings[()]
dynamic_lmk_faces_idx = np.array(
conture_embeddings['lmk_face_idx']).astype(np.int64)
dynamic_lmk_faces_idx = torch.tensor(dynamic_lmk_faces_idx,
dtype=torch.long)
self.register_buffer('dynamic_lmk_faces_idx',
dynamic_lmk_faces_idx)
dynamic_lmk_bary_coords = conture_embeddings['lmk_b_coords']
dynamic_lmk_bary_coords = torch.tensor(dynamic_lmk_bary_coords,
dtype=self.dtype)
self.register_buffer('dynamic_lmk_bary_coords',
dynamic_lmk_bary_coords)
neck_kin_chain = []
curr_idx = torch.tensor(self.NECK_IDX, dtype=torch.long)
while curr_idx != -1:
neck_kin_chain.append(curr_idx)
curr_idx = self.parents[curr_idx]
self.register_buffer('neck_kin_chain', torch.stack(neck_kin_chain))