lmd/Pose_to_SMPL_an_230402
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update code

Browse Source
This commit is contained in:
Iridoudou
2021-08-05 16:22:06 +08:00
parent 791f02f280
commit bc0766fc76
11 changed files with 47 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
fit/tools/main.py
View File

@ -26,7 +26,7 @@ from display_utils import display_model
from smplpytorch.pytorch.smpl_layer import SMPL_Layer
from train import train
from transform import transform
from save import save_pic
from save import save_pic,save_params
torch.backends.cudnn.benchmark=True
def parse_args():
@ -104,7 +104,7 @@ if __name__ == "__main__":
logger.info('Processing file: {}'.format(file))
target_path=os.path.join(root,file)
target = np.array(transform(np.load(cfg.TARGET_PATH)))
target = np.array(transform(np.load(target_path)))
logger.info('File shape: {}'.format(target.shape))
target = torch.from_numpy(target).float()
@ -112,5 +112,6 @@ if __name__ == "__main__":
logger,writer,device,
args,cfg)
# save_pic(target,res,smpl_layer,file)
# save_pic(target,res,smpl_layer,file,logger)
save_params(res,file,logger)

32
fit/tools/save.py
View File

@ -1,6 +1,10 @@
import sys
import os
import re
from tqdm import tqdm
import numpy as np
import json
sys.path.append(os.getcwd())
from display_utils import display_model
@ -9,14 +13,15 @@ def create_dir_not_exist(path):
if not os.path.exists(path):
os.mkdir(path)
def save_pic(target, res, smpl_layer, file):
def save_pic(target, res, smpl_layer, file,logger):
pose_params, shape_params, verts, Jtr = res
name=re.split('[/.]',file)[-2]
gt_path="fit/output/HumanAct12/picture/gt/{}".format(name)
fit_path="fit/output/HumanAct12/picture/fit/{}".format(name)
create_dir_not_exist(gt_path)
create_dir_not_exist(fit_path)
for i in range(target.shape[0]):
logger.info('Saving pictures at {} and {}'.format(gt_path,fit_path))
for i in tqdm(range(target.shape[0])):
display_model(
{'verts': verts.cpu().detach(),
'joints': target.cpu().detach()},
@ -36,3 +41,26 @@ def save_pic(target, res, smpl_layer, file):
savepath=os.path.join(fit_path+"/frame_{}".format(i)),
batch_idx=i,
show=False)
logger.info('Pictures saved')
def save_params(res,file,logger):
pose_params, shape_params, verts, Jtr = res
name=re.split('[/.]',file)[-2]
fit_path="fit/output/HumanAct12/params/"
create_dir_not_exist(fit_path)
logger.info('Saving params at {}'.format(fit_path))
pose_params=pose_params.cpu().detach()
pose_params=pose_params.numpy().tolist()
shape_params=shape_params.cpu().detach()
shape_params=shape_params.numpy().tolist()
Jtr=Jtr.cpu().detach()
Jtr=Jtr.numpy().tolist()
params={}
params["pose_params"]=pose_params
params["shape_params"]=shape_params
params["Jtr"]=Jtr
f=open(os.path.join((fit_path),
"{}_params.json".format(name)),'w')
json.dump(params,f)
logger.info('Params saved')

9
fit/tools/train.py
View File

@ -30,7 +30,7 @@ def train(smpl_layer, target,
args, cfg):
res = []
pose_params = torch.rand(target.shape[0], 72) * 0.0
shape_params = torch.rand(target.shape[0], 10) * 0.1
shape_params = torch.rand(target.shape[0], 10) * 0.03
scale = torch.ones([1])
smpl_layer = smpl_layer.to(device)
@ -41,9 +41,10 @@ def train(smpl_layer, target,
pose_params.requires_grad = True
shape_params.requires_grad = True
scale.requires_grad = True
scale.requires_grad = False
smpl_layer.requires_grad = False
optimizer = optim.Adam([pose_params],
optimizer = optim.Adam([pose_params, shape_params],
lr=cfg.TRAIN.LEARNING_RATE)
min_loss = float('inf')
@ -62,5 +63,5 @@ def train(smpl_layer, target,
writer.add_scalar('loss', float(loss), epoch)
writer.add_scalar('learning_rate', float(
optimizer.state_dict()['param_groups'][0]['lr']), epoch)
logger.info('Train ended, loss = {:.9f}'.format(float(loss)))
logger.info('Train ended, min_loss = {:.9f}'.format(float(min_loss)))
return res