File size: 4,842 Bytes
823807d |
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 |
import sys
import os
from os.path import join as pjoin
import torch
from models.vq.model import RVQVAE
from options.vq_option import arg_parse
from motion_loaders.dataset_motion_loader import get_dataset_motion_loader
import utils.eval_t2m as eval_t2m
from utils.get_opt import get_opt
from models.t2m_eval_wrapper import EvaluatorModelWrapper
import warnings
warnings.filterwarnings('ignore')
import numpy as np
from utils.word_vectorizer import WordVectorizer
def load_vq_model(vq_opt, which_epoch):
# opt_path = pjoin(opt.checkpoints_dir, opt.dataset_name, opt.vq_name, 'opt.txt')
vq_model = RVQVAE(vq_opt,
dim_pose,
vq_opt.nb_code,
vq_opt.code_dim,
vq_opt.code_dim,
vq_opt.down_t,
vq_opt.stride_t,
vq_opt.width,
vq_opt.depth,
vq_opt.dilation_growth_rate,
vq_opt.vq_act,
vq_opt.vq_norm)
ckpt = torch.load(pjoin(vq_opt.checkpoints_dir, vq_opt.dataset_name, vq_opt.name, 'model', which_epoch),
map_location='cpu')
model_key = 'vq_model' if 'vq_model' in ckpt else 'net'
vq_model.load_state_dict(ckpt[model_key])
vq_epoch = ckpt['ep'] if 'ep' in ckpt else -1
print(f'Loading VQ Model {vq_opt.name} Completed!, Epoch {vq_epoch}')
return vq_model, vq_epoch
if __name__ == "__main__":
##### ---- Exp dirs ---- #####
args = arg_parse(False)
args.device = torch.device("cpu" if args.gpu_id == -1 else "cuda:" + str(args.gpu_id))
args.out_dir = pjoin(args.checkpoints_dir, args.dataset_name, args.name, 'eval')
os.makedirs(args.out_dir, exist_ok=True)
f = open(pjoin(args.out_dir, '%s.log'%args.ext), 'w')
dataset_opt_path = 'checkpoints/kit/Comp_v6_KLD005/opt.txt' if args.dataset_name == 'kit' \
else 'checkpoints/t2m/Comp_v6_KLD005/opt.txt'
wrapper_opt = get_opt(dataset_opt_path, torch.device('cuda'))
eval_wrapper = EvaluatorModelWrapper(wrapper_opt)
##### ---- Dataloader ---- #####
args.nb_joints = 21 if args.dataset_name == 'kit' else 22
dim_pose = 251 if args.dataset_name == 'kit' else 263
eval_val_loader, _ = get_dataset_motion_loader(dataset_opt_path, 32, 'test', device=args.device)
print(len(eval_val_loader))
##### ---- Network ---- #####
vq_opt_path = pjoin(args.checkpoints_dir, args.dataset_name, args.name, 'opt.txt')
vq_opt = get_opt(vq_opt_path, device=args.device)
# net = load_vq_model()
model_dir = pjoin(args.checkpoints_dir, args.dataset_name, args.name, 'model')
for file in os.listdir(model_dir):
# if not file.endswith('tar'):
# continue
# if not file.startswith('net_best_fid'):
# continue
if args.which_epoch != "all" and args.which_epoch not in file:
continue
print(file)
net, ep = load_vq_model(vq_opt, file)
net.eval()
net.cuda()
fid = []
div = []
top1 = []
top2 = []
top3 = []
matching = []
mae = []
repeat_time = 20
for i in range(repeat_time):
best_fid, best_div, Rprecision, best_matching, l1_dist = \
eval_t2m.evaluation_vqvae_plus_mpjpe(eval_val_loader, net, i, eval_wrapper=eval_wrapper, num_joint=args.nb_joints)
fid.append(best_fid)
div.append(best_div)
top1.append(Rprecision[0])
top2.append(Rprecision[1])
top3.append(Rprecision[2])
matching.append(best_matching)
mae.append(l1_dist)
fid = np.array(fid)
div = np.array(div)
top1 = np.array(top1)
top2 = np.array(top2)
top3 = np.array(top3)
matching = np.array(matching)
mae = np.array(mae)
print(f'{file} final result, epoch {ep}')
print(f'{file} final result, epoch {ep}', file=f, flush=True)
msg_final = f"\tFID: {np.mean(fid):.3f}, conf. {np.std(fid)*1.96/np.sqrt(repeat_time):.3f}\n" \
f"\tDiversity: {np.mean(div):.3f}, conf. {np.std(div)*1.96/np.sqrt(repeat_time):.3f}\n" \
f"\tTOP1: {np.mean(top1):.3f}, conf. {np.std(top1)*1.96/np.sqrt(repeat_time):.3f}, TOP2. {np.mean(top2):.3f}, conf. {np.std(top2)*1.96/np.sqrt(repeat_time):.3f}, TOP3. {np.mean(top3):.3f}, conf. {np.std(top3)*1.96/np.sqrt(repeat_time):.3f}\n" \
f"\tMatching: {np.mean(matching):.3f}, conf. {np.std(matching)*1.96/np.sqrt(repeat_time):.3f}\n" \
f"\tMAE:{np.mean(mae):.3f}, conf.{np.std(mae)*1.96/np.sqrt(repeat_time):.3f}\n\n"
# logger.info(msg_final)
print(msg_final)
print(msg_final, file=f, flush=True)
f.close()
|