File size: 1,636 Bytes
45950ff | 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 | import joblib
import random
import numpy as np
from tqdm import tqdm
data_root = 'data/final_data'
all_data = []
all_g1_motion = []
all_smplx_motion = []
with open(f'{data_root}/all.txt', 'r') as f:
data_paths = f.readlines()
for data_path in tqdm(data_paths):
data = joblib.load(data_path.strip())
all_data.append(data)
all_g1_motion.append(data['g1_motion'])
all_smplx_motion.append(data['smplx_motion'])
g1_data = np.concatenate(all_g1_motion, axis=0)
Mean = g1_data.mean(axis=0)
Std = g1_data.std(axis=0)
joints_num = 30
Std[:2] = Std[:2].mean() / 1.0
Std[2:8] = Std[2:8].mean() / 1.0
Std[8:8+3*joints_num] = Std[8:8+3*joints_num].mean() / 1.0
Std[8+3*joints_num:8+6*joints_num] = Std[8+3*joints_num:8+6*joints_num].mean() / 1.0
Std[8+6*joints_num:] = Std[8+6*joints_num:].mean() / 1.0
np.save(f'{data_root}/g1_mean.npy', Mean)
np.save(f'{data_root}/g1_std.npy', Std)
smplx_data = np.concatenate(all_smplx_motion, axis=0)
Mean = smplx_data.mean(axis=0)
Std = smplx_data.std(axis=0)
joints_num = 22
Std[:2] = Std[:2].mean() / 1.0
Std[2:8] = Std[2:8].mean() / 1.0
Std[8:8+3*joints_num] = Std[8:8+3*joints_num].mean() / 1.0
Std[8+3*joints_num:8+6*joints_num] = Std[8+3*joints_num:8+6*joints_num].mean() / 1.0
Std[8+6*joints_num:] = Std[8+6*joints_num:].mean() / 1.0
np.save(f'{data_root}/smplx_mean.npy', Mean)
np.save(f'{data_root}/smplx_std.npy', Std)
random.shuffle(all_data)
random.shuffle(all_data)
random.shuffle(all_data)
random.shuffle(all_data)
random.shuffle(all_data)
joblib.dump(all_data[:12500], f'{data_root}/train.pkl')
joblib.dump(all_data[12500:], f'{data_root}/test.pkl') |