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()