fulyvv/EasyMocap
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

🐛 add num_shapes for smpl/smplh

Browse Source
This commit is contained in:
shuaiqing 2021-05-11 15:13:05 +08:00
parent 436c03c149
commit fb1d2908bd
1 changed files with 12 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
easymocap/smplmodel/body_model.py
View File

@ -2,8 +2,8 @@
@ Date: 2020-11-18 14:04:10 @ Date: 2020-11-18 14:04:10
@ Author: Qing Shuai @ Author: Qing Shuai
@ LastEditors: Qing Shuai @ LastEditors: Qing Shuai
@ LastEditTime: 2021-03-15 22:04:32 @ LastEditTime: 2021-05-11 15:09:44
@ FilePath: /EasyMocap/code/smplmodel/body_model.py @ FilePath: /EasyMocap/easymocap/smplmodel/body_model.py
''' '''
import torch import torch
import torch.nn as nn import torch.nn as nn
@ -40,6 +40,7 @@ def load_regressor(regressor_path):
return X_regressor return X_regressor
NUM_POSES = {'smpl': 72, 'smplh': 78, 'smplx': 66 + 12 + 9} NUM_POSES = {'smpl': 72, 'smplh': 78, 'smplx': 66 + 12 + 9}
NUM_SHAPES = 10
NUM_EXPR = 10 NUM_EXPR = 10
class SMPLlayer(nn.Module): class SMPLlayer(nn.Module):
def __init__(self, model_path, model_type='smpl', gender='neutral', device=None, def __init__(self, model_path, model_type='smpl', gender='neutral', device=None,
@ -82,6 +83,8 @@ class SMPLlayer(nn.Module):
self.num_expression_coeffs = 10 self.num_expression_coeffs = 10
self.num_shapes = 10 self.num_shapes = 10
self.shapedirs = self.shapedirs[:, :, :self.num_shapes+self.num_expression_coeffs] self.shapedirs = self.shapedirs[:, :, :self.num_shapes+self.num_expression_coeffs]
elif self.model_type in ['smpl', 'smplh']:
self.shapedirs = self.shapedirs[:, :, :NUM_SHAPES]
# joints regressor # joints regressor
if regressor_path is not None: if regressor_path is not None:
X_regressor = load_regressor(regressor_path) X_regressor = load_regressor(regressor_path)
@ -222,8 +225,10 @@ class SMPLlayer(nn.Module):
dtype, device = self.dtype, self.device dtype, device = self.dtype, self.device
poses = to_tensor(poses, dtype, device) poses = to_tensor(poses, dtype, device)
shapes = to_tensor(shapes, dtype, device) shapes = to_tensor(shapes, dtype, device)
Rh = to_tensor(Rh, dtype, device) if Rh is not None:
Th = to_tensor(Th, dtype, device) Rh = to_tensor(Rh, dtype, device)
if Th is not None:
Th = to_tensor(Th, dtype, device)
if expression is not None: if expression is not None:
expression = to_tensor(expression, dtype, device) expression = to_tensor(expression, dtype, device)
@ -231,6 +236,9 @@ class SMPLlayer(nn.Module):
# process Rh, Th # process Rh, Th
if Rh is None: if Rh is None:
Rh = torch.zeros(bn, 3, device=poses.device) Rh = torch.zeros(bn, 3, device=poses.device)
if Th is None:
Th = torch.zeros(bn, 3, device=poses.device)
if len(Rh.shape) == 2: # angle-axis if len(Rh.shape) == 2: # angle-axis
rot = batch_rodrigues(Rh) rot = batch_rodrigues(Rh)
else: else: