def validate(m, opt, heatmap_to_coord, batch_size=64):
det_dataset = builder.build_dataset(cfg.DATASET.TEST,
preset_cfg=cfg.DATA_PRESET,
train=False,
opt=opt)
det_loader = torch.utils.data.DataLoader(det_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=opt.nThreads,
drop_last=False)
kpt_json = []
eval_joints = det_dataset.EVAL_JOINTS
m.eval()
norm_type = cfg.LOSS.get('NORM_TYPE', None)
hm_size = cfg.DATA_PRESET.HEATMAP_SIZE
for index, (inps, crop_bboxes, bboxes, img_ids, scores, imghts,
imgwds) in tqdm(enumerate(det_loader), dynamic_ncols=True):
if opt.device.type != "cpu":
if isinstance(inps, list):
inps = [inp.cuda() for inp in inps]
else:
inps = inps.cuda()
full_output = m(inps)
joints_map = full_output.joints_map
pred = joints_map
assert pred.dim() == 4
pred = pred[:, eval_joints, :, :]
for i in range(joints_map.shape[0]):
bbox = crop_bboxes[i].tolist()
pose_coords, pose_scores = heatmap_to_coord(
pred[i][det_dataset.EVAL_JOINTS],
bbox,
hm_shape=hm_size,
norm_type=norm_type)
keypoints = np.concatenate((pose_coords, pose_scores), axis=1)
keypoints = keypoints.reshape(-1).tolist()
data = dict()
data['bbox'] = bboxes[i, 0].tolist()
data['image_id'] = int(img_ids[i])
data['score'] = float(scores[i] + np.mean(pose_scores) +
np.max(pose_scores))
data['category_id'] = 1
data['keypoints'] = keypoints
kpt_json.append(data)
with open(os.path.join(opt.work_dir, 'test_kpt.json'), 'w') as fid:
json.dump(kpt_json, fid)
res = evaluate_mAP(os.path.join(opt.work_dir, 'test_kpt.json'),
ann_type='keypoints',
ann_file=os.path.join(cfg.DATASET.VAL.ROOT,
cfg.DATASET.VAL.ANN))
return res
def main():
### Get Dataset
with torch.no_grad():
## build dataset
test_dataset = builder.build_dataset(cfg.DATASET.TEST,
preset_cfg=cfg.DATA_PRESET,
train=False)
eval_joints = test_dataset.EVAL_JOINTS
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=batchSize,
shuffle=False,
num_workers=20,
drop_last=False)
kpt_json = []
for inps, labels, label_masks, img_ids, bboxes in tqdm(
test_loader, dynamic_ncols=True):
start_time = getTime()
if isinstance(inps, list):
inps = [inp.cuda for inp in inps]
else:
inps = inps.cuda()
output = pose_model(inps)
pred = output.cpu().data.numpy()
assert pred.ndim == 4
pred = pred[:, eval_joints, :, :]
for i in range(output.shape[0]):
bbox = bboxes[i].tolist()
pose_coords, pose_scores = heatmap_to_coord(
pred[i][test_dataset.EVAL_JOINTS], bbox)
## concatenate - pose coordinates and pose scores
keypoints = np.concatenate((pose_coords, pose_scores), axis=1)
## keypoints
keypoints = keypoints.reshape(-1).tolist()
data = dict()
data['bbox'] = bboxes[i].tolist()
data['image_id'] = img_ids[i]
data['score'] = float(
np.mean(pose_scores) + np.max(pose_scores))
data['category_id'] = 22
data['keypoints'] = keypoints
kpt_json.append(data)
with open(os.path.join(args.outputpath, 'test_gt_kpt.json'), 'w') as fid:
json.dump(kpt_json, fid)
## Visualization
images, annotations = load_annotations("test_annot_keypoint.pkl")
## visualize
cmap = plt.cm.get_cmap("hsv", 1)
visualize(cfg.DATASET.TEST, args.outputpath, kpt_json, cmap, annotations)
def validate_gt(pose_model, cfg, heatmap_to_coord, batch_size=20):
gt_val_dataset = builder.build_dataset(cfg.DATASET.VAL, preset_cfg=cfg.DATA_PRESET, train=False)
eval_joints = gt_val_dataset.EVAL_JOINTS
gt_val_loader = torch.utils.data.DataLoader(
gt_val_dataset, batch_size=batch_size, shuffle=False, num_workers=20, drop_last=False)
kpt_json = []
# m.eval()
norm_type = cfg.LOSS.get('NORM_TYPE', None)
hm_size = cfg.DATA_PRESET.HEATMAP_SIZE
for inps, labels, label_masks, img_ids, bboxes in tqdm(gt_val_loader, dynamic_ncols=True):
if isinstance(inps, list):
inps = [inp.cuda() for inp in inps]
else:
inps = inps.cuda()
# 使用static shape的fast pose的engine进行推理
output = pose_model.detect_context(inps.cpu().numpy())
# 一定要转到cuda上，结果才不会出错
output = torch.from_numpy(output).to(opt.device)
if opt.flip_test:
if isinstance(inps, list):
inps_flip = [flip(inp).cuda() for inp in inps]
else:
inps_flip = flip(inps).cuda()
inps_flip = pose_model.detect_context(inps_flip.cpu().numpy())
inps_flip = torch.from_numpy(inps_flip).to(opt.device)
# output_flip = flip_heatmap(m(inps_flip), gt_val_dataset.joint_pairs, shift=True)
output_flip = flip_heatmap(inps_flip, gt_val_dataset.joint_pairs, shift=True)
pred_flip = output_flip[:, eval_joints, :, :]
else:
output_flip = None
pred = output
assert pred.dim() == 4
pred = pred[:, eval_joints, :, :]
for i in range(bboxes.shape[0]):
bbox = bboxes[i].tolist()
pose_coords, pose_scores = heatmap_to_coord(
pred[i], bbox, hms_flip=pred_flip[i], hm_shape=hm_size, norm_type=norm_type)
keypoints = np.concatenate((pose_coords, pose_scores), axis=1)
keypoints = keypoints.reshape(-1).tolist()
data = dict()
data['bbox'] = bboxes[i].tolist()
data['image_id'] = int(img_ids[i])
data['score'] = float(np.mean(pose_scores) + np.max(pose_scores))
data['category_id'] = 1
data['keypoints'] = keypoints
kpt_json.append(data)
with open('./exp/json/trt/validate_gt_kpt.json', 'w') as fid:
json.dump(kpt_json, fid)
res = evaluate_mAP('./exp/json/trt/validate_gt_kpt.json', ann_type='keypoints',
ann_file=os.path.join(cfg.DATASET.VAL.ROOT, cfg.DATASET.VAL.ANN))
return res
def validate(m, heatmap_to_coord, batch_size=20):
det_dataset = builder.build_dataset(cfg.DATASET.TEST,
preset_cfg=cfg.DATA_PRESET,
train=False,
opt=opt)
eval_joints = det_dataset.EVAL_JOINTS
det_loader = torch.utils.data.DataLoader(det_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=20,
drop_last=False)
kpt_json = []
m.eval()
for inps, crop_bboxes, bboxes, img_ids, scores, imghts, imgwds in tqdm(
det_loader, dynamic_ncols=True):
if isinstance(inps, list):
inps = [inp.cuda() for inp in inps]
else:
inps = inps.cuda()
output = m(inps)
if opt.flip_test:
if isinstance(inps, list):
inps_flip = [flip(inp) for inp in inps]
else:
inps_flip = flip(inps)
output_flip = flip_heatmap(m(inps_flip),
det_dataset.joint_pairs,
shift=True)
output = (output + output_flip) / 2
pred = output.cpu().data.numpy()
assert pred.ndim == 4
pred = pred[:, eval_joints, :, :]
for i in range(output.shape[0]):
bbox = crop_bboxes[i].tolist()
pose_coords, pose_scores = heatmap_to_coord(
pred[i][det_dataset.EVAL_JOINTS], bbox)
keypoints = np.concatenate((pose_coords, pose_scores), axis=1)
keypoints = keypoints.reshape(-1).tolist()
data = dict()
data['bbox'] = bboxes[i, 0].tolist()
data['image_id'] = int(img_ids[i])
data['score'] = float(scores[i] + np.mean(pose_scores) +
np.max(pose_scores))
data['category_id'] = 1
data['keypoints'] = keypoints
kpt_json.append(data)
with open('./exp/json/validate_rcnn_kpt.json', 'w') as fid:
json.dump(kpt_json, fid)
res = evaluate_mAP('./exp/json/validate_rcnn_kpt.json',
ann_type='keypoints')
return res['AP']
def validate(m, heatmap_to_coord, batch_size=20):
det_dataset = builder.build_dataset(cfg.DATASET.TEST, preset_cfg=cfg.DATA_PRESET, detector_cfg=cfg.DETECTOR, train=False, opt=opt)
eval_joints = det_dataset.EVAL_JOINTS
det_loader = torch.utils.data.DataLoader(
det_dataset, batch_size=batch_size, shuffle=False, num_workers=20, drop_last=False)
kpt_json = []
m.eval()
norm_type = cfg.LOSS.get('NORM_TYPE', None)
hm_size = cfg.DATA_PRESET.HEATMAP_SIZE
for inps, crop_bboxes, bboxes, img_ids, scores, imghts, imgwds in tqdm(det_loader, dynamic_ncols=True):
if isinstance(inps, list):
inps = [inp.cuda() for inp in inps]
else:
inps = inps.cuda()
output = m(inps)
if opt.flip_test:
if isinstance(inps, list):
inps_flip = [flip(inp).cuda() for inp in inps]
else:
inps_flip = flip(inps).cuda()
output_flip = flip_heatmap(m(inps_flip), det_dataset.joint_pairs, shift=True)
pred_flip = output_flip[:, eval_joints, :, :]
else:
output_flip = None
pred = output
assert pred.dim() == 4
pred = pred[:, eval_joints, :, :]
for i in range(output.shape[0]):
bbox = crop_bboxes[i].tolist()
pose_coords, pose_scores = heatmap_to_coord(
pred[i], bbox, hms_flip=pred_flip[i], hm_shape=hm_size, norm_type=norm_type)
keypoints = np.concatenate((pose_coords, pose_scores), axis=1)
keypoints = keypoints.reshape(-1).tolist()
data = dict()
data['bbox'] = bboxes[i, 0].tolist()
data['image_id'] = int(img_ids[i])
data['area'] = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
# data['score'] = float(scores[i] + np.mean(pose_scores) + np.max(pose_scores))
data['score'] = float(scores[i])
data['category_id'] = 1
data['keypoints'] = keypoints
kpt_json.append(data)
kpt_json = oks_pose_nms(kpt_json)
with open('./exp/json/validate_rcnn_kpt.json', 'w') as fid:
json.dump(kpt_json, fid)
res = evaluate_mAP('./exp/json/validate_rcnn_kpt.json', ann_type='keypoints')
return res['AP']
def validate_gt(m, opt, cfg, heatmap_to_coord, criterion, batch_size=20):
loss_logger = DataLogger()
acc_logger = DataLogger()
gt_val_dataset = builder.build_dataset(cfg.DATASET.VAL, preset_cfg=cfg.DATA_PRESET, train=False)
eval_joints = gt_val_dataset.EVAL_JOINTS
gt_val_loader = torch.utils.data.DataLoader(
gt_val_dataset, batch_size=batch_size, shuffle=False, num_workers=20, drop_last=False)
kpt_json = []
m.eval()
for inps, labels, label_masks, img_ids, bboxes in tqdm(gt_val_loader, dynamic_ncols=True):
if isinstance(inps, list):
inps = [inp.cuda() for inp in inps]
else:
inps = inps.cuda()
output = m(inps)
##
labels = labels.cuda()
label_masks = label_masks.cuda()
loss = 0.5 * criterion(output.mul(label_masks), labels.mul(label_masks))
acc = calc_accuracy(output.mul(label_masks), labels.mul(label_masks))
pred = output.cpu().data.numpy()
assert pred.ndim == 4
pred = pred[:, eval_joints, :, :]
loss_logger.update(loss.item(), batch_size)
acc_logger.update(acc, batch_size)
for i in range(output.shape[0]):
bbox = bboxes[i].tolist()
pose_coords, pose_scores = heatmap_to_coord(pred[i][gt_val_dataset.EVAL_JOINTS], bbox)
keypoints = np.concatenate((pose_coords, pose_scores), axis=1)
keypoints = keypoints.reshape(-1).tolist()
data = dict()
data['bbox'] = bboxes[i].tolist()
data['image_id'] = img_ids[i]
data['score'] = float(np.mean(pose_scores) + np.max(pose_scores))
data['category_id'] = 1
data['keypoints'] = keypoints
kpt_json.append(data)
with open(os.path.join(opt.work_dir, 'test_gt_kpt.json'), 'w') as fid:
json.dump(kpt_json, fid)
#res = evaluate_mAP(os.path.join(opt.work_dir, 'test_gt_kpt.json'), ann_type='keypoints')
#return res['AP']
return loss_logger.avg, acc_logger.avg
def validate_gt(m, opt, cfg, heatmap_to_coord, batch_size=20):
gt_val_dataset = builder.build_dataset(cfg.DATASET.VAL,
preset_cfg=cfg.DATA_PRESET,
train=False)
eval_joints = gt_val_dataset.EVAL_JOINTS
gt_val_loader = torch.utils.data.DataLoader(gt_val_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=20,
drop_last=False)
kpt_json = []
m.eval()
norm_type = cfg.LOSS.get('NORM_TYPE', None)
hm_size = cfg.DATA_PRESET.HEATMAP_SIZE
for inps, labels, label_masks, img_ids, bboxes in tqdm(gt_val_loader,
dynamic_ncols=True):
if isinstance(inps, list):
inps = [inp.cuda() for inp in inps]
else:
inps = inps.cuda()
output = m(inps)
pred = output
assert pred.dim() == 4
pred = pred[:, eval_joints, :, :]
for i in range(output.shape[0]):
bbox = bboxes[i].tolist()
pose_coords, pose_scores = heatmap_to_coord(
pred[i][gt_val_dataset.EVAL_JOINTS],
bbox,
hm_shape=hm_size,
norm_type=norm_type)
keypoints = np.concatenate((pose_coords, pose_scores), axis=1)
keypoints = keypoints.reshape(-1).tolist()
data = dict()
data['bbox'] = bboxes[i].tolist()
data['image_id'] = int(img_ids[i])
data['score'] = float(np.mean(pose_scores) + np.max(pose_scores))
data['category_id'] = 1
data['keypoints'] = keypoints
kpt_json.append(data)
with open(os.path.join(opt.work_dir, 'test_gt_kpt.json'), 'w') as fid:
json.dump(kpt_json, fid)
res = evaluate_mAP(os.path.join(opt.work_dir, 'test_gt_kpt.json'),
ann_type='keypoints',
ann_file=os.path.join(cfg.DATASET.VAL.ROOT,
cfg.DATASET.VAL.ANN))
return res
def validate(m, opt, heatmap_to_coord, batch_size=20):
det_dataset = builder.build_dataset(cfg.DATASET.TEST, preset_cfg=cfg.DATA_PRESET, train=False, opt=opt)
det_loader = torch.utils.data.DataLoader(
det_dataset, batch_size=batch_size, shuffle=False, num_workers=20, drop_last=False)
kpt_json = []
eval_joints = det_dataset.EVAL_JOINTS
test_branch = cfg.OTHERS.get('TEST_BRANCH',True)
m.eval()
norm_type = cfg.LOSS.get('NORM_TYPE', None)
hm_size = cfg.DATA_PRESET.HEATMAP_SIZE
for inps, crop_bboxes, bboxes, img_ids, scores, imghts, imgwds in tqdm(det_loader, dynamic_ncols=True):
if isinstance(inps, list):
inps = [inp.cuda() for inp in inps]
else:
inps = inps.cuda()
output,_ = m(inps,crop_bboxes[:,1,:],crop_bboxes[:,2,:],crop_bboxes[:,3,:])
pred = output
assert pred.dim() == 4
pred = pred[:, eval_joints, :, :]
for i in range(output.shape[0]):
bbox = crop_bboxes[i][0].tolist()
pose_coords, pose_scores = heatmap_to_coord(
pred[i][det_dataset.EVAL_JOINTS], bbox, hm_shape=hm_size, norm_type=norm_type)
keypoints = np.concatenate((pose_coords, pose_scores), axis=1)
keypoints = keypoints.reshape(-1).tolist()
data = dict()
#data['bbox'] = bboxes[i, 0].tolist()
data['bbox'] = bbox
data['image_id'] = int(img_ids[i])
data['score'] = float(scores[i] + np.mean(pose_scores) + np.max(pose_scores))
data['category_id'] = 1
data['keypoints'] = keypoints
kpt_json.append(data)
with open(os.path.join(opt.work_dir, 'test_kpt.json'), 'w') as fid:
json.dump(kpt_json, fid)
res = evaluate_mAP(os.path.join(opt.work_dir, 'test_kpt.json'), ann_type='keypoints', ann_file='/ssd3/Benchmark/coco/annotations/coco_wholebody_val_133.json')#ann_file=os.path.join(cfg.DATASET.VAL.ROOT, cfg.DATASET.VAL.ANN))
return res
def __init__(self, train=True, dpg=False, skip_empty=True, **cfg):
self._cfg = cfg
self._subset_cfg_list = cfg['SET_LIST']
self._preset_cfg = cfg['PRESET']
self._mask_id = [item['MASK_ID'] for item in self._subset_cfg_list]
self.num_joints = self._preset_cfg['NUM_JOINTS']
self._subsets = []
self._subset_size = [0]
for _subset_cfg in self._subset_cfg_list:
subset = build_dataset(_subset_cfg,
preset_cfg=self._preset_cfg,
train=train)
self._subsets.append(subset)
self._subset_size.append(len(subset))
self.cumulative_sizes = self.cumsum(self._subset_size)
def main():
logger.info('******************************')
logger.info(opt)
logger.info('******************************')
logger.info(cfg)
logger.info('******************************')
# Model Initialize
m = preset_model(cfg)
m = nn.DataParallel(m).cuda()
criterion = builder.build_loss(cfg.LOSS).cuda()
if cfg.TRAIN.OPTIMIZER == 'adam':
optimizer = torch.optim.Adam(m.parameters(), lr=cfg.TRAIN.LR)
elif cfg.TRAIN.OPTIMIZER == 'rmsprop':
optimizer = torch.optim.RMSprop(m.parameters(), lr=cfg.TRAIN.LR)
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=cfg.TRAIN.LR_STEP, gamma=cfg.TRAIN.LR_FACTOR)
writer = SummaryWriter('.tensorboard/{}-{}'.format(opt.exp_id,
cfg.FILE_NAME))
train_dataset = builder.build_dataset(cfg.DATASET.TRAIN,
preset_cfg=cfg.DATA_PRESET,
train=True)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE * num_gpu,
shuffle=True,
num_workers=opt.nThreads)
heatmap_to_coord = get_func_heatmap_to_coord(cfg)
opt.trainIters = 0
for i in range(cfg.TRAIN.BEGIN_EPOCH, cfg.TRAIN.END_EPOCH):
opt.epoch = i
current_lr = optimizer.state_dict()['param_groups'][0]['lr']
logger.info(
f'############# Starting Epoch {opt.epoch} | LR: {current_lr} #############'
)
# Training
loss, miou = train(opt, train_loader, m, criterion, optimizer, writer)
logger.epochInfo('Train', opt.epoch, loss, miou)
lr_scheduler.step()
if (i + 1) % opt.snapshot == 0:
# Save checkpoint
torch.save(
m.module.state_dict(),
'./exp/{}-{}/model_{}.pth'.format(opt.exp_id, cfg.FILE_NAME,
opt.epoch))
# Prediction Test
with torch.no_grad():
gt_AP = validate_gt(m.module, opt, cfg, heatmap_to_coord)
rcnn_AP = validate(m.module, opt, heatmap_to_coord)
logger.info(
f'##### Epoch {opt.epoch} | gt mAP: {gt_AP} | rcnn mAP: {rcnn_AP} #####'
)
# Time to add DPG
if i == cfg.TRAIN.DPG_MILESTONE:
torch.save(
m.module.state_dict(),
'./exp/{}-{}/final.pth'.format(opt.exp_id, cfg.FILE_NAME))
# Adjust learning rate
for param_group in optimizer.param_groups:
param_group['lr'] = cfg.TRAIN.LR
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=cfg.TRAIN.DPG_STEP, gamma=0.1)
# Reset dataset
train_dataset = builder.build_dataset(cfg.DATASET.TRAIN,
preset_cfg=cfg.DATA_PRESET,
train=True,
dpg=True)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE * num_gpu,
shuffle=True,
num_workers=opt.nThreads)
torch.save(m.module.state_dict(),
'./exp/{}-{}/final_DPG.pth'.format(opt.exp_id, cfg.FILE_NAME))
def main():
logger.info('******************************')
logger.info(opt)
logger.info('******************************')
logger.info(cfg)
logger.info('******************************')
# Model Initialize
m = preset_model(cfg)
# todo: try to replace with distributedDataParallel to see if it is faster
m = nn.DataParallel(m)
if opt.device.type != 'cpu':
m = m.cuda()
criterion = builder.build_loss(cfg.LOSS)
if opt.device.type != 'cpu':
criterion = criterion.cuda()
if cfg.TRAIN.OPTIMIZER == 'adam':
optimizer = torch.optim.Adam(m.parameters(), lr=cfg.TRAIN.LR)
elif cfg.TRAIN.OPTIMIZER == 'rmsprop':
optimizer = torch.optim.RMSprop(m.parameters(), lr=cfg.TRAIN.LR)
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=cfg.TRAIN.LR_STEP, gamma=cfg.TRAIN.LR_FACTOR)
if opt.clean:
if opt.tensorboard_path.exists():
shutil.rmtree(opt.tensorboard_path)
if opt.experiment_path.exists():
shutil.rmtree(opt.experiment_path)
opt.tensorboard_path.mkdir(exist_ok=True, parents=True)
opt.experiment_path.mkdir(exist_ok=True, parents=True)
writer = SummaryWriter(str(opt.tensorboard_path))
train_dataset = builder.build_dataset(cfg.DATASET.TRAIN,
preset_cfg=cfg.DATA_PRESET,
train=True)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE * max(1, num_gpu),
shuffle=True,
num_workers=opt.nThreads)
heatmap_to_coord = get_func_heatmap_to_coord(cfg)
opt.trainIters = 0
scaler = GradScaler()
for i in range(cfg.TRAIN.BEGIN_EPOCH, cfg.TRAIN.END_EPOCH):
opt.epoch = i
current_lr = optimizer.state_dict()['param_groups'][0]['lr']
logger.info(
f'############# Starting Epoch {opt.epoch} | LR: {current_lr} #############'
)
# Training
loggers = train(opt, train_loader, m, criterion, optimizer, writer,
scaler)
logger.info(
f'Train-{opt.epoch:d} epoch | '
f'{" | ".join(f"{name}:{l.avg:.07f}" for name, l in loggers.items())}'
)
lr_scheduler.step()
if (i + 1) % opt.snapshot == 0:
# Save checkpoint
torch.save(m.module.state_dict(),
str(opt.experiment_path / f'model_{opt.epoch}.pth'))
# Prediction Test
with torch.no_grad():
metrics_on_true_box = validate_gt(m.module, opt, cfg,
heatmap_to_coord)
gt_AP = metrics_on_true_box["map"]
gt_radius_mse = metrics_on_true_box["radius_mse"]
rcnn_AP = validate(m.module, opt, heatmap_to_coord)
logger.info(f'##### Epoch {opt.epoch} | '
f'gt mAP: {gt_AP} | '
f'rcnn mAP: {rcnn_AP} | '
f'gt radius_mse {gt_radius_mse}'
f' #####')
writer.add_scalar(f'Validation/mAP_on_gt_box', gt_AP,
opt.trainIters)
writer.add_scalar(f'Validation/mAP_on_pred_box', rcnn_AP,
opt.trainIters)
writer.add_scalar(f'Validation/radius_mse_on_gt_box',
gt_radius_mse, opt.trainIters)
# Time to add DPG
if i == cfg.TRAIN.DPG_MILESTONE:
torch.save(m.module.state_dict(),
str(opt.experiment_path / "final.pth"))
# Adjust learning rate
for param_group in optimizer.param_groups:
param_group['lr'] = cfg.TRAIN.LR
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=cfg.TRAIN.DPG_STEP, gamma=0.1)
# Reset dataset
train_dataset = builder.build_dataset(cfg.DATASET.TRAIN,
preset_cfg=cfg.DATA_PRESET,
train=True,
dpg=True)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE * max(1, num_gpu),
shuffle=True,
num_workers=opt.nThreads)
torch.save(m.module.state_dict(),
str(opt.experiment_path / 'final_DPG.pth'))
def validate(m, heatmap_to_coord, batch_size=20):
det_dataset = builder.build_dataset(cfg.DATASET.TEST,
preset_cfg=cfg.DATA_PRESET,
train=False,
opt=opt)
eval_joints = det_dataset.EVAL_JOINTS
det_loader = torch.utils.data.DataLoader(det_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=20,
drop_last=False)
kpt_json = []
m.eval()
norm_type = cfg.LOSS.get('NORM_TYPE', None)
hm_size = cfg.DATA_PRESET.HEATMAP_SIZE
combined_loss = (cfg.LOSS.get('TYPE') == 'Combined')
halpe = (cfg.DATA_PRESET.NUM_JOINTS == 133) or (cfg.DATA_PRESET.NUM_JOINTS
== 136)
for inps, crop_bboxes, bboxes, img_ids, scores, imghts, imgwds in tqdm(
det_loader, dynamic_ncols=True):
if isinstance(inps, list):
inps = [inp.cuda() for inp in inps]
else:
inps = inps.cuda()
output = m(inps)
if opt.flip_test:
if isinstance(inps, list):
inps_flip = [flip(inp).cuda() for inp in inps]
else:
inps_flip = flip(inps).cuda()
output_flip = flip_heatmap(m(inps_flip),
det_dataset.joint_pairs,
shift=True)
pred_flip = output_flip[:, eval_joints, :, :]
else:
output_flip = None
pred_flip = None
pred = output
assert pred.dim() == 4
pred = pred[:, eval_joints, :, :]
if output.size()[1] == 68:
face_hand_num = 42
else:
face_hand_num = 110
for i in range(output.shape[0]):
bbox = crop_bboxes[i].tolist()
if combined_loss:
pose_coords_body_foot, pose_scores_body_foot = heatmap_to_coord[
0](pred[i][det_dataset.EVAL_JOINTS[:-face_hand_num]],
bbox,
hm_shape=hm_size,
norm_type=norm_type,
hms_flip=pred_flip[i][
det_dataset.EVAL_JOINTS[:-face_hand_num]]
if pred_flip is not None else None)
pose_coords_face_hand, pose_scores_face_hand = heatmap_to_coord[
1](pred[i][det_dataset.EVAL_JOINTS[-face_hand_num:]],
bbox,
hm_shape=hm_size,
norm_type=norm_type,
hms_flip=pred_flip[i][
det_dataset.EVAL_JOINTS[-face_hand_num:]]
if pred_flip is not None else None)
pose_coords = np.concatenate(
(pose_coords_body_foot, pose_coords_face_hand), axis=0)
pose_scores = np.concatenate(
(pose_scores_body_foot, pose_scores_face_hand), axis=0)
else:
pose_coords, pose_scores = heatmap_to_coord(
pred[i][det_dataset.EVAL_JOINTS],
bbox,
hm_shape=hm_size,
norm_type=norm_type,
hms_flip=pred_flip[i][det_dataset.EVAL_JOINTS]
if pred_flip is not None else None)
keypoints = np.concatenate((pose_coords, pose_scores), axis=1)
keypoints = keypoints.reshape(-1).tolist()
data = dict()
data['bbox'] = bboxes[i, 0].tolist()
data['image_id'] = int(img_ids[i])
data['area'] = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
data['score'] = float(scores[i] + np.mean(pose_scores) +
1.25 * np.max(pose_scores))
# data['score'] = float(scores[i])
data['category_id'] = 1
data['keypoints'] = keypoints
kpt_json.append(data)
if opt.ppose_nms:
from alphapose.utils.pPose_nms import ppose_nms_validate_preprocess, pose_nms, write_json
final_result = []
tmp_data = ppose_nms_validate_preprocess(kpt_json)
for key in tmp_data:
boxes, scores, ids, preds_img, preds_scores = tmp_data[key]
boxes, scores, ids, preds_img, preds_scores, pick_ids = \
pose_nms(boxes, scores, ids, preds_img, preds_scores, 0, cfg.LOSS.get('TYPE') == 'MSELoss')
_result = []
for k in range(len(scores)):
_result.append({
'keypoints':
preds_img[k],
'kp_score':
preds_scores[k],
'proposal_score':
torch.mean(preds_scores[k]) + scores[k] +
1.25 * max(preds_scores[k]),
'idx':
ids[k],
'box': [
boxes[k][0], boxes[k][1], boxes[k][2] - boxes[k][0],
boxes[k][3] - boxes[k][1]
]
})
im_name = str(key).zfill(12) + '.jpg'
result = {'imgname': im_name, 'result': _result}
final_result.append(result)
write_json(final_result,
'./exp/json/',
form='coco',
for_eval=True,
outputfile='validate_rcnn_kpt.json')
else:
if opt.oks_nms:
from alphapose.utils.pPose_nms import oks_pose_nms
kpt_json = oks_pose_nms(kpt_json)
with open('./exp/json/validate_rcnn_kpt.json', 'w') as fid:
json.dump(kpt_json, fid)
sysout = sys.stdout
res = evaluate_mAP('./exp/json/validate_rcnn_kpt.json',
ann_type='keypoints',
ann_file=os.path.join(cfg.DATASET.TEST.ROOT,
cfg.DATASET.TEST.ANN),
halpe=halpe)
sys.stdout = sysout
return res
def validate_gt(m, opt, cfg, heatmap_to_coord, batch_size=64):
joint_radius_mse = DataLogger()
gt_val_dataset = builder.build_dataset(cfg.DATASET.VAL,
preset_cfg=cfg.DATA_PRESET,
train=False)
eval_joints = gt_val_dataset.EVAL_JOINTS
gt_val_loader = torch.utils.data.DataLoader(gt_val_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=opt.nThreads,
drop_last=False)
kpt_json = []
m.eval()
norm_type = cfg.LOSS.get('NORM_TYPE', None)
hm_size = cfg.DATA_PRESET.HEATMAP_SIZE
mse_loss = nn.MSELoss()
for index, (inps, labels, label_masks, joint_radius_gt, img_ids,
bboxes) in tqdm(enumerate(gt_val_loader), dynamic_ncols=True):
if opt.device.type != 'cpu':
if isinstance(inps, list):
inps = [inp.cuda() for inp in inps]
else:
inps = inps.cuda()
full_output = m(inps)
joints_map = full_output.joints_map
joints_radius = full_output.joints_radius
pred = joints_map
assert pred.dim() == 4
pred = pred[:, eval_joints, :, :]
for i in range(joints_map.shape[0]):
bbox = bboxes[i].tolist()
pose_coords, pose_scores = heatmap_to_coord(
pred[i][gt_val_dataset.EVAL_JOINTS],
bbox,
hm_shape=hm_size,
norm_type=norm_type)
keypoints = np.concatenate((pose_coords, pose_scores), axis=1)
keypoints = keypoints.reshape(-1).tolist()
data = dict()
data['bbox'] = bboxes[i].tolist()
data['image_id'] = int(img_ids[i])
data['score'] = float(np.mean(pose_scores) + np.max(pose_scores))
data['category_id'] = 1
data['keypoints'] = keypoints
kpt_json.append(data)
radius_masks = (label_masks[:, :, 0, 0] != 0) & (joint_radius_gt != -1)
joints_radius = joints_radius[:, eval_joints]
joint_radius_gt = joint_radius_gt[:, eval_joints]
joint_radius_gt = joint_radius_gt[radius_masks]
joints_radius = joints_radius[radius_masks]
joints_radius_error = mse_loss(joint_radius_gt, joints_radius.cpu())
joint_radius_mse.update(joints_radius_error, joint_radius_gt.shape[0])
with open(os.path.join(opt.work_dir, 'test_gt_kpt.json'), 'w') as fid:
json.dump(kpt_json, fid)
res = evaluate_mAP(os.path.join(opt.work_dir, 'test_gt_kpt.json'),
ann_type='keypoints',
ann_file=os.path.join(cfg.DATASET.VAL.ROOT,
cfg.DATASET.VAL.ANN))
return {
"map": res,
"radius_mse": joint_radius_mse.avg,
}
def validate_gt(m, opt, cfg, heatmap_to_coord, batch_size=20):
gt_val_dataset = builder.build_dataset(cfg.DATASET.VAL, preset_cfg=cfg.DATA_PRESET, train=False)
eval_joints = gt_val_dataset.EVAL_JOINTS
test_branch = cfg.OTHERS.get('TEST_BRANCH',True)
gt_val_loader = torch.utils.data.DataLoader(
gt_val_dataset, batch_size=batch_size, shuffle=False, num_workers=20, drop_last=False)
kpt_json = []
kpt_json_branch = []
m.eval()
norm_type = cfg.LOSS.get('NORM_TYPE', None)
hm_size = cfg.DATA_PRESET.HEATMAP_SIZE
for inps, labels, label_masks, img_ids, bboxes in tqdm(gt_val_loader, dynamic_ncols=True):
if isinstance(inps, list):
inps = [inp.cuda() for inp in inps]
else:
inps = inps.cuda()
output,feature = m(inps)
pred = copy.deepcopy(output)
assert pred.dim() == 4
pred = pred[:, eval_joints, :, :]
for i in range(output.shape[0]):
bbox = bboxes[i][0].tolist()
pose_coords, pose_scores = heatmap_to_coord(
pred[i][gt_val_dataset.EVAL_JOINTS], bbox, hm_shape=hm_size, norm_type=norm_type)
keypoints = np.concatenate((pose_coords, pose_scores), axis=1)
keypoints = keypoints.reshape(-1).tolist()
data = dict()
#data['bbox'] = bboxes[i, 0].tolist()
data['bbox'] = bbox
data['image_id'] = int(img_ids[i])
data['score'] = float(np.mean(pose_scores) + np.max(pose_scores))
data['category_id'] = 1
data['keypoints'] = keypoints
kpt_json.append(data)
if test_branch:
hm_height, hm_width = hm_size
# regression the joints of wholeboy in stage1
pred_jts, pred_score = _integral_tensor(
pred, 133, False, hm_width, hm_height, 1, integral_operation=integral_op, norm_type='sigmoid')
pred_jts = pred_jts.reshape(pred_jts.shape[0], 133, 2)
# get the coords with the size of heatmap
coords_x = (pred_jts[:, :, 0] + 0.5) * hm_width
coords_y = (pred_jts[:, :, 1] + 0.5) * hm_height
# get the box of hands for roi align
lefthand_boxes = get_box_for_align(coords_x[:,-42:-21],coords_y[:,-42:-21])
righthand_boxes = get_box_for_align(coords_x[:,-21:],coords_y[:,-21:])
# stage2 testing
fine_out = m.forward_branch(output, feature, lefthand_boxes, righthand_boxes)
# output contains the finer and amplified hands kpts, need to apply aff
fine_pred_jts, fine_pred_score = _integral_tensor(
fine_out[:,-42:,:,:], 42, False, hm_width, hm_height, 1, integral_operation=integral_op, norm_type='sigmoid')
fine_pred_jts = fine_pred_jts.reshape(fine_pred_jts.shape[0], 42, 2)
lefthand_jts = fine_pred_jts[:,:21,:]
righthand_jts = fine_pred_jts[:,21:,:]
lefthand_jts[:,:,0] = (lefthand_jts[:,:,0]+0.5)*hm_width
lefthand_jts[:,:,1] = (lefthand_jts[:,:,1]+0.5)*hm_height
righthand_jts[:,:,0] = (righthand_jts[:,:,0]+0.5)*hm_width
righthand_jts[:,:,1] = (righthand_jts[:,:,1]+0.5)*hm_height
center_hm = np.array([hm_width/2.0,hm_height/2.0])
scale_hm = np.array([hm_size[1],hm_size[0]])
lefthand_kpts = copy.deepcopy(lefthand_jts.cpu().numpy().astype(np.float32))
righthand_kpts = copy.deepcopy(righthand_jts.cpu().numpy().astype(np.float32))
# apply affine trans to lefthand and add offset
for j in range(lefthand_jts.shape[0]):
box = lefthand_boxes[j].tolist()
width = np.array(box[2] - box[0])
height = np.array(box[3] - box[1])
output_size = [box[2]-box[0],box[3]-box[1]]
offset = np.array([box[0],box[1]])
trans = get_affine_transform(center_hm,scale_hm,0,output_size)
for k in range(21):
lefthand_kpts[j ,k, 0:2] = affine_transform(lefthand_kpts[j ,k, 0:2], trans)
lefthand_kpts[j,:,0] = (lefthand_kpts[j,:,0]) + offset[0]
lefthand_kpts[j,:,1] = (lefthand_kpts[j,:,1])+ offset[1]
#--------------------------------------------------
# apply affine trans to righthand and add offset
for j in range(righthand_jts.shape[0]):
box = righthand_boxes[j].tolist()
width = np.array(box[2] - box[0])
height = np.array(box[3] - box[1])
output_size = [box[2]-box[0],box[3]-box[1]]
offset = np.array([box[0],box[1]])
trans = get_affine_transform(center_hm,scale_hm,0,output_size)
for k in range(21):
righthand_kpts[j,k, 0:2] = affine_transform(righthand_kpts[j ,k, 0:2], trans)
righthand_kpts[j,:,0] = (righthand_kpts[j,:,0]) + offset[0]
righthand_kpts[j,:,1] = (righthand_kpts[j,:,1]) + offset[1]
#--------------------------------------------------
bodyface_kpts = copy.deepcopy(pred_jts[:,:-42,:].cpu().numpy().astype(np.float32))
bodyface_kpts[:,:,0] = (bodyface_kpts[:,:,0]+0.5)*hm_width
bodyface_kpts[:,:,1] = (bodyface_kpts[:,:,1]+0.5)*hm_height
fine_kpts = np.concatenate((bodyface_kpts,lefthand_kpts,righthand_kpts), axis=1)
fine_socre = np.concatenate((pred_score[:,:-42,:].cpu().numpy(),fine_pred_score.cpu().numpy()), axis=1)
for n in range(output.shape[0]):
bbox = bboxes[n][0].tolist()
xmin, ymin, xmax, ymax = bbox
w = xmax - xmin
h = ymax - ymin
center = np.array([xmin + w * 0.5, ymin + h * 0.5])
scale = np.array([w, h])
for l in range(fine_kpts.shape[1]):
fine_kpts[n, l, 0:2] = transform_preds(fine_kpts[n, l, 0:2], center, scale,
[hm_size[1],hm_size[0]])
keypoints = np.concatenate((fine_kpts[n], fine_socre[n]), axis=1)
keypoints = keypoints.reshape(-1).tolist()
data_branch = dict()
#data['bbox'] = bboxes[i, 0].tolist()
data_branch['bbox'] = bbox
data_branch['image_id'] = int(img_ids[n])
data_branch['score'] = float(np.mean(fine_socre) + np.max(fine_socre))
data_branch['category_id'] = 1
data_branch['keypoints'] = keypoints
kpt_json_branch.append(data_branch)
with open(os.path.join(opt.work_dir, 'test_gt_kpt.json'), 'w') as fid:
json.dump(kpt_json, fid)
res = evaluate_mAP(os.path.join(opt.work_dir, 'test_gt_kpt.json'), ann_type='keypoints', ann_file='/ssd3/Benchmark/coco/annotations/coco_wholebody_val_133.json')#ann_file=os.path.join(cfg.DATASET.VAL.ROOT, cfg.DATASET.VAL.ANN))
if test_branch:
with open(os.path.join(opt.work_dir, 'test_gt_kpt_2branch.json'), 'w') as fid2:
json.dump(kpt_json_branch, fid2)
res_branch = evaluate_mAP(os.path.join(opt.work_dir, 'test_gt_kpt_2branch.json'), ann_type='keypoints', ann_file='/ssd3/Benchmark/coco/annotations/coco_wholebody_val_133.json')
return res,res_branch
else:
return res, 0
def validate_gt(cfg,
batchsize=1,
engine_file_path=None,
m=None,
heatmap_to_coord=None):
gt_val_dataset = builder.build_dataset(cfg.DATASET.VAL,
preset_cfg=cfg.DATA_PRESET,
train=False)
eval_joints = gt_val_dataset.EVAL_JOINTS
gt_val_loader = torch.utils.data.DataLoader(gt_val_dataset,
batch_size=batchsize,
shuffle=False,
num_workers=20,
drop_last=True)
kpt_json = []
if m:
m.eval()
norm_type = cfg.LOSS.get('NORM_TYPE', None)
hm_size = cfg.DATA_PRESET.HEATMAP_SIZE
average_time = 0
pytorch_all_infer_time = 0
trt_all_infer_time = 0
data_num = 0
for inps, labels, label_masks, img_ids, bboxes in tqdm(gt_val_loader,
dynamic_ncols=True):
data_num += 1
if engine_file_path:
# hm_data = []
inps = inps.numpy()
np.copyto(inputs[0].host, inps.ravel())
trt_infer_before_time = count_time()
trt_outputs = trt_common.do_inference(context,
bindings=bindings,
inputs=inputs,
outputs=outputs,
stream=stream)
trt_infer_after_time = count_time()
trt_all_infer_time += trt_infer_after_time - trt_infer_before_time
pred = trt_outputs[0].reshape(-1, 17, 64, 48)
else:
if isinstance(inps, list):
inps = [inp.cuda() for inp in inps]
else:
inps = inps.cuda()
pytorch_infer_before_time = count_time()
output = m(inps)
pytorch_infer_after_time = count_time()
pytorch_all_infer_time += pytorch_infer_after_time - pytorch_infer_before_time
if opt.flip_test:
if isinstance(inps, list):
inps_flip = [flip(inp).cuda() for inp in inps]
else:
inps_flip = flip(inps).cuda()
output_flip = flip_heatmap(m(inps_flip),
gt_val_dataset.joint_pairs,
shift=True)
pred_flip = output_flip[:, eval_joints, :, :]
else:
output_flip = None
pred = output
assert pred.dim() == 4
pred = pred[:, eval_joints, :, :]
for i in range(pred.shape[0]): #后处理过程
bbox = bboxes[i].tolist()
if engine_file_path:
pose_coords, pose_scores = heatmap_to_coord_simple(
pred[i], bbox, hm_shape=hm_size, norm_type=norm_type)
else:
pose_coords, pose_scores = heatmap_to_coord(
pred[i], bbox, hm_shape=hm_size, norm_type=norm_type)
keypoints = np.concatenate((pose_coords, pose_scores), axis=1)
keypoints = keypoints.reshape(-1).tolist()
data = dict()
data['bbox'] = bboxes[i].tolist()
data['image_id'] = int(img_ids[i])
data['score'] = float(np.mean(pose_scores) + np.max(pose_scores))
data['category_id'] = 1
data['keypoints'] = keypoints
kpt_json.append(data)
if engine_file_path:
average_time = float(trt_all_infer_time / data_num)
else:
average_time = (pytorch_all_infer_time / data_num)
print("average_time:", average_time)
res_file = r"data/coco/res.json"
with open(res_file, 'w') as F:
json.dump(kpt_json, F)
res = evaluate_mAP(res_file,
ann_type='keypoints',
ann_file=os.path.join(cfg.DATASET.VAL.ROOT,
cfg.DATASET.VAL.ANN))
return res, average_time
def validate(m, opt, heatmap_to_coord, batch_size=20):
det_dataset = builder.build_dataset(cfg.DATASET.TEST, preset_cfg=cfg.DATA_PRESET, train=False, opt=opt)
det_loader = torch.utils.data.DataLoader(
det_dataset, batch_size=batch_size, shuffle=False, num_workers=20, drop_last=False)
kpt_json = []
eval_joints = det_dataset.EVAL_JOINTS
m.eval()
norm_type = cfg.LOSS.get('NORM_TYPE', None)
hm_size = cfg.DATA_PRESET.HEATMAP_SIZE
combined_loss = (cfg.LOSS.get('TYPE') == 'Combined')
halpe = (cfg.DATA_PRESET.NUM_JOINTS == 133) or (cfg.DATA_PRESET.NUM_JOINTS == 136)
for inps, crop_bboxes, bboxes, img_ids, scores, imghts, imgwds in tqdm(det_loader, dynamic_ncols=True):
if isinstance(inps, list):
inps = [inp.cuda() for inp in inps]
else:
inps = inps.cuda()
output = m(inps)
pred = output
assert pred.dim() == 4
pred = pred[:, eval_joints, :, :]
if output.size()[1] == 68:
face_hand_num = 42
else:
face_hand_num = 110
for i in range(output.shape[0]):
bbox = crop_bboxes[i].tolist()
if combined_loss:
pose_coords_body_foot, pose_scores_body_foot = heatmap_to_coord[0](
pred[i][det_dataset.EVAL_JOINTS[:-face_hand_num]], bbox, hm_shape=hm_size, norm_type=norm_type)
pose_coords_face_hand, pose_scores_face_hand = heatmap_to_coord[1](
pred[i][det_dataset.EVAL_JOINTS[-face_hand_num:]], bbox, hm_shape=hm_size, norm_type=norm_type)
pose_coords = np.concatenate((pose_coords_body_foot, pose_coords_face_hand), axis=0)
pose_scores = np.concatenate((pose_scores_body_foot, pose_scores_face_hand), axis=0)
else:
pose_coords, pose_scores = heatmap_to_coord(
pred[i][det_dataset.EVAL_JOINTS], bbox, hm_shape=hm_size, norm_type=norm_type)
keypoints = np.concatenate((pose_coords, pose_scores), axis=1)
keypoints = keypoints.reshape(-1).tolist()
data = dict()
data['bbox'] = bboxes[i, 0].tolist()
data['image_id'] = int(img_ids[i])
data['score'] = float(scores[i] + np.mean(pose_scores) + 1.25 * np.max(pose_scores))
data['category_id'] = 1
data['keypoints'] = keypoints
kpt_json.append(data)
sysout = sys.stdout
with open(os.path.join(opt.work_dir, 'test_kpt.json'), 'w') as fid:
json.dump(kpt_json, fid)
res = evaluate_mAP(os.path.join(opt.work_dir, 'test_kpt.json'), ann_type='keypoints', ann_file=os.path.join(cfg.DATASET.VAL.ROOT, cfg.DATASET.VAL.ANN), halpe=halpe)
sys.stdout = sysout
return res
def validate_gt(m, cfg, heatmap_to_coord, batch_size=20):
gt_val_dataset = builder.build_dataset(cfg.DATASET.VAL,
preset_cfg=cfg.DATA_PRESET,
train=False)
eval_joints = gt_val_dataset.EVAL_JOINTS
gt_val_loader = torch.utils.data.DataLoader(gt_val_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=20,
drop_last=False)
kpt_json = []
m.eval()
norm_type = cfg.LOSS.get('NORM_TYPE', None)
hm_size = cfg.DATA_PRESET.HEATMAP_SIZE
combined_loss = (cfg.LOSS.get('TYPE') == 'Combined')
halpe = (cfg.DATA_PRESET.NUM_JOINTS == 133) or (cfg.DATA_PRESET.NUM_JOINTS
== 136)
for inps, labels, label_masks, img_ids, bboxes in tqdm(gt_val_loader,
dynamic_ncols=True):
if isinstance(inps, list):
inps = [inp.cuda() for inp in inps]
else:
inps = inps.cuda()
output = m(inps)
if opt.flip_test:
if isinstance(inps, list):
inps_flip = [flip(inp).cuda() for inp in inps]
else:
inps_flip = flip(inps).cuda()
output_flip = flip_heatmap(m(inps_flip),
gt_val_dataset.joint_pairs,
shift=True)
pred_flip = output_flip[:, eval_joints, :, :]
else:
output_flip = None
pred_flip = None
pred = output
assert pred.dim() == 4
pred = pred[:, eval_joints, :, :]
if output.size()[1] == 68:
face_hand_num = 42
else:
face_hand_num = 110
for i in range(output.shape[0]):
bbox = bboxes[i].tolist()
if combined_loss:
pose_coords_body_foot, pose_scores_body_foot = heatmap_to_coord[
0](pred[i][gt_val_dataset.EVAL_JOINTS[:-face_hand_num]],
bbox,
hm_shape=hm_size,
norm_type=norm_type,
hms_flip=pred_flip[i][
gt_val_dataset.EVAL_JOINTS[:-face_hand_num]]
if pred_flip is not None else None)
pose_coords_face_hand, pose_scores_face_hand = heatmap_to_coord[
1](pred[i][gt_val_dataset.EVAL_JOINTS[-face_hand_num:]],
bbox,
hm_shape=hm_size,
norm_type=norm_type,
hms_flip=pred_flip[i][
gt_val_dataset.EVAL_JOINTS[-face_hand_num:]]
if pred_flip is not None else None)
pose_coords = np.concatenate(
(pose_coords_body_foot, pose_coords_face_hand), axis=0)
pose_scores = np.concatenate(
(pose_scores_body_foot, pose_scores_face_hand), axis=0)
else:
pose_coords, pose_scores = heatmap_to_coord(
pred[i][gt_val_dataset.EVAL_JOINTS],
bbox,
hm_shape=hm_size,
norm_type=norm_type,
hms_flip=pred_flip[i][gt_val_dataset.EVAL_JOINTS]
if pred_flip is not None else None)
keypoints = np.concatenate((pose_coords, pose_scores), axis=1)
keypoints = keypoints.reshape(-1).tolist()
data = dict()
data['bbox'] = bboxes[i].tolist()
data['image_id'] = int(img_ids[i])
data['score'] = float(
np.mean(pose_scores) + 1.25 * np.max(pose_scores))
data['category_id'] = 1
data['keypoints'] = keypoints
kpt_json.append(data)
sysout = sys.stdout
with open('./exp/json/validate_gt_kpt.json', 'w') as fid:
json.dump(kpt_json, fid)
res = evaluate_mAP('./exp/json/validate_gt_kpt.json',
ann_type='keypoints',
ann_file=os.path.join(cfg.DATASET.VAL.ROOT,
cfg.DATASET.VAL.ANN),
halpe=halpe)
sys.stdout = sysout
return res
