Spaces:
Sleeping
Sleeping
| import os, argparse | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| 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 | |
| from l2cs import select_device, natural_keys, gazeto3d, angular, getArch, L2CS, Gaze360, Mpiigaze | |
| def parse_args(): | |
| """Parse input arguments.""" | |
| parser = argparse.ArgumentParser( | |
| description='Gaze estimation using L2CSNet .') | |
| # Gaze360 | |
| parser.add_argument( | |
| '--gaze360image_dir', dest='gaze360image_dir', help='Directory path for gaze images.', | |
| default='datasets/Gaze360/Image', type=str) | |
| parser.add_argument( | |
| '--gaze360label_dir', dest='gaze360label_dir', help='Directory path for gaze labels.', | |
| default='datasets/Gaze360/Label/test.label', type=str) | |
| # mpiigaze | |
| parser.add_argument( | |
| '--gazeMpiimage_dir', dest='gazeMpiimage_dir', help='Directory path for gaze images.', | |
| default='datasets/MPIIFaceGaze/Image', type=str) | |
| parser.add_argument( | |
| '--gazeMpiilabel_dir', dest='gazeMpiilabel_dir', help='Directory path for gaze labels.', | |
| default='datasets/MPIIFaceGaze/Label', type=str) | |
| # Important args ------------------------------------------------------------------------------------------------------- | |
| # ---------------------------------------------------------------------------------------------------------------------- | |
| parser.add_argument( | |
| '--dataset', dest='dataset', help='gaze360, mpiigaze', | |
| default= "gaze360", type=str) | |
| parser.add_argument( | |
| '--snapshot', dest='snapshot', help='Path to the folder contains models.', | |
| default='output/snapshots/L2CS-gaze360-_loader-180-4-lr', type=str) | |
| parser.add_argument( | |
| '--evalpath', dest='evalpath', help='path for the output evaluating gaze test.', | |
| default="evaluation/L2CS-gaze360-_loader-180-4-lr", type=str) | |
| parser.add_argument( | |
| '--gpu',dest='gpu_id', help='GPU device id to use [0]', | |
| default="0", type=str) | |
| parser.add_argument( | |
| '--batch_size', dest='batch_size', help='Batch size.', | |
| default=100, type=int) | |
| parser.add_argument( | |
| '--arch', dest='arch', help='Network architecture, can be: ResNet18, ResNet34, [ResNet50], ''ResNet101, ResNet152, Squeezenet_1_0, Squeezenet_1_1, MobileNetV2', | |
| default='ResNet50', type=str) | |
| # --------------------------------------------------------------------------------------------------------------------- | |
| # Important args ------------------------------------------------------------------------------------------------------ | |
| args = parser.parse_args() | |
| return args | |
| def getArch(arch,bins): | |
| # Base network structure | |
| if arch == 'ResNet18': | |
| model = L2CS( torchvision.models.resnet.BasicBlock,[2, 2, 2, 2], bins) | |
| elif arch == 'ResNet34': | |
| model = L2CS( torchvision.models.resnet.BasicBlock,[3, 4, 6, 3], bins) | |
| elif arch == 'ResNet101': | |
| model = L2CS( torchvision.models.resnet.Bottleneck,[3, 4, 23, 3], bins) | |
| elif arch == 'ResNet152': | |
| model = L2CS( torchvision.models.resnet.Bottleneck,[3, 8, 36, 3], bins) | |
| else: | |
| if arch != 'ResNet50': | |
| print('Invalid value for architecture is passed! ' | |
| 'The default value of ResNet50 will be used instead!') | |
| model = L2CS( torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], bins) | |
| return model | |
| if __name__ == '__main__': | |
| args = parse_args() | |
| cudnn.enabled = True | |
| gpu = select_device(args.gpu_id, batch_size=args.batch_size) | |
| batch_size=args.batch_size | |
| arch=args.arch | |
| data_set=args.dataset | |
| evalpath =args.evalpath | |
| snapshot_path = args.snapshot | |
| bins=args.bins | |
| angle=args.angle | |
| bin_width=args.bin_width | |
| transformations = transforms.Compose([ | |
| transforms.Resize(448), | |
| transforms.ToTensor(), | |
| transforms.Normalize( | |
| mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225] | |
| ) | |
| ]) | |
| if data_set=="gaze360": | |
| gaze_dataset=Gaze360(args.gaze360label_dir,args.gaze360image_dir, transformations, 180, 4, train=False) | |
| test_loader = torch.utils.data.DataLoader( | |
| dataset=gaze_dataset, | |
| batch_size=batch_size, | |
| shuffle=False, | |
| num_workers=4, | |
| pin_memory=True) | |
| if not os.path.exists(evalpath): | |
| os.makedirs(evalpath) | |
| # list all epochs for testing | |
| folder = os.listdir(snapshot_path) | |
| folder.sort(key=natural_keys) | |
| softmax = nn.Softmax(dim=1) | |
| with open(os.path.join(evalpath,data_set+".log"), 'w') as outfile: | |
| configuration = f"\ntest configuration = gpu_id={gpu}, batch_size={batch_size}, model_arch={arch}\nStart testing dataset={data_set}----------------------------------------\n" | |
| print(configuration) | |
| outfile.write(configuration) | |
| epoch_list=[] | |
| avg_yaw=[] | |
| avg_pitch=[] | |
| avg_MAE=[] | |
| for epochs in folder: | |
| # Base network structure | |
| model=getArch(arch, 90) | |
| saved_state_dict = torch.load(os.path.join(snapshot_path, epochs)) | |
| model.load_state_dict(saved_state_dict) | |
| model.cuda(gpu) | |
| model.eval() | |
| total = 0 | |
| idx_tensor = [idx for idx in range(90)] | |
| idx_tensor = torch.FloatTensor(idx_tensor).cuda(gpu) | |
| avg_error = .0 | |
| with torch.no_grad(): | |
| for j, (images, labels, cont_labels, name) in enumerate(test_loader): | |
| images = Variable(images).cuda(gpu) | |
| total += cont_labels.size(0) | |
| label_pitch = cont_labels[:,0].float()*np.pi/180 | |
| label_yaw = cont_labels[:,1].float()*np.pi/180 | |
| gaze_pitch, gaze_yaw = model(images) | |
| # Binned predictions | |
| _, pitch_bpred = torch.max(gaze_pitch.data, 1) | |
| _, yaw_bpred = torch.max(gaze_yaw.data, 1) | |
| # Continuous predictions | |
| pitch_predicted = softmax(gaze_pitch) | |
| yaw_predicted = softmax(gaze_yaw) | |
| # mapping from binned (0 to 28) to angels (-180 to 180) | |
| pitch_predicted = torch.sum(pitch_predicted * idx_tensor, 1).cpu() * 4 - 180 | |
| yaw_predicted = torch.sum(yaw_predicted * idx_tensor, 1).cpu() * 4 - 180 | |
| pitch_predicted = pitch_predicted*np.pi/180 | |
| yaw_predicted = yaw_predicted*np.pi/180 | |
| for p,y,pl,yl in zip(pitch_predicted,yaw_predicted,label_pitch,label_yaw): | |
| avg_error += angular(gazeto3d([p,y]), gazeto3d([pl,yl])) | |
| x = ''.join(filter(lambda i: i.isdigit(), epochs)) | |
| epoch_list.append(x) | |
| avg_MAE.append(avg_error/total) | |
| loger = f"[{epochs}---{args.dataset}] Total Num:{total},MAE:{avg_error/total}\n" | |
| outfile.write(loger) | |
| print(loger) | |
| fig = plt.figure(figsize=(14, 8)) | |
| plt.xlabel('epoch') | |
| plt.ylabel('avg') | |
| plt.title('Gaze angular error') | |
| plt.legend() | |
| plt.plot(epoch_list, avg_MAE, color='k', label='mae') | |
| fig.savefig(os.path.join(evalpath,data_set+".png"), format='png') | |
| plt.show() | |
| elif data_set=="mpiigaze": | |
| model_used=getArch(arch, bins) | |
| for fold in range(15): | |
| folder = os.listdir(args.gazeMpiilabel_dir) | |
| folder.sort() | |
| testlabelpathombined = [os.path.join(args.gazeMpiilabel_dir, j) for j in folder] | |
| gaze_dataset=Mpiigaze(testlabelpathombined,args.gazeMpiimage_dir, transformations, False, angle, fold) | |
| test_loader = torch.utils.data.DataLoader( | |
| dataset=gaze_dataset, | |
| batch_size=batch_size, | |
| shuffle=True, | |
| num_workers=4, | |
| pin_memory=True) | |
| if not os.path.exists(os.path.join(evalpath, f"fold"+str(fold))): | |
| os.makedirs(os.path.join(evalpath, f"fold"+str(fold))) | |
| # list all epochs for testing | |
| folder = os.listdir(os.path.join(snapshot_path,"fold"+str(fold))) | |
| folder.sort(key=natural_keys) | |
| softmax = nn.Softmax(dim=1) | |
| with open(os.path.join(evalpath, os.path.join("fold"+str(fold), data_set+".log")), 'w') as outfile: | |
| configuration = f"\ntest configuration equal gpu_id={gpu}, batch_size={batch_size}, model_arch={arch}\nStart testing dataset={data_set}, fold={fold}---------------------------------------\n" | |
| print(configuration) | |
| outfile.write(configuration) | |
| epoch_list=[] | |
| avg_MAE=[] | |
| for epochs in folder: | |
| model=model_used | |
| saved_state_dict = torch.load(os.path.join(snapshot_path+"/fold"+str(fold),epochs)) | |
| model= nn.DataParallel(model,device_ids=[0]) | |
| model.load_state_dict(saved_state_dict) | |
| model.cuda(gpu) | |
| model.eval() | |
| total = 0 | |
| idx_tensor = [idx for idx in range(28)] | |
| idx_tensor = torch.FloatTensor(idx_tensor).cuda(gpu) | |
| avg_error = .0 | |
| with torch.no_grad(): | |
| for j, (images, labels, cont_labels, name) in enumerate(test_loader): | |
| images = Variable(images).cuda(gpu) | |
| total += cont_labels.size(0) | |
| label_pitch = cont_labels[:,0].float()*np.pi/180 | |
| label_yaw = cont_labels[:,1].float()*np.pi/180 | |
| gaze_pitch, gaze_yaw = model(images) | |
| # Binned predictions | |
| _, pitch_bpred = torch.max(gaze_pitch.data, 1) | |
| _, yaw_bpred = torch.max(gaze_yaw.data, 1) | |
| # Continuous predictions | |
| pitch_predicted = softmax(gaze_pitch) | |
| yaw_predicted = softmax(gaze_yaw) | |
| # mapping from binned (0 to 28) to angels (-42 to 42) | |
| pitch_predicted = \ | |
| torch.sum(pitch_predicted * idx_tensor, 1).cpu() * 3 - 42 | |
| yaw_predicted = \ | |
| torch.sum(yaw_predicted * idx_tensor, 1).cpu() * 3 - 42 | |
| pitch_predicted = pitch_predicted*np.pi/180 | |
| yaw_predicted = yaw_predicted*np.pi/180 | |
| for p,y,pl,yl in zip(pitch_predicted, yaw_predicted, label_pitch, label_yaw): | |
| avg_error += angular(gazeto3d([p,y]), gazeto3d([pl,yl])) | |
| x = ''.join(filter(lambda i: i.isdigit(), epochs)) | |
| epoch_list.append(x) | |
| avg_MAE.append(avg_error/ total) | |
| loger = f"[{epochs}---{args.dataset}] Total Num:{total},MAE:{avg_error/total} \n" | |
| outfile.write(loger) | |
| print(loger) | |
| fig = plt.figure(figsize=(14, 8)) | |
| plt.xlabel('epoch') | |
| plt.ylabel('avg') | |
| plt.title('Gaze angular error') | |
| plt.legend() | |
| plt.plot(epoch_list, avg_MAE, color='k', label='mae') | |
| fig.savefig(os.path.join(evalpath, os.path.join("fold"+str(fold), data_set+".png")), format='png') | |
| # plt.show() | |