From 61526433d2f56a669dd077de7920ada32b6008ad Mon Sep 17 00:00:00 2001 From: natanielruiz <nataniel777@hotmail.com> Date: 星期二, 11 七月 2017 11:21:08 +0800 Subject: [PATCH] next --- code/test_resnet_bins.py | 38 ++++++++++++++++++++++++++++---------- 1 files changed, 28 insertions(+), 10 deletions(-) diff --git a/code/test_resnet_bins.py b/code/test_resnet_bins.py index 30aa158..8d0eaec 100644 --- a/code/test_resnet_bins.py +++ b/code/test_resnet_bins.py @@ -31,6 +31,8 @@ default='', type=str) parser.add_argument('--batch_size', dest='batch_size', help='Batch size.', default=1, type=int) + parser.add_argument('--save_viz', dest='save_viz', help='Save images with pose cube.', + default=False, type=bool) args = parser.parse_args() @@ -92,33 +94,49 @@ label_roll = labels[:,2] yaw, pitch, roll = model(images) - # _, yaw_predicted = torch.max(yaw.data, 1) - # _, pitch_predicted = torch.max(pitch.data, 1) - # _, roll_predicted = torch.max(roll.data, 1) + + # Binned predictions + _, yaw_bpred = torch.max(yaw.data, 1) + _, pitch_bpred = torch.max(pitch.data, 1) + _, roll_bpred = torch.max(roll.data, 1) yaw_predicted = F.softmax(yaw) pitch_predicted = F.softmax(pitch) roll_predicted = F.softmax(roll) + # Continuous predictions yaw_predicted = torch.sum(yaw_predicted.data[0] * idx_tensor) pitch_predicted = torch.sum(pitch_predicted.data[0] * idx_tensor) roll_predicted = torch.sum(roll_predicted.data[0] * idx_tensor) + # Mean absolute error yaw_error += abs(yaw_predicted - label_yaw[0]) * 3 pitch_error += abs(pitch_predicted - label_pitch[0]) * 3 roll_error += abs(roll_predicted - label_roll[0]) * 3 - # print yaw_predicted * 3, label_yaw[0] * 3, abs(yaw_predicted - label_yaw[0]) * 3 + # Binned Accuracy + # for er in xrange(n_margins): + # yaw_bpred[er] += (label_yaw[0] in range(yaw_bpred[0,0] - er, yaw_bpred[0,0] + er + 1)) + # pitch_bpred[er] += (label_pitch[0] in range(pitch_bpred[0,0] - er, pitch_bpred[0,0] + er + 1)) + # roll_bpred[er] += (label_roll[0] in range(roll_bpred[0,0] - er, roll_bpred[0,0] + er + 1)) - # for er in xrange(0,n_margins): - # yaw_correct[er] += (label_yaw[0] in range(yaw_predicted[0,0] - er, yaw_predicted[0,0] + er + 1)) - # pitch_correct[er] += (label_pitch[0] in range(pitch_predicted[0,0] - er, pitch_predicted[0,0] + er + 1)) - # roll_correct[er] += (label_roll[0] in range(roll_predicted[0,0] - er, roll_predicted[0,0] + er + 1)) + # print label_yaw[0], yaw_bpred[0,0] - # print label_yaw[0], yaw_predicted[0,0] - # 4 -> 15 + # Save images with pose cube. + if args.save_viz: + name = name[0] + cv2_img = cv2.imread(os.path.join(args.data_dir, name + '.jpg')) + #cv2_img = cv2.cvtColor(cv2_img, cv2.COLOR_RGB2BGR) + #print name + #print label_yaw[0] * 3 - 99, label_pitch[0] * 3 - 99, label_roll[0] * 3 - 99 + #print yaw_predicted * 3 - 99, pitch_predicted * 3 - 99, roll_predicted * 3 - 99 + utils.plot_pose_cube(cv2_img, yaw_predicted * 3 - 99, pitch_predicted * 3 - 99, roll_predicted * 3 - 99) + cv2.imwrite(os.path.join('output/images', name + '.jpg'), cv2_img) + print('Test error in degrees of the model on the ' + str(total) + ' test images. Yaw: %.4f, Pitch: %.4f, Roll: %.4f' % (yaw_error / total, pitch_error / total, roll_error / total)) + + # Binned accuracy # for idx in xrange(len(yaw_correct)): # print yaw_correct[idx] / total, pitch_correct[idx] / total, roll_correct[idx] / total -- Gitblit v1.8.0