support NTU

This commit is contained in:
Iridoudou
2021-09-03 11:06:26 +08:00
parent 0702081dbe
commit 16387a7afe
8 changed files with 236 additions and 23 deletions

View File

@ -1,3 +1,5 @@
from display_utils import display_model
from label import get_label
import sys
import os
import re
@ -6,8 +8,6 @@ import numpy as np
import pickle
sys.path.append(os.getcwd())
from label import get_label
from display_utils import display_model
def create_dir_not_exist(path):
@ -18,10 +18,8 @@ def create_dir_not_exist(path):
def save_pic(res, smpl_layer, file, logger, dataset_name, target):
_, _, verts, Jtr = res
file_name = re.split('[/.]', file)[-2]
fit_path = "fit/output/{}/picture/fit/{}".format(dataset_name, file_name)
# gt_path = "fit/output/{}/picture/gt/{}".format(dataset_name, file_name)
fit_path = "fit/output/{}/picture/{}".format(dataset_name, file_name)
create_dir_not_exist(fit_path)
# create_dir_not_exist(gt_path)
logger.info('Saving pictures at {}'.format(fit_path))
for i in tqdm(range(Jtr.shape[0])):
display_model(
@ -32,18 +30,8 @@ def save_pic(res, smpl_layer, file, logger, dataset_name, target):
kintree_table=smpl_layer.kintree_table,
savepath=os.path.join(fit_path+"/frame_{}".format(i)),
batch_idx=i,
show=True,
show=False,
only_joint=True)
# display_model(
# {'verts': verts.cpu().detach(),
# 'joints': target.cpu().detach()},
# model_faces=smpl_layer.th_faces,
# with_joints=True,
# kintree_table=smpl_layer.kintree_table,
# savepath=os.path.join(gt_path+"/frame_{}".format(i)),
# batch_idx=i,
# show=False,
# only_joint=True)
logger.info('Pictures saved')
@ -63,6 +51,7 @@ def save_params(res, file, logger, dataset_name):
params["pose_params"] = pose_params
params["shape_params"] = shape_params
params["Jtr"] = Jtr
print("label:{}".format(label))
with open(os.path.join((fit_path),
"{}_params.pkl".format(file_name)), 'wb') as f:
pickle.dump(params, f)