import PIL from PIL import Image import numpy as np from os import listdir, mkdir from os.path import isfile, join, isdir from tqdm import tqdm import glob import skimage from utils import _load_img_array, _rgb2ycbcr, PSNR, im2tensor ### Quality mesuare ### # LPIPS import LPIPS.models.dist_model as dm model_LPIPS = dm.DistModel() model_LPIPS.initialize(model='net-lin', net='alex', use_gpu=True) ### Generator ### ## RRDB import RRDBNet_arch as arch model_RRDB = arch.RRDBNet(3, 3, 64, 23, gc=32).cuda() lm = torch.load('./models/RRDB_BlurBicubicGAN.pth') model_RRDB.load_state_dict(lm.state_dict(), strict=True) ## Create output folder if not isdir('./output'): mkdir('./output') if not isdir('./output/Bic'): mkdir('./output/Bic')
dataset = Dataset(config['data_root'], config, split='val') dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, drop_last=True, shuffle=False, num_workers=int( config['num_workers'])) with torch.no_grad(): # Start testing # load the LPIPS model from LPIPS.models import dist_model as dm dis_model = dm.DistModel() dis_model.initialize(model='net-lin', net='alex', use_gpu=True) img_collect_a2b = [] img_collect_b2a = [] dis_array_a2b = np.zeros(__TRY_NUM) dis_array_b2a = np.zeros(__TRY_NUM) for j in trange(__TRY_NUM, desc='try_num'): dis_accu_a2b = 0 dis_accu_b2a = 0 for _ in trange(__EVAL_NUM, desc='eval_num'): input_a = (dataset.sample_a().cuda()[0]).unsqueeze(dim=0) input_b = (dataset.sample_b().cuda()[0]).unsqueeze(dim=0)
def LPIPS_eval(self, data_loader, eval_num=100, eval_name='LPIPS_eval', try_num=20): """ LPIPS evaluation function Parameters: eval_num: Total number of samples needed to be evaluated eval_name: Evaluation name for different kinds of evaluation try_num: How many times to run the evaluation on the whole """ import warnings warnings.filterwarnings("ignore", category=UserWarning) netG, _ = load_networks(self.model_dir) netG.eval() pbar = tqdm(total=try_num) # load the LPIPS feature extract network model from LPIPS.models import dist_model as dm dis_model = dm.DistModel() dis_model.initialize(model='net-lin', net='alex', use_gpu=True) print("\033[1;31m Doing Evaluation, Name: **%s**. \033[0m" % eval_name) img_collect = [] dis_array = np.zeros(try_num) for j in range(try_num): dis_accu = 0 pbar.update(1) for i, data in enumerate(data_loader): if i > eval_num: break data = data[:-1] # exclude the filename if data[2].dim() > 2: data[2] = get_edges(data[2]) real_imgs, mask_tensor, edge_tensor = var_and_cuda(data) if edge_tensor.dim() == 2: edge_tensor = None inputs = (mask_tensor, edge_tensor, None) for _ in range(cfg.SAMPLE_NUM): while len(img_collect) != 2: img_collect.append(netG(*inputs)[0].data) d = dis_model.forward(*img_collect) dis_accu += d.sum() img_collect = [] dis_array[j] = dis_accu / (cfg.SAMPLE_NUM * eval_num) pbar.close() output_str = '''For total %d pairs of random samples from %d condition, trial num %d, the LPIPS score is %.4f, std %.4f .\n''' % (cfg.SAMPLE_NUM * eval_num, eval_num, try_num, dis_array.mean(), dis_array.std()) print('''\033[1;31m %s \033[0m''' % output_str) with open(os.path.join(self.output_dir, 'LPIPS_scores.txt'), 'a') as f: date_str = datetime.datetime.now().strftime('%b-%d-%I%M%p-%G') f.write('%s, %s' % (date_str, output_str)) return dis_array.mean()