def main():
args = parse_args()
cfg.set_args(args.gpu_ids)
cudnn.fastest = True
cudnn.benchmark = True
cudnn.deterministic = False
cudnn.enabled = True
tester = Tester(args.backbone)
tester._make_batch_generator()
for epoch in range(args.model_epoch[0], args.model_epoch[1]):
tester._make_model(epoch)
preds = []
with torch.no_grad():
for itr, input_img in enumerate(tqdm(tester.batch_generator)):
# forward
coord_out = tester.model(input_img)
if cfg.flip_test:
flipped_input_img = flip(input_img, dims=3)
flipped_coord_out = tester.model(flipped_input_img)
flipped_coord_out[:, :, 0] = cfg.output_shape[
1] - flipped_coord_out[:, :, 0] - 1
for pair in tester.flip_pairs:
flipped_coord_out[:, pair[
0], :], flipped_coord_out[:, pair[
1], :] = flipped_coord_out[:, pair[1], :].clone(
), flipped_coord_out[:, pair[0], :].clone()
coord_out = (coord_out + flipped_coord_out) / 2.
vis = False
if vis:
filename = str(itr)
tmpimg = input_img[0].cpu().numpy()
tmpimg = tmpimg * np.array(cfg.pixel_std).reshape(
3, 1, 1) + np.array(cfg.pixel_mean).reshape(3, 1, 1)
tmpimg = tmpimg.astype(np.uint8)
tmpimg = tmpimg[::-1, :, :]
tmpimg = np.transpose(tmpimg, (1, 2, 0)).copy()
tmpkps = np.zeros((3, tester.joint_num))
tmpkps[:2, :] = coord_out[
0, :, :2].cpu().numpy().transpose(
1, 0) / cfg.output_shape[0] * cfg.input_shape[0]
tmpkps[2, :] = 1
tmpimg = vis_keypoints(tmpimg, tmpkps, tester.skeleton)
cv2.imwrite(filename + '_output.jpg', tmpimg)
coord_out = coord_out.cpu().numpy()
preds.append(coord_out)
# evaluate
preds = np.concatenate(preds, axis=0)
tester._evaluate(preds, cfg.result_dir)
def main():
args = parse_args()
cfg.set_args(args.gpu_ids)
cudnn.fastest = True
cudnn.benchmark = True
cudnn.deterministic = False
cudnn.enabled = True
tester = Tester(args.test_epoch)
tester._make_batch_generator()
tester._make_model()
preds = []
tmpp = 0
with torch.no_grad():
for itr, input_img in enumerate(tqdm(tester.batch_generator)):
# forward
coord_out = tester.model(input_img)
if cfg.flip_test:
flipped_input_img = flip(input_img, dims=3)
flipped_coord_out = tester.model(flipped_input_img)
flipped_coord_out[:, :, 0] = cfg.output_shape[
1] - flipped_coord_out[:, :, 0] - 1
for pair in tester.flip_pairs:
flipped_coord_out[:, pair[0], :], flipped_coord_out[:, pair[
1], :] = flipped_coord_out[:, pair[1], :].clone(
), flipped_coord_out[:, pair[0], :].clone()
coord_out = (coord_out + flipped_coord_out) / 2.
vis = True
if vis:
filename = str(itr)
tmpimg = input_img[0].cpu().numpy()
tmpimg = tmpimg * np.array(cfg.pixel_std).reshape(
3, 1, 1) + np.array(cfg.pixel_mean).reshape(3, 1, 1)
tmpimg = tmpimg.astype(np.uint8)
tmpimg = tmpimg[::-1, :, :]
tmpimg = np.transpose(tmpimg, (1, 2, 0)).copy()
tmpkps = np.zeros((3, tester.joint_num))
tmpkps[:2, :] = coord_out[0, :, :2].cpu().numpy().transpose(
1, 0) / cfg.output_shape[0] * cfg.input_shape[0]
tmpkps[2, :] = 1
tmpimg = vis_keypoints(tmpimg, tmpkps, tester.skeleton)
# print(tmpkps)
cv2.imwrite(filename + '_output.jpg', tmpimg)
coord_out = coord_out.cpu().numpy()
preds.append(coord_out)
tmpp += 1
if tmpp == 5: break
# evaluate
preds = np.concatenate(preds, axis=0)
# print(preds.shape)
# It takes testest from common/base.py
# It then calls evaluate of that dataset
# from data/MSCOCO/MSCOCO.py call visualise keypoints
# from common/utils/vis.py the plts are plotted
tester._evaluate(preds, cfg.result_dir)