Spaces:
Runtime error
Runtime error
| # 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) | |