import argparse from sata.utils import tensor_utils import numpy as np import torch from sata.mypath import * from sata.mymodel import make_load_model from sata.mydataset import PairedDataset, get_mi_src_tgt_all_graph from sata.skel_pose_graph import SkelPoseGraph, rnd_mask from sata.utils.skel_gen_utils import create_random_skel from sata.conversions.graph_to_motion import graph_2_skel from fairmotion.core import motion as motion_class def prepare_model_test(model_epoch, device): # device, printoptions tensor_utils.set_device(device) np.set_printoptions(precision=5, suppress=True) torch.set_printoptions(precision=5, sci_mode=False) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False # Model model, cfg = make_load_model(model_epoch, device) model.eval() load_dir = os.path.join(RESULT_DIR, model_epoch.split("/")[0]) ms_dict = torch.load(os.path.join(load_dir, "ms_dict.pt")) # set SkelPoseGraph class variables SkelPoseGraph.skel_cfg = cfg["representation"]["skel"] SkelPoseGraph.pose_cfg = cfg["representation"]["pose"] SkelPoseGraph.ms_dict = ms_dict return model, cfg, ms_dict """ ================= basic functions commonly needed for tasks ================= """ from sata.conversions.graph_to_motion import gt_recon_motion, hatD_recon_motion from sata.conversions.motion_to_graph import bvh_2_graph, skel_2_graph from torch_geometric.data import Batch def retarget(model, src_batch, tgt_batch, ms_dict, out_rep_cfg, consq_n): # src ground truth src_motion_list, src_contact_list = gt_recon_motion(src_batch, consq_n) # predicted result z, hatD = model(src_batch, tgt_batch) out_motion_list, out_contact_list = hatD_recon_motion( hatD, tgt_batch, out_rep_cfg, ms_dict, consq_n ) # when tgt ground-truth motion is available if hasattr(tgt_batch, "q"): tgt_motion_list, tgt_contact_list = gt_recon_motion(tgt_batch, consq_n) return src_motion_list[0], tgt_motion_list[0], out_motion_list[0] else: tgt_skel = graph_2_skel(tgt_batch, 1)[0] tgt_motion = motion_class.Motion(skel=tgt_skel) tpose = np.eye(4)[None, ...].repeat(tgt_skel.num_joints(), 0) tpose[0, 1, 3] = tgt_batch.go[0, 1] # root height tgt_motion.add_one_frame(tpose) return src_motion_list[0], tgt_motion, out_motion_list[0] ##### motion to z ##### def bvh_2_graph_z(model, bvh_filepath): graph_batch = bvh_2_graph(bvh_filepath).to(device=model.device) z = model.encoder[0](graph_batch) motion_list, contact_list = gt_recon_motion(graph_batch, len(z)) return motion_list[0], graph_batch, z ##### z to motion ##### def decode_z_skel(model, z, skel, ms_dict): return decode_z_skelgraph(model, z, skel_2_graph(skel), ms_dict) def decode_z_skelgraph(model, z, skel_graph, ms_dict): B_skel_graph = Batch.from_data_list([skel_graph] * len(z)).to(device=model.device) hatD = model.decoder[0](z, B_skel_graph) out_motion_list, out_contact_list = hatD_recon_motion( hatD, B_skel_graph, model.rep_cfg["out"], ms_dict, len(z) ) return out_motion_list[0], out_contact_list[0] ##### convert all bvh to z and save as npy ##### def list_bvh_files(directory): bvh_files = [] for root, dirs, files in os.walk(directory): if not dirs: # leaf directory relative_path = os.path.relpath(root, directory) for file in files: if file.endswith(".bvh"): bvh_files.append(os.path.join(relative_path, file)) return bvh_files import tqdm, gc def save_bvh_z(model_epoch, bvh_dir, npy_dir): model, cfg, ms_dict = prepare_model_test(model_epoch, "cuda:0") bvh_files = list_bvh_files(bvh_dir) for bvh_rel_fn in tqdm.tqdm(bvh_files): bvh_fp = os.path.join(bvh_dir, bvh_rel_fn) npy_fp = os.path.join(npy_dir, bvh_rel_fn[:-4] + ".npy") if not os.path.exists(os.path.dirname(npy_fp)): os.makedirs(os.path.dirname(npy_fp)) motion, graph_batch, z = bvh_2_graph_z(model, bvh_fp) # print(bvh_fp, npy_fp) np.save(npy_fp, z.cpu().detach().numpy()) del motion, graph_batch, z # gc.collect() # torch.cuda.empty_cache() if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--model_epoch", type=str, default="ckpt0") parser.add_argument("--device", type=str, default="cuda:0") parser.add_argument("--data_dir", type=str, default="test/motion/processed/") parser.add_argument("--rnd_tgt", type=int, default=0) parser.add_argument("--src_mask", type=int, default=0) parser.add_argument("--tgt_mask", type=int, default=0) args = parser.parse_args() model, cfg, ms_dict = prepare_model_test(args.model_epoch, args.device) # Dataset ds = PairedDataset() data_dir = os.path.join(DATA_DIR, args.data_dir) ds.load_data_dir_pairs(data_dir) from sata.default_veiwer import get_default_viewer viewer = get_default_viewer(argparse.Namespace(imgui=False)) def retarget_mi(mi): R = len(ds.mi_ri_2_fi[mi]) src_ri, tgt_ri = np.random.randint(0, R, size=2) (src_batch, tgt_batch), consq_n = get_mi_src_tgt_all_graph( dataset=ds, mi=mi, src_ri=src_ri, tgt_ri=tgt_ri, device=args.device ) # option 1) test with data tgt skeleton # option 2) test with random skeleton if args.rnd_tgt: rnd_tgt_skel = create_random_skel() rnd_tgt_skel_graph = skel_2_graph(rnd_tgt_skel) rnd_tgt_batch = Batch.from_data_list([rnd_tgt_skel_graph] * consq_n).to( device=model.device ) tgt_batch = rnd_tgt_batch if args.src_mask: src_batch.mask = rnd_mask(src_batch, consq_n=consq_n) if args.tgt_mask: tgt_batch.mask = rnd_mask(tgt_batch, consq_n=consq_n) src_motion, tgt_motion, out_motion = retarget( model, src_batch, tgt_batch, ms_dict, out_rep_cfg=cfg["representation"]["out"], consq_n=consq_n, ) # update viewer viewer.update_motions([src_motion, tgt_motion, out_motion], 150, linear=True) viewer.mi = mi retarget_mi(0) def extra_key_callback(key): if key == b"m": next_mi = (viewer.mi + 1) % len(ds.mi_ri_2_fi) retarget_mi(next_mi) return True return False viewer.extra_key_callback = extra_key_callback viewer.run()