Spaces:
Sleeping
Sleeping
File size: 4,482 Bytes
99afdfe |
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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
import os
import sys
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
sys.path.append(os.getcwd())
from glob import glob
import numpy as np
import json
import smplx as smpl
from nets import *
from repro_nets import *
from trainer.options import parse_args
from data_utils import torch_data
from trainer.config import load_JsonConfig
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils import data
def init_model(model_name, model_path, args, config):
if model_name == 'freeMo':
# generator = freeMo_Generator(args)
# generator = freeMo_Generator(args)
generator = freeMo_dev(args, config)
# generator.load_state_dict(torch.load(model_path)['generator'])
elif model_name == 'smplx_S2G':
generator = smplx_S2G(args, config)
elif model_name == 'StyleGestures':
generator = StyleGesture_Generator(
args,
config
)
elif model_name == 'Audio2Gestures':
config.Train.using_mspec_stat = False
generator = Audio2Gesture_Generator(
args,
config,
torch.zeros([1, 1, 108]),
torch.ones([1, 1, 108])
)
elif model_name == 'S2G':
generator = S2G_Generator(
args,
config,
)
elif model_name == 'Tmpt':
generator = S2G_Generator(
args,
config,
)
else:
raise NotImplementedError
model_ckpt = torch.load(model_path, map_location=torch.device('cpu'))
if model_name == 'smplx_S2G':
generator.generator.load_state_dict(model_ckpt['generator']['generator'])
elif 'generator' in list(model_ckpt.keys()):
generator.load_state_dict(model_ckpt['generator'])
else:
model_ckpt = {'generator': model_ckpt}
generator.load_state_dict(model_ckpt)
return generator
def prevar_loader(data_root, speakers, args, config, model_path, device, generator):
path = model_path.split('ckpt')[0]
file = os.path.join(os.path.dirname(path), "pre_variable.npy")
data_base = torch_data(
data_root=data_root,
speakers=speakers,
split='pre',
limbscaling=False,
normalization=config.Data.pose.normalization,
norm_method=config.Data.pose.norm_method,
split_trans_zero=False,
num_pre_frames=config.Data.pose.pre_pose_length,
num_generate_length=config.Data.pose.generate_length,
num_frames=15,
aud_feat_win_size=config.Data.aud.aud_feat_win_size,
aud_feat_dim=config.Data.aud.aud_feat_dim,
feat_method=config.Data.aud.feat_method,
smplx=True,
audio_sr=22000,
convert_to_6d=config.Data.pose.convert_to_6d,
expression=config.Data.pose.expression
)
data_base.get_dataset()
pre_set = data_base.all_dataset
pre_loader = data.DataLoader(pre_set, batch_size=config.DataLoader.batch_size, shuffle=False, drop_last=True)
total_pose = []
with torch.no_grad():
for bat in pre_loader:
pose = bat['poses'].to(device).to(torch.float32)
expression = bat['expression'].to(device).to(torch.float32)
pose = pose.permute(0, 2, 1)
pose = torch.cat([pose[:, :15], pose[:, 15:30], pose[:, 30:45], pose[:, 45:60], pose[:, 60:]], dim=0)
expression = expression.permute(0, 2, 1)
expression = torch.cat([expression[:, :15], expression[:, 15:30], expression[:, 30:45], expression[:, 45:60], expression[:, 60:]], dim=0)
pose = torch.cat([pose, expression], dim=-1)
pose = pose.reshape(pose.shape[0], -1, 1)
pose_code = generator.generator.pre_pose_encoder(pose).squeeze().detach().cpu()
total_pose.append(np.asarray(pose_code))
total_pose = np.concatenate(total_pose, axis=0)
mean = np.mean(total_pose, axis=0)
std = np.std(total_pose, axis=0)
prevar = (mean, std)
np.save(file, prevar, allow_pickle=True)
return mean, std
def main():
parser = parse_args()
args = parser.parse_args()
device = torch.device(args.gpu)
torch.cuda.set_device(device)
config = load_JsonConfig(args.config_file)
print('init model...')
generator = init_model(config.Model.model_name, args.model_path, args, config)
print('init pre-pose vectors...')
mean, std = prevar_loader(config.Data.data_root, args.speakers, args, config, args.model_path, device, generator)
main() |