101 lines
3.9 KiB
Python
101 lines
3.9 KiB
Python
'''
|
|
@ Date: 2020-11-20 13:34:54
|
|
@ Author: Qing Shuai
|
|
@ LastEditors: Qing Shuai
|
|
@ LastEditTime: 2021-04-13 20:31:49
|
|
@ FilePath: /EasyMocapRelease/easymocap/smplmodel/body_param.py
|
|
'''
|
|
import numpy as np
|
|
from os.path import join
|
|
|
|
def merge_params(param_list, share_shape=True):
|
|
output = {}
|
|
for key in ['poses', 'shapes', 'Rh', 'Th', 'expression']:
|
|
if key in param_list[0].keys():
|
|
output[key] = np.vstack([v[key] for v in param_list])
|
|
if share_shape:
|
|
output['shapes'] = output['shapes'].mean(axis=0, keepdims=True)
|
|
return output
|
|
|
|
def select_nf(params_all, nf):
|
|
output = {}
|
|
for key in ['poses', 'Rh', 'Th']:
|
|
output[key] = params_all[key][nf:nf+1, :]
|
|
if 'expression' in params_all.keys():
|
|
output['expression'] = params_all['expression'][nf:nf+1, :]
|
|
if params_all['shapes'].shape[0] == 1:
|
|
output['shapes'] = params_all['shapes']
|
|
else:
|
|
output['shapes'] = params_all['shapes'][nf:nf+1, :]
|
|
return output
|
|
|
|
NUM_POSES = {'smpl': 72, 'smplh': 78, 'smplx': 66 + 12 + 9}
|
|
NUM_EXPR = 10
|
|
|
|
def init_params(nFrames=1, model_type='smpl'):
|
|
params = {
|
|
'poses': np.zeros((nFrames, NUM_POSES[model_type])),
|
|
'shapes': np.zeros((1, 10)),
|
|
'Rh': np.zeros((nFrames, 3)),
|
|
'Th': np.zeros((nFrames, 3)),
|
|
}
|
|
if model_type == 'smplx':
|
|
params['expression'] = np.zeros((nFrames, NUM_EXPR))
|
|
return params
|
|
|
|
def check_params(body_params, model_type):
|
|
nFrames = body_params['poses'].shape[0]
|
|
if body_params['poses'].shape[1] != NUM_POSES[model_type]:
|
|
body_params['poses'] = np.hstack((body_params['poses'], np.zeros((nFrames, NUM_POSES[model_type] - body_params['poses'].shape[1]))))
|
|
if model_type == 'smplx' and 'expression' not in body_params.keys():
|
|
body_params['expression'] = np.zeros((nFrames, NUM_EXPR))
|
|
return body_params
|
|
|
|
def load_model(gender='neutral', use_cuda=True, model_type='smpl', skel_type='body25', device=None, model_path='data/smplx'):
|
|
# prepare SMPL model
|
|
# print('[Load model {}/{}]'.format(model_type, gender))
|
|
import torch
|
|
if device is None:
|
|
if use_cuda and torch.cuda.is_available():
|
|
device = torch.device('cuda')
|
|
else:
|
|
device = torch.device('cpu')
|
|
from .body_model import SMPLlayer
|
|
if model_type == 'smpl':
|
|
if skel_type == 'body25':
|
|
reg_path = join(model_path, 'J_regressor_body25.npy')
|
|
elif skel_type == 'h36m':
|
|
reg_path = join(model_path, 'J_regressor_h36m.npy')
|
|
else:
|
|
raise NotImplementedError
|
|
body_model = SMPLlayer(join(model_path, 'smpl'), gender=gender, device=device,
|
|
regressor_path=reg_path)
|
|
elif model_type == 'smplh':
|
|
body_model = SMPLlayer(join(model_path, 'smplh/SMPLH_MALE.pkl'), model_type='smplh', gender=gender, device=device,
|
|
regressor_path=join(model_path, 'J_regressor_body25_smplh.txt'))
|
|
elif model_type == 'smplx':
|
|
body_model = SMPLlayer(join(model_path, 'smplx/SMPLX_{}.pkl'.format(gender.upper())), model_type='smplx', gender=gender, device=device,
|
|
regressor_path=join(model_path, 'J_regressor_body25_smplx.txt'))
|
|
else:
|
|
body_model = None
|
|
body_model.to(device)
|
|
return body_model
|
|
|
|
def check_keypoints(keypoints2d, WEIGHT_DEBUFF=1, min_conf=0.3):
|
|
# keypoints2d: nFrames, nJoints, 3
|
|
#
|
|
# wrong feet
|
|
# if keypoints2d.shape[-2] > 25 + 42:
|
|
# keypoints2d[..., 0, 2] = 0
|
|
# keypoints2d[..., [15, 16, 17, 18], -1] = 0
|
|
# keypoints2d[..., [19, 20, 21, 22, 23, 24], -1] /= 2
|
|
if keypoints2d.shape[-2] > 25:
|
|
# set the hand keypoints
|
|
keypoints2d[..., 25, :] = keypoints2d[..., 7, :]
|
|
keypoints2d[..., 46, :] = keypoints2d[..., 4, :]
|
|
keypoints2d[..., 25:, -1] *= WEIGHT_DEBUFF
|
|
# reduce the confidence of hand and face
|
|
MIN_CONF = min_conf
|
|
conf = keypoints2d[..., -1]
|
|
conf[conf<MIN_CONF] = 0
|
|
return keypoints2d |