Spaces:
Runtime error
Runtime error
File size: 2,563 Bytes
c3d0293 |
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 |
# merges the output of the main transfer_model script
import torch
from pathlib import Path
import pickle
from scipy.spatial.transform import Rotation as R
import numpy as np
KEYS = [
"transl",
"betas",
"full_pose",
]
def aggregate_rotmats(x):
x = np.concatenate(x, axis=0)
s = x.shape[:-2]
try:
x = R.from_matrix(x.reshape(-1, 3, 3)).as_rotvec()
except:
pass
x = x.reshape(s[0], -1)
return x
aggregate_function = {k: lambda x: np.concatenate(x, axis=0) for k in KEYS}
aggregate_function["betas"] = lambda x: np.concatenate(x, axis=0).mean(0)
for k in ["global_orient", "body_pose", "left_hand_pose", "right_hand_pose", "jaw_pose", "full_pose"]:
aggregate_function[k] = aggregate_rotmats
def merge(output_dir, gender):
output_dir = Path(output_dir)
assert output_dir.exists()
assert output_dir.is_dir()
# get list of all pkl files in output_dir with fixed length numeral names
pkl_files = [f for f in output_dir.glob("*.pkl") if f.stem != "merged"]
pkl_files = [f for f in sorted(pkl_files, key=lambda x: int(x.stem))]
assert "merged.pkl" not in [f.name for f in pkl_files]
merged = {}
# iterate over keys and put all values in lists
keys = set(KEYS)
for k in keys:
merged[k] = []
for pkl_file in pkl_files:
with open(pkl_file, "rb") as f:
data = pickle.load(f)
for k in keys:
if k in data:
merged[k].append(data[k])
b = np.concatenate(merged["betas"], axis=0)
print("betas:")
for mu, sigma in zip(b.mean(0), b.std(0)):
print(" {:.3f} +/- {:.3f}".format(mu, sigma))
# aggregate all values
for k in keys:
merged[k] = aggregate_function[k](merged[k])
# add gender
poses = merged["full_pose"]
trans = merged["transl"]
if gender == "female":
gender = np.zeros([poses.shape[0], 1])
elif gender == "male":
gender = np.ones([poses.shape[0], 1])
else:
gender = np.ones([poses.shape[0], 1]) * 2
merged = np.concatenate([poses, trans, gender], axis=1)
return merged
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser(description='Merge output of transfer_model script')
parser.add_argument('output_dir', type=str, help='output directory of transfer_model script')
parser.add_argument('--gender', type=str, choices=['male', 'female', 'neutral'], help='gender of actor in motion sequence')
args = parser.parse_args()
merge(args.output_dir, args.gender)
|