import numpy as np import torch import torch.nn as nn from torch.autograd import Variable from torch.utils.data import DataLoader from torchvision import transforms import torch.backends.cudnn as cudnn import torchvision import torch.nn.functional as F import cv2 import matplotlib.pyplot as plt import sys import os import argparse import datasets import hopenet import utils from PIL import Image def parse_args(): """Parse input arguments.""" parser = argparse.ArgumentParser(description='Head pose estimation using the Hopenet network.') parser.add_argument('--gpu', dest='gpu_id', help='GPU device id to use [0]', default=0, type=int) parser.add_argument('--data_dir', dest='data_dir', help='Directory path for data.', default='', type=str) parser.add_argument('--filename_list', dest='filename_list', help='Path to text file containing relative paths for every example.', default='', type=str) parser.add_argument('--snapshot', dest='snapshot', help='Name of model snapshot.', 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) parser.add_argument('--dataset', dest='dataset', help='Dataset type.', default='AFLW2000', type=str) args = parser.parse_args() return args if __name__ == '__main__': args = parse_args() cudnn.enabled = True gpu = args.gpu_id snapshot_path = args.snapshot model = hopenet.Hopenet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], 66, 0) print 'Loading snapshot.' # Load snapshot saved_state_dict = torch.load(snapshot_path) model.load_state_dict(saved_state_dict) print 'Loading data.' transformations = transforms.Compose([transforms.Scale(224), transforms.CenterCrop(224), transforms.ToTensor()]) if args.dataset == 'AFLW2000': pose_dataset = datasets.AFLW2000(args.data_dir, args.filename_list, transformations) elif args.dataset == 'AFLW2000_ds': pose_dataset = datasets.AFLW2000_ds(args.data_dir, args.filename_list, transformations) elif args.dataset == 'BIWI': pose_dataset = datasets.BIWI(args.data_dir, args.filename_list, transformations) elif args.dataset == 'AFLW': pose_dataset = datasets.AFLW(args.data_dir, args.filename_list, transformations) elif args.dataset == 'Pose_300W_LP': pose_dataset = datasets.Pose_300W_LP(args.data_dir, args.filename_list, transformations) elif args.dataset == 'AFW': pose_dataset = datasets.AFW(args.data_dir, args.filename_list, transformations) else: print 'Error: not a valid dataset name' sys.exit() test_loader = torch.utils.data.DataLoader(dataset=pose_dataset, batch_size=args.batch_size, num_workers=2) model.cuda(gpu) print 'Ready to test network.' # Super-resolution model sr_model = torch.load('data/sr_model/model_epoch_50.pth')["model"] sr_model = sr_model.cuda(gpu) # Test the Model model.eval() # Change model to 'eval' mode (BN uses moving mean/var). total = 0 idx_tensor = [idx for idx in xrange(66)] idx_tensor = torch.FloatTensor(idx_tensor).cuda(gpu) yaw_error = .0 pitch_error = .0 roll_error = .0 l1loss = torch.nn.L1Loss(size_average=False) for i, (images, labels, cont_labels, name) in enumerate(test_loader): ### START Super-resolution ### # To new color space img = transforms.ToPILImage()(images[0]) print img img = img.convert('YCbCr') img_y, img_cb, img_cr = img.split() # Super-resolution img_y_var = Variable(transforms.ToTensor()(img_y)).view(1, -1, img_y.size[0], img_y.size[1]).cuda(gpu) / 255. out_sr = sr_model(img_y_var) img_h_y = out_sr.data[0].cpu().numpy().astype(np.float32) img_h_y = img_h_y * 255 img_h_y[img_h_y<0] = 0 img_h_y[img_h_y>255.] = 255. img_h_y = img_h_y[0] img_new = np.zeros((img_h_y.shape[0], img_h_y.shape[1], 3), np.uint8) img_new[:,:,0] = img_h_y # img_new[:,:,0] = np.asarray(img_y) img_new[:,:,1] = np.asarray(img_cb) img_new[:,:,2] = np.asarray(img_cr) img_new = Image.fromarray(img_new, "YCbCr").convert("RGB") # To tensor and normalize img_new.save('output/test_superres/' + name[0] + '.jpg', "JPEG") img = transforms.ToTensor()(img_new) img = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(img) images = Variable(img.view(1,-1,img.shape[1],img.shape[2])).cuda(gpu) ### END Super-resolution ### total += cont_labels.size(0) label_yaw = cont_labels[:,0].float() label_pitch = cont_labels[:,1].float() label_roll = cont_labels[:,2].float() yaw, pitch, roll, angles = model(images) # Binned predictions _, yaw_bpred = torch.max(yaw.data, 1) _, pitch_bpred = torch.max(pitch.data, 1) _, roll_bpred = torch.max(roll.data, 1) # Continuous predictions yaw_predicted = utils.softmax_temperature(yaw.data, 1) pitch_predicted = utils.softmax_temperature(pitch.data, 1) roll_predicted = utils.softmax_temperature(roll.data, 1) yaw_predicted = torch.sum(yaw_predicted * idx_tensor, 1).cpu() * 3 - 99 pitch_predicted = torch.sum(pitch_predicted * idx_tensor, 1).cpu() * 3 - 99 roll_predicted = torch.sum(roll_predicted * idx_tensor, 1).cpu() * 3 - 99 # Mean absolute error yaw_error += torch.sum(torch.abs(yaw_predicted - label_yaw)) pitch_error += torch.sum(torch.abs(pitch_predicted - label_pitch)) roll_error += torch.sum(torch.abs(roll_predicted - label_roll)) # Save images with pose cube. # TODO: fix for larger batch size if args.save_viz: name = name[0] if args.dataset == 'BIWI': cv2_img = cv2.imread(os.path.join(args.data_dir, name + '_rgb.png')) else: cv2_img = cv2.imread(os.path.join(args.data_dir, name + '.jpg')) if args.batch_size == 1: error_string = 'y %.2f, p %.2f, r %.2f' % (torch.sum(torch.abs(yaw_predicted - label_yaw)), torch.sum(torch.abs(pitch_predicted - label_pitch)), torch.sum(torch.abs(roll_predicted - label_roll))) cv2.putText(cv2_img, error_string, (30, cv2_img.shape[0]- 30), fontFace=1, fontScale=1, color=(0,0,255), thickness=1) utils.plot_pose_cube(cv2_img, yaw_predicted[0], pitch_predicted[0], roll_predicted[0]) 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))