Spaces:
Sleeping
Sleeping
File size: 2,417 Bytes
5ac1897 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
from lib.kits.basic import *
from smplx import SMPL
from lib.platform import PM
from lib.body_models.skel_wrapper import SKELWrapper as SKEL
from lib.body_models.smpl_wrapper import SMPLWrapper
def make_SMPL(gender='neutral', device='cuda:0'):
return SMPL(
gender = gender,
model_path = PM.inputs / 'body_models' / 'smpl',
).to(device)
def make_SMPL_hmr2(gender='neutral', device='cuda:0'):
''' SKEL doesn't have neutral model, so align with SKEL, using male. '''
return SMPLWrapper(
gender = gender,
model_path = PM.inputs / 'body_models' / 'smpl',
num_body_joints = 23,
joint_regressor_extra = PM.inputs / 'body_models/SMPL_to_J19.pkl',
).to(device)
def make_SKEL(gender='male', device='cuda:0'):
''' We don't have neutral model for SKEL, so use male for now. '''
return make_SKEL_mix_joints(gender, device)
def make_SKEL_smpl_joints(gender='male', device='cuda:0'):
''' We don't have neutral model for SKEL, so use male for now. '''
return SKEL(
gender = gender,
model_path = PM.inputs / 'body_models' / 'skel',
joint_regressor_extra = PM.inputs / 'body_models' / 'SMPL_to_J19.pkl',
joint_regressor_custom = PM.inputs / 'body_models' / 'J_regressor_SMPL_MALE.pkl',
).to(device)
def make_SKEL_skel_joints(gender='male', device='cuda:0'):
''' We don't have neutral model for SKEL, so use male for now. '''
return SKEL(
gender = gender,
model_path = PM.inputs / 'body_models' / 'skel',
joint_regressor_extra = PM.inputs / 'body_models' / 'SMPL_to_J19.pkl',
).to(device)
def make_SKEL_mix_joints(gender='male', device='cuda:0'):
''' We don't have neutral model for SKEL, so use male for now. '''
return SKEL(
gender = gender,
model_path = PM.inputs / 'body_models' / 'skel',
joint_regressor_extra = PM.inputs / 'body_models' / 'SMPL_to_J19.pkl',
joint_regressor_custom = PM.inputs / 'body_models' / 'J_regressor_SKEL_mix_MALE.pkl',
).to(device)
def make_SMPLX_moyo(v_template_path:Union[str, Path], batch_size:int=1, device='cuda:0'):
from lib.body_models.moyo_smplx_wrapper import MoYoSMPLX
return MoYoSMPLX(
model_path = PM.inputs / 'body_models' / 'smplx',
v_template_path = v_template_path,
batch_size = batch_size,
device = device,
) |