def _load_pretrained_model(self):
print("Loading pretrained posenet...")
#ld = "/home/ubuntu/project/Pose2Mesh_RELEASE/experiment/posenet_human36J_train_human36"
checkpoint = load_checkpoint(load_dir=cfg.MODEL.posenet_path,
pick_best=True)
#checkpoint = load_checkpoint(load_dir=ld, pick_best=True)
self.load_state_dict(checkpoint['model_state_dict'])
def prepare_network(args, load_dir='', is_train=True):
dataset_names = cfg.DATASET.train_list if is_train else cfg.DATASET.test_list
dataset_list, dataloader = get_dataloader(args, dataset_names, is_train)
model, criterion, optimizer, lr_scheduler = None, None, None, None
loss_history, test_error_history = [], {'surface': [], 'joint': []}
main_dataset = dataset_list[0]
if is_train or load_dir:
print(f"==> Preparing {cfg.MODEL.name} MODEL...")
if cfg.MODEL.name == 'pose2mesh_net':
model = models.pose2mesh_net.get_model(
num_joint=main_dataset.joint_num, graph_L=main_dataset.graph_L)
elif cfg.MODEL.name == 'posenet':
model = models.posenet.get_model(main_dataset.joint_num,
hid_dim=4096,
num_layer=2,
p_dropout=0.5)
print('# of model parameters: {}'.format(count_parameters(model)))
if is_train:
criterion = get_loss(faces=main_dataset.mesh_model.face)
optimizer = get_optimizer(model=model)
lr_scheduler = get_scheduler(optimizer=optimizer)
if load_dir and (not is_train or args.resume_training):
print('==> Loading checkpoint')
checkpoint = load_checkpoint(load_dir=load_dir,
pick_best=(cfg.MODEL.name == 'posenet'))
model.load_state_dict(checkpoint['model_state_dict'])
if is_train:
optimizer.load_state_dict(checkpoint['optim_state_dict'])
for state in optimizer.state.values():
for k, v in state.items():
if torch.is_tensor(v):
state[k] = v.cuda()
curr_lr = 0.0
for param_group in optimizer.param_groups:
curr_lr = param_group['lr']
lr_state = checkpoint['scheduler_state_dict']
# update lr_scheduler
lr_state['milestones'], lr_state['gamma'] = Counter(
cfg.TRAIN.lr_step), cfg.TRAIN.lr_factor
lr_scheduler.load_state_dict(lr_state)
loss_history = checkpoint['train_log']
test_error_history = checkpoint['test_log']
cfg.TRAIN.begin_epoch = checkpoint['epoch'] + 1
print(
'===> resume from epoch {:d}, current lr: {:.0e}, milestones: {}, lr factor: {:.0e}'
.format(cfg.TRAIN.begin_epoch, curr_lr, lr_state['milestones'],
lr_state['gamma']))
return dataloader, dataset_list, model, criterion, optimizer, lr_scheduler, loss_history, test_error_history
def get_joint_setting(mesh_model, joint_category='human36'):
joint_regressor, joint_num, skeleton, graph_L, graph_perm_reverse = None, None, None, None, None
if joint_category == 'human36':
joint_regressor = mesh_model.joint_regressor_h36m
joint_num = 17
skeleton = (
(0, 7), (7, 8), (8, 9), (9, 10), (8, 11), (11, 12), (12, 13), (8, 14), (14, 15), (15, 16), (0, 1), (1, 2),
(2, 3), (0, 4), (4, 5), (5, 6))
flip_pairs = ((1, 4), (2, 5), (3, 6), (14, 11), (15, 12), (16, 13))
graph_Adj, graph_L, graph_perm,graph_perm_reverse = \
build_coarse_graphs(mesh_model.face, joint_num, skeleton, flip_pairs, levels=9)
model_chk_path = './experiment/pose2mesh_human36J_train_human36/final.pth.tar'
elif joint_category == 'coco':
joint_regressor = mesh_model.joint_regressor_coco
joint_num = 19 # add pelvis and neck
skeleton = (
(1, 2), (0, 1), (0, 2), (2, 4), (1, 3), (6, 8), (8, 10), (5, 7), (7, 9), (12, 14), (14, 16), (11, 13),
(13, 15), # (5, 6), #(11, 12),
(17, 11), (17, 12), (17, 18), (18, 5), (18, 6), (18, 0))
flip_pairs = ((1, 2), (3, 4), (5, 6), (7, 8), (9, 10), (11, 12), (13, 14), (15, 16))
graph_Adj, graph_L, graph_perm, graph_perm_reverse = \
build_coarse_graphs(mesh_model.face, joint_num, skeleton, flip_pairs, levels=9)
model_chk_path = './experiment/pose2mesh_cocoJ_train_human36_coco_muco/final.pth.tar'
elif joint_category == 'smpl':
joint_regressor = mesh_model.layer['neutral'].th_J_regressor.numpy().astype(np.float32)
joint_num = 24
skeleton = (
(0, 1), (1, 4), (4, 7), (7, 10), (0, 2), (2, 5), (5, 8), (8, 11), (0, 3), (3, 6), (6, 9), (9, 14), (14, 17),
(17, 19), (19, 21), (21, 23), (9, 13), (13, 16), (16, 18), (18, 20), (20, 22), (9, 12), (12, 15))
flip_pairs = ((1, 2), (4, 5), (7, 8), (10, 11), (13, 14), (16, 17), (18, 19), (20, 21), (22, 23))
graph_Adj, graph_L, graph_perm, graph_perm_reverse = \
build_coarse_graphs(mesh_model.face, joint_num, skeleton, flip_pairs, levels=9)
model_chk_path = './experiment/pose2mesh_smplJ_train_surreal/final.pth.tar'
elif joint_category == 'mano':
joint_regressor = mesh_model.joint_regressor
joint_num = 21
skeleton = ( (0,1), (0,5), (0,9), (0,13), (0,17), (1,2), (2,3), (3,4), (5,6), (6,7), (7,8), (9,10), (10,11), (11,12), (13,14), (14,15), (15,16), (17,18), (18,19), (19,20) )
hori_conn = (
(1, 5), (5, 9), (9, 13), (13, 17), (2, 6), (6, 10), (10, 14), (14, 18), (3, 7), (7, 11), (11, 15), (15, 19),
(4, 8), (8, 12), (12, 16), (16, 20))
graph_Adj, graph_L, graph_perm, graph_perm_reverse = \
build_coarse_graphs(mesh_model.face, joint_num, skeleton, hori_conn, levels=6)
model_chk_path = './experiment/pose2mesh_manoJ_train_freihand/final.pth.tar'
else:
raise NotImplementedError(f"{joint_category}: unknown joint set category")
model = models.pose2mesh_net.get_model(joint_num, graph_L)
checkpoint = load_checkpoint(load_dir=model_chk_path)
model.load_state_dict(checkpoint['model_state_dict'])
return model, joint_regressor, joint_num, skeleton, graph_L, graph_perm_reverse
def _load_pretrained_model(self):
print("Loading pretrained posenet...")
checkpoint = load_checkpoint(load_dir=cfg.MODEL.posenet_path,
pick_best=True)
self.load_state_dict(checkpoint['model_state_dict'])