|
- import numpy as np
- import paddle_lpips as lpips
- from data import data_loader as dl
- import argparse
- from IPython import embed
-
- parser = argparse.ArgumentParser()
- parser.add_argument('--dataset_mode', type=str, default='2afc', help='[2afc,jnd]')
- parser.add_argument('--datasets', type=str, nargs='+', default=['val/traditional','val/cnn','val/superres','val/deblur','val/color','val/frameinterp'], help='datasets to test - for jnd mode: [val/traditional],[val/cnn]; for 2afc mode: [train/traditional],[train/cnn],[train/mix],[val/traditional],[val/cnn],[val/color],[val/deblur],[val/frameinterp],[val/superres]')
- parser.add_argument('--model', type=str, default='lpips', help='distance model type [lpips] for linearly calibrated net, [baseline] for off-the-shelf network, [l2] for euclidean distance, [ssim] for Structured Similarity Image Metric')
- parser.add_argument('--net', type=str, default='alex', help='[squeeze], [alex], or [vgg] for network architectures')
- parser.add_argument('--colorspace', type=str, default='Lab', help='[Lab] or [RGB] for colorspace to use for l2, ssim model types')
- parser.add_argument('--batch_size', type=int, default=50, help='batch size to test image patches in')
- parser.add_argument('--gpu_ids', type=int, nargs='+', default=[0], help='gpus to use')
- parser.add_argument('--nThreads', type=int, default=4, help='number of threads to use in data loader')
-
- parser.add_argument('--model_path', type=str, default=None, help='location of model, will default to ./weights/v[version]/[net_name].pth')
-
- parser.add_argument('--from_scratch', action='store_true', help='model was initialized from scratch')
- parser.add_argument('--train_trunk', action='store_true', help='model trunk was trained/tuned')
- parser.add_argument('--version', type=str, default='0.1', help='v0.1 is latest, v0.0 was original release')
-
- opt = parser.parse_args()
- if(opt.model in ['l2','ssim']):
- opt.batch_size = 1
-
- # initialize model
- trainer = lpips.Trainer()
- # trainer.initialize(model=opt.model,net=opt.net,colorspace=opt.colorspace,model_path=opt.model_path)
- trainer.initialize(model=opt.model, net=opt.net, colorspace=opt.colorspace,
- model_path=opt.model_path, pnet_rand=opt.from_scratch, pnet_tune=opt.train_trunk,
- version=opt.version, gpu_ids=opt.gpu_ids)
-
- if(opt.model in ['net-lin','net']):
- print('Testing model [%s]-[%s]'%(opt.model,opt.net))
- elif(opt.model in ['l2','ssim']):
- print('Testing model [%s]-[%s]'%(opt.model,opt.colorspace))
-
- # initialize data loader
- for dataset in opt.datasets:
- data_loader = dl.CreateDataLoader(dataset,dataset_mode=opt.dataset_mode, batch_size=opt.batch_size, nThreads=opt.nThreads)
-
- # evaluate model on data
- if(opt.dataset_mode=='2afc'):
- (score, results_verbose) = lpips.score_2afc_dataset(data_loader, trainer.forward, name=dataset)
- elif(opt.dataset_mode=='jnd'):
- (score, results_verbose) = lpips.score_jnd_dataset(data_loader, trainer.forward, name=dataset)
-
- # print results
- print(' Dataset [%s]: %.2f'%(dataset,100.*score))
|