def __init__(
self,
model='net-lin',
net='alex',
colorspace='rgb',
spatial=False,
use_gpu=True,
gpu_ids=[
0
]): # VGG using our perceptually-learned weights (LPIPS metric)
# def __init__(self, model='net', net='vgg', use_gpu=True): # "default" way of using VGG as a perceptual loss
super(PerceptualLoss, self).__init__()
print('Setting up Perceptual loss...')
self.use_gpu = use_gpu
self.spatial = spatial
self.gpu_ids = gpu_ids
self.model = dist_model.DistModel()
self.model.initialize(model=model,
net=net,
use_gpu=use_gpu,
colorspace=colorspace,
spatial=self.spatial,
gpu_ids=gpu_ids)
print('...[%s] initialized' % self.model.name())
print('...Done')
def init_lpips_eval():
LPIPS_dir = "../PerceptualSimilarity"
LPIPS_net = "squeeze"
sys.path.append(LPIPS_dir)
from models import dist_model as dm
print("Initialize Distance model from %s" % LPIPS_net)
lpips_model = dm.DistModel()
lpips_model.initialize(model='net-lin',
net='squeeze',
use_gpu=True,
model_path=os.path.join(
LPIPS_dir, 'weights/v0.1/%s.pth' % LPIPS_net))
return lpips_model
def compute(image_file1, image_file2, gpu_flag): ## Initializing the model model = dm.DistModel() model.initialize(model='net-lin', net='squeeze', use_gpu=gpu_flag) # Load images img0 = util.im2tensor(util.load_image(image_file1)) # RGB image from [-1,1] img1 = util.im2tensor(util.load_image(image_file2)) # Compute distance dist01 = model.forward(img0, img1) return float(dist01[0])
def __init__(self):
mode = 'net-lin' # ['net', 'net-lin'] 'net-lin' for Linearly calibrated models, 'net' for Off-the-shelf uncalibrated networks
net = 'alex' # ['squeeze', 'alex', 'vgg']
use_gpu = True # if cuda.is_avaliable() else False
mode_low = 'l2' # ['l2', 'ssim']
colorspace = 'RGB' # ['Lab', 'RGB']
## modify configuration
self.folder_root = '../Results'
self.dataset = 'Set5'
self.methods = {'FASRGAN', 'Fs-SRGAN', 'FA+Fs-SRGAN'}
self.model = dm.DistModel()
self.model.initialize(model=mode,
net=net,
use_gpu=use_gpu,
colorspace=colorspace)
import argparse
from models import dist_model as dm
from util import util
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--path0', type=str, default='./imgs/ex_ref.png')
parser.add_argument('--path1', type=str, default='./imgs/ex_p0.png')
parser.add_argument('--use_gpu', action='store_true', help='turn on flag to use GPU')
opt = parser.parse_args()
## Initializing the model
print("Initializing the model")
model = dm.DistModel()
model.initialize(model='net-lin',net='squeeze',use_gpu=opt.use_gpu)
# Load images
print("Load images")
img0 = util.im2tensor(util.load_image(opt.path0)) # RGB image from [-1,1]
img1 = util.im2tensor(util.load_image(opt.path1))
# Compute distance
print("Compute distance")
dist01 = model.forward(img0,img1)
print('Distance: %.3f'%dist01)
def __init__(self, opt):
BaseModel.__init__(self, opt)
# specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
self.loss_names = [
'D_A', 'G_A', 'cycle_A', 'idt_A', 'D_B', 'G_B', 'cycle_B', 'idt_B'
]
# specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
visual_names_A = ['real_A', 'fake_B', 'rec_A']
visual_names_B = ['real_B', 'fake_A', 'rec_B']
if self.isTrain and self.opt.lambda_identity > 0.0: # if identity loss is used, we also visualize idt_B=G_A(B) ad idt_A=G_A(B)
visual_names_A.append('idt_B')
visual_names_B.append('idt_A')
if self.isTrain:
visual_names_A.append('real_A_hed')
visual_names_A.append('rec_A_hed')
if self.isTrain and self.opt.ntrunc_trunc:
visual_names_A.append('rec_At')
visual_names_A.append('rec_At_hed')
self.loss_names = [
'D_A', 'G_A', 'cycle_A', 'cycle_A2', 'idt_A', 'D_B', 'G_B',
'cycle_B', 'idt_B', 'G'
]
if self.isTrain and self.opt.use_mask:
visual_names_A.append('fake_B_l')
visual_names_A.append('real_B_l')
self.loss_names += ['D_A_l', 'G_A_l']
if self.isTrain and self.opt.use_eye_mask:
visual_names_A.append('fake_B_le')
visual_names_A.append('real_B_le')
self.loss_names += ['D_A_le', 'G_A_le']
if self.isTrain and self.opt.use_lip_mask:
visual_names_A.append('fake_B_ll')
visual_names_A.append('real_B_ll')
self.loss_names += ['D_A_ll', 'G_A_ll']
if not self.isTrain and self.opt.use_mask:
visual_names_A.append('fake_B_l')
visual_names_A.append('real_B_l')
if not self.isTrain and self.opt.use_eye_mask:
visual_names_A.append('fake_B_le')
visual_names_A.append('real_B_le')
if not self.isTrain and self.opt.use_lip_mask:
visual_names_A.append('fake_B_ll')
visual_names_A.append('real_B_ll')
self.loss_names += ['D_A_cls', 'G_A_cls']
self.visual_names = visual_names_A + visual_names_B # combine visualizations for A and B
print(self.visual_names)
# specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>.
if self.isTrain:
self.model_names = ['G_A', 'G_B', 'D_A', 'D_B']
if self.opt.use_mask:
self.model_names += ['D_A_l']
if self.opt.use_eye_mask:
self.model_names += ['D_A_le']
if self.opt.use_lip_mask:
self.model_names += ['D_A_ll']
else: # during test time, only load Gs
self.model_names = ['G_A', 'G_B']
# define networks (both Generators and discriminators)
# The naming is different from those used in the paper.
# Code (vs. paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)
if not self.opt.style_control:
self.netG_A = networks.define_G(opt.input_nc, opt.output_nc,
opt.ngf, opt.netG, opt.norm,
not opt.no_dropout, opt.init_type,
opt.init_gain, self.gpu_ids)
else:
print(opt.netga)
print('model0_res', opt.model0_res)
print('model1_res', opt.model1_res)
self.netG_A = networks.define_G(opt.input_nc, opt.output_nc,
opt.ngf, opt.netga, opt.norm,
not opt.no_dropout, opt.init_type,
opt.init_gain, self.gpu_ids,
opt.model0_res, opt.model1_res)
self.netG_B = networks.define_G(opt.output_nc, opt.input_nc, opt.ngf,
opt.netG, opt.norm, not opt.no_dropout,
opt.init_type, opt.init_gain,
self.gpu_ids)
if self.isTrain: # define discriminators
self.netD_A = networks.define_D(opt.output_nc,
opt.ndf,
opt.netda,
opt.n_layers_D,
opt.norm,
opt.init_type,
opt.init_gain,
self.gpu_ids,
n_class=3)
self.netD_B = networks.define_D(opt.input_nc, opt.ndf, opt.netD,
opt.n_layers_D, opt.norm,
opt.init_type, opt.init_gain,
self.gpu_ids)
if self.opt.use_mask:
if self.opt.mask_type in [2, 3]:
output_nc = opt.output_nc + 1
else:
output_nc = opt.output_nc
self.netD_A_l = networks.define_D(output_nc, opt.ndf, opt.netD,
opt.n_layers_D, opt.norm,
opt.init_type, opt.init_gain,
self.gpu_ids)
if self.opt.use_eye_mask:
if self.opt.mask_type in [2, 3]:
output_nc = opt.output_nc + 1
else:
output_nc = opt.output_nc
self.netD_A_le = networks.define_D(output_nc, opt.ndf,
opt.netD, opt.n_layers_D,
opt.norm, opt.init_type,
opt.init_gain, self.gpu_ids)
if self.opt.use_lip_mask:
if self.opt.mask_type in [2, 3]:
output_nc = opt.output_nc + 1
else:
output_nc = opt.output_nc
self.netD_A_ll = networks.define_D(output_nc, opt.ndf,
opt.netD, opt.n_layers_D,
opt.norm, opt.init_type,
opt.init_gain, self.gpu_ids)
if not self.isTrain:
self.criterionGAN = networks.GANLoss('lsgan').to(self.device)
if self.isTrain:
if opt.lambda_identity > 0.0: # only works when input and output images have the same number of channels
assert (opt.input_nc == opt.output_nc)
self.fake_A_pool = ImagePool(
opt.pool_size
) # create image buffer to store previously generated images
self.fake_B_pool = ImagePool(
opt.pool_size
) # create image buffer to store previously generated images
# define loss functions
self.criterionGAN = networks.GANLoss(opt.gan_mode).to(
self.device) # define GAN loss.
self.criterionCycle = torch.nn.L1Loss()
self.criterionIdt = torch.nn.L1Loss()
self.criterionCls = torch.nn.CrossEntropyLoss()
self.criterionCls2 = torch.nn.CrossEntropyLoss(reduction='none')
# initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
self.optimizer_G = torch.optim.Adam(itertools.chain(
self.netG_A.parameters(), self.netG_B.parameters()),
lr=opt.lr,
betas=(opt.beta1, 0.999))
if not self.opt.use_mask:
self.optimizer_D = torch.optim.Adam(itertools.chain(
self.netD_A.parameters(), self.netD_B.parameters()),
lr=opt.lr,
betas=(opt.beta1, 0.999))
elif not self.opt.use_eye_mask:
D_params = list(self.netD_A.parameters()) + list(
self.netD_B.parameters()) + list(
self.netD_A_l.parameters())
self.optimizer_D = torch.optim.Adam(D_params,
lr=opt.lr,
betas=(opt.beta1, 0.999))
elif not self.opt.use_lip_mask:
D_params = list(self.netD_A.parameters()) + list(
self.netD_B.parameters()) + list(
self.netD_A_l.parameters()) + list(
self.netD_A_le.parameters())
self.optimizer_D = torch.optim.Adam(D_params,
lr=opt.lr,
betas=(opt.beta1, 0.999))
else:
D_params = list(self.netD_A.parameters()) + list(
self.netD_B.parameters()) + list(
self.netD_A_l.parameters()) + list(
self.netD_A_le.parameters()) + list(
self.netD_A_ll.parameters())
self.optimizer_D = torch.optim.Adam(D_params,
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
self.optimizers.append(self.optimizer_D)
self.lpips = dm.DistModel(opt,
model='net-lin',
net='alex',
use_gpu=True)
self.set_requires_grad(self.lpips, False)
self.hed = networks.define_HED(
init_weights_=opt.hed_pretrained_mode,
gpu_ids_=self.opt.gpu_ids_p)
self.set_requires_grad(self.hed, False)
def compute_metric(dir1, dir2, mode, mask=None, thre=0):
if mode == 'perceptual':
from models import dist_model as dm
import torch
from util import util
global model
if model is None:
cwd = os.getcwd()
os.chdir('../../PerceptualSimilarity')
model = dm.DistModel()
#model.initialize(model='net-lin',net='alex',use_gpu=True, spatial=True)
model.initialize(model='net-lin', net='alex', use_gpu=True)
print('Model [%s] initialized' % model.name())
os.chdir(cwd)
if mode.startswith('perceptual_tf'):
sys.path += ['../../lpips-tensorflow']
import lpips_tf
import tensorflow as tf
image0_ph = tf.placeholder(tf.float32, [1, None, None, 3])
image1_ph = tf.placeholder(tf.float32, [1, None, None, 3])
if mode == 'perceptual_tf':
distance_t = lpips_tf.lpips(image0_ph,
image1_ph,
model='net-lin',
net='alex')
elif mode == 'perceptual_tf_vgg':
distance_t = lpips_tf.lpips(image0_ph,
image1_ph,
model='net-lin',
net='vgg')
else:
raise
sess = tf.Session()
if mode == 'l2_with_gradient':
import demo
import tensorflow as tf
output = tf.placeholder(tf.float32, shape=[1, None, None, None])
gradient = demo.image_gradients(output)
sess = tf.Session()
files1 = os.listdir(dir1)
files2 = os.listdir(dir2)
img_files1 = sorted([
file for file in files1
if file.endswith('.png') or file.endswith('.jpg')
])
img_files2 = sorted([
file for file in files2
if file.endswith('.png') or file.endswith('.jpg')
])
if '--prefix' in sys.argv:
prefix_idx = sys.argv.index('--prefix')
prefix = sys.argv[prefix_idx + 1]
img_files2 = [file for file in img_files2 if file.startswith(prefix)]
if mask is not None:
mask_files = []
for dir in sorted(mask):
files3 = os.listdir(dir)
add_mask_files = sorted([
os.path.join(dir, file) for file in files3
if file.startswith('mask')
])
if len(add_mask_files) == 0:
files_to_dilate = sorted([
file for file in files3
if file.startswith('g_intermediates')
])
for file in files_to_dilate:
mask_arr = numpy.load(os.path.join(dir, file))
mask_arr = numpy.squeeze(mask_arr)
mask_arr = mask_arr >= thre
dilated_mask = dilation(mask_arr, disk(10))
dilated_filename = 'mask_' + file
numpy.save(os.path.join(dir, dilated_filename),
dilated_mask.astype('f'))
add_mask_files.append(os.path.join(dir, dilated_filename))
mask_files += add_mask_files
print(mask_files)
assert len(mask_files) == len(img_files1)
skip_first_n = 0
if '--skip_first_n' in sys.argv:
try:
skip_first_n = int(sys.argv[sys.argv.index('--skip_first_n') + 1])
except:
skip_first_n = 0
skip_last_n = 0
if '--skip_last_n' in sys.argv:
try:
skip_last_n = int(sys.argv[sys.argv.index('--skip_last_n') + 1])
except:
skip_last_n = 0
img_files2 = img_files2[skip_first_n:]
if skip_last_n > 0:
img_files2 = img_files2[:-skip_last_n]
assert len(img_files1) == len(img_files2)
# locate GT gradient directory
if mode == 'l2_with_gradient':
head, tail = os.path.split(dir2)
gradient_gt_dir = os.path.join(head, tail[:-3] + 'grad')
if not os.path.exists(gradient_gt_dir):
printf("dir not found,", gradient_gt_dir)
raise
gradient_gt_files = os.listdir(gradient_gt_dir)
gradient_gt_files = sorted(
[file for file in gradient_gt_files if file.endswith('.npy')])
assert len(img_files1) == len(gradient_gt_files)
vals = numpy.empty(len(img_files1))
#if mode == 'perceptual':
# global model
for ind in range(len(img_files1)):
if mode == 'ssim' or mode == 'l2' or mode == 'l2_with_gradient':
img1 = skimage.img_as_float(
skimage.io.imread(os.path.join(dir1, img_files1[ind])))
img2 = skimage.img_as_float(
skimage.io.imread(os.path.join(dir2, img_files2[ind])))
if mode == 'ssim':
#vals[ind] = skimage.measure.compare_ssim(img1, img2, datarange=img2.max()-img2.min(), multichannel=True)
metric_val = skimage.measure.compare_ssim(
img1,
img2,
datarange=img2.max() - img2.min(),
multichannel=True)
else:
#vals[ind] = numpy.mean((img1 - img2) ** 2) * 255.0 * 255.0
metric_val = ((img1 - img2)**2) * 255.0 * 255.0
if mode == 'l2_with_gradient':
metric_val = numpy.mean(metric_val, axis=2)
gradient_gt = numpy.load(
os.path.join(gradient_gt_dir, gradient_gt_files[ind]))
dx, dy = sess.run(gradient,
feed_dict={
output:
numpy.expand_dims(img1[..., ::-1],
axis=0)
})
#is_edge = skimage.feature.canny(skimage.color.rgb2gray(img1))
dx_ground = gradient_gt[:, :, :, 1:4]
dy_ground = gradient_gt[:, :, :, 4:]
edge_ground = gradient_gt[:, :, :, 0]
gradient_loss_term = numpy.mean(
(dx - dx_ground)**2.0 + (dy - dy_ground)**2.0, axis=3)
metric_val += numpy.squeeze(
0.2 * 255.0 * 255.0 * gradient_loss_term * edge_ground *
edge_ground.size / numpy.sum(edge_ground))
#if mode == 'l2' and mask is not None:
# img_diff = (img1 - img2) ** 2.0
# mask_img = numpy.load(mask_files[ind])
# img_diff *= numpy.expand_dims(mask_img, axis=2)
# vals[ind] = (numpy.sum(img_diff) / numpy.sum(mask_img * 3)) * 255.0 * 255.0
elif mode == 'perceptual':
img1 = util.im2tensor(
util.load_image(os.path.join(dir1, img_files1[ind])))
img2 = util.im2tensor(
util.load_image(os.path.join(dir2, img_files2[ind])))
#vals[ind] = numpy.mean(model.forward(img1, img2)[0])
metric_val = numpy.expand_dims(model.forward(img1, img2), axis=2)
elif mode.startswith('perceptual_tf'):
img1 = np.expand_dims(skimage.img_as_float(
skimage.io.imread(os.path.join(dir1, img_files1[ind]))),
axis=0)
img2 = np.expand_dims(skimage.img_as_float(
skimage.io.imread(os.path.join(dir2, img_files2[ind]))),
axis=0)
metric_val = sess.run(distance_t,
feed_dict={
image0_ph: img1,
image1_ph: img2
})
else:
raise
if mask is not None:
assert mode in ['l2', 'perceptual']
mask_img = numpy.load(mask_files[ind])
metric_val *= numpy.expand_dims(mask_img, axis=2)
vals[ind] = numpy.sum(metric_val) / (numpy.sum(mask_img) *
metric_val.shape[2])
else:
vals[ind] = numpy.mean(metric_val)
mode = mode + ('_mask' if mask is not None else '')
filename_all = mode + '_all.txt'
filename_breakdown = mode + '_breakdown.txt'
filename_single = mode + '.txt'
numpy.savetxt(os.path.join(dir1, filename_all), vals, fmt="%f, ")
target = open(os.path.join(dir1, filename_single), 'w')
target.write("%f" % numpy.mean(vals))
target.close()
if len(img_files1) == 30:
target = open(os.path.join(dir1, filename_breakdown), 'w')
target.write("%f, %f, %f" % (numpy.mean(
vals[:5]), numpy.mean(vals[5:10]), numpy.mean(vals[10:])))
target.close()
if mode in ['l2_with_gradient', 'perceptual_tf']:
sess.close()
return vals
def __init__(self, model='net-lin', net='vgg', use_gpu=True): # VGG using our perceptually-learned weights (LPIPS metric)
# def __init__(self, model='net', net='vgg', use_gpu=True): # "default" way of using VGG
print('Setting up Perceptual loss..')
self.model = dist_model.DistModel()
self.model.initialize(model=model, net=net, use_gpu=True)
print('Done')
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--model',
choices=['net-lin', 'net'],
default='net-lin',
help='net-lin or net')
parser.add_argument('--net',
choices=['squeeze', 'alex', 'vgg'],
default='alex',
help='squeeze, alex, or vgg')
parser.add_argument('--version', type=str, default='0.1')
parser.add_argument('--image_height', type=int, default=64)
parser.add_argument('--image_width', type=int, default=64)
args = parser.parse_args()
model = dm.DistModel()
model.initialize(model=args.model,
net=args.net,
use_gpu=False,
version=args.version)
print('Model [%s] initialized' % model.name())
dummy_im0 = torch.Tensor(
1, 3, args.image_height,
args.image_width) # image should be RGB, normalized to [-1, 1]
dummy_im1 = torch.Tensor(1, 3, args.image_height, args.image_width)
cache_dir = os.path.expanduser('~/.lpips')
os.makedirs(cache_dir, exist_ok=True)
onnx_fname = os.path.join(
cache_dir, '%s_%s_v%s.onnx' % (args.model, args.net, args.version))
# export model to onnx format
torch.onnx.export(model.net, (dummy_im0, dummy_im1),
onnx_fname,
verbose=True)
# load and change dimensions to be dynamic
model = onnx.load(onnx_fname)
for dim in (0, 2, 3):
model.graph.input[0].type.tensor_type.shape.dim[dim].dim_param = '?'
model.graph.input[1].type.tensor_type.shape.dim[dim].dim_param = '?'
# needs to be imported after all the pytorch stuff, otherwise this causes a segfault
from onnx_tf.backend import prepare
tf_rep = prepare(model, device='CPU')
producer_version = tf_rep.graph.graph_def_versions.producer
pb_fname = os.path.join(
cache_dir, '%s_%s_v%s_%d.pb' %
(args.model, args.net, args.version, producer_version))
tf_rep.export_graph(pb_fname)
input0_name, input1_name = [
tf_rep.tensor_dict[input_name].name for input_name in tf_rep.inputs
]
(output_name, ) = [
tf_rep.tensor_dict[output_name].name for output_name in tf_rep.outputs
]
# ensure these are the names of the 2 inputs, since that will be assumed when loading the pb file
assert input0_name == '0:0'
assert input1_name == '1:0'
# ensure that the only output is the output of the last op in the graph, since that will be assumed later
(last_output_name, ) = [
output.name for output in tf_rep.graph.get_operations()[-1].outputs
]
assert output_name == last_output_name
def __init__(self, model='net-lin', net='alex', use_cuda=True):
super(PerceptualSimilarityLoss, self).__init__()
self.model = dm.DistModel()
self.model.initialize(model=model, net=net, use_gpu=use_cuda)
# Deeploss competitor 2: lin tuned squeeze
m = dm.DistModel()
path=os.path.join('./checkpoints/', 'rgb_pnet_lin_squeeze_trial0', 'latest_net_.pth')
m.initialize(model='net-lin', net='squeeze', model_path=path, colorspace='RGB',use_gpu=is_cuda, batch_size=64)
m.model_name = 'Deeploss-squeeze'
m.batch_size = m.batch_size // 4
models.append(m)
"""
# # Adaptive loss, as per reviewer request
# In[13]:
m = dm.DistModel()
path = os.path.join('./checkpoints/', 'gray_adaptive_trial0',
'latest_net_.pth')
m.initialize(model='adaptive', colorspace='Gray', use_gpu=is_cuda)
m.net.load_state_dict(torch.load(path, map_location=device))
m.model_name = 'Adaptive [LA]'
models.append(m)
m = dm.DistModel()
path = os.path.join('./checkpoints/', 'rgb_adaptive_trial0', 'latest_net_.pth')
m.initialize(model='adaptive', colorspace='RGB', use_gpu=is_cuda)
m.net.load_state_dict(torch.load(path, map_location=device))
m.model_name = 'Adaptive'
models.append(m)
# # Evaluate
def test_training_images2(model, test_loader, device):
#model = model.eval()
lpipsloss = dm.DistModel()
lpipsloss.initialize(model='net-lin',
net='alex',
use_gpu=True,
version='0.1')
mse_loss = torch.nn.MSELoss(size_average=None)
loss_dict = {
'mse': [],
'mse_avg': 0,
'psnr': [],
'psnr_avg': 0,
'lpips': [],
'lpips_avg': 0,
'lpips_center': [],
'lpips_center_avg': 0
}
with torch.no_grad():
for i_batch, sample_batched in enumerate(test_loader):
print('\r', 'running test images, image:', i_batch, end='')
inputs = sample_batched['image'].to(device)
labels = sample_batched['label'].to(device)
output, _ = model(inputs)
if not (np.any(output.cpu().detach().numpy() == -np.inf)):
mse_batch = mse_loss(output, labels)
lpips_batch = lpipsloss.forward_pair(output, labels)
c1 = 270 // 2
c2 = 480 // 2
sz = 75
lpips_center = lpipsloss.forward_pair(
output[:, c1 - sz:c1 + sz, c2 - sz:c2 + sz],
inputs[:, c1 - sz:c1 + sz, c2 - sz:c2 + sz])
psnr_batch = 20 * torch.log10(1 / torch.sqrt(mse_batch))
loss_dict['mse'].append(
mse_batch.cpu().detach().numpy().squeeze())
loss_dict['lpips'].append(
lpips_batch.cpu().detach().numpy().squeeze())
loss_dict['lpips_center'].append(
lpips_center.cpu().detach().numpy().squeeze())
#loss_dict['lpips'].append(lpips_batch)
loss_dict['psnr'].append(
psnr_batch.cpu().detach().numpy().squeeze())
if i_batch == 63:
loss_dict['sample_image'] = preplot(
output.detach().cpu().numpy()[0])
loss_dict['mse_avg'] = np.average(loss_dict['mse']).squeeze()
loss_dict['psnr_avg'] = np.average(loss_dict['psnr']).squeeze()
loss_dict['lpips_avg'] = np.average(loss_dict['lpips']).squeeze()
loss_dict['lpips_center_avg'] = np.average(
loss_dict['lpips_center']).squeeze()
print('\r', 'avg mse:', loss_dict['mse_avg'], 'avg psnr:',
loss_dict['psnr_avg'], 'avg lpips:', loss_dict['lpips_avg'],
'avg lpips center:', loss_dict['lpips_center_avg'])
return loss_dict
def test_training_images(model, model_admm, test_loader, device):
lpipsloss = dm.DistModel()
lpipsloss.initialize(model='net-lin',
net='alex',
use_gpu=True,
version='0.1')
mse_loss = torch.nn.MSELoss(size_average=None)
loss_dict = {
'mse': [],
'mse_avg': 0,
'psnr': [],
'psnr_avg': 0,
'lpips': [],
'lpips_avg': 0,
'data_loss': [],
'data_loss_avg': 0,
'lpips_center': [],
'lpips_center_avg': 0,
'mse_center': [],
'mse_center_avg': 0,
'sample_images': []
}
with torch.no_grad():
for i_batch, sample_batched in enumerate(test_loader):
print('\r', 'running test images, image:', i_batch, end='')
# Get input and label batch
inputs = sample_batched['image'].to(device)
labels = sample_batched['label'].to(device)
output = model(inputs)
if isinstance(output, tuple):
output = output[0]
# Check if image is bad
if not (np.any(output.cpu().detach().numpy() == -np.inf)):
mse_batch = mse_loss(output, labels) # MSE loss
lpips_batch = lpipsloss.forward_pair(output,
labels) # lpips loss
psnr_batch = 20 * torch.log10(
1 / torch.sqrt(mse_batch)) # PSNR
# Center region
c1 = 270 // 2
c2 = 480 // 2
sz = 75
lpips_center = lpipsloss.forward_pair(
output[:, :, c1 - sz:c1 + sz, c2 - sz:c2 + sz],
labels[:, :, c1 - sz:c1 + sz, c2 - sz:c2 + sz])
mse_center = mse_loss(
output[:, :, c1 - sz:c1 + sz, c2 - sz:c2 + sz],
labels[:, :, c1 - sz:c1 + sz, c2 - sz:c2 + sz])
# Data fidelity loss
input_image = normalize_image(inputs)
hfor = normalize_image(
Hfor(model_admm, pad_zeros_torch(model_admm, output)))
data_loss = torch.sum(
torch.norm(crop(model_admm, hfor) - input_image)**2)
loss_dict['mse'].append(
mse_batch.cpu().detach().numpy().squeeze())
loss_dict['lpips'].append(
lpips_batch.cpu().detach().numpy().squeeze())
loss_dict['psnr'].append(
psnr_batch.cpu().detach().numpy().squeeze())
loss_dict['data_loss'].append(
psnr_batch.cpu().detach().numpy().squeeze())
loss_dict['lpips_center'].append(
lpips_center.cpu().detach().numpy().squeeze())
loss_dict['mse_center'].append(
data_loss.cpu().detach().numpy().squeeze())
inds = [
63, 41, 88, 123, 134, 135, 151, 155, 160, 163, 178, 180,
187, 198, 202, 212, 224, 227, 239, 250, 253, 261, 271, 274,
281, 283, 396, 394, 392, 385, 376, 372, 366, 340, 336, 325,
324, 323, 400, 406, 419, 461, 502, 546, 549, 595, 641, 653,
693, 695, 712, 732, 738, 741, 757, 809, 984
]
if i_batch in inds:
loss_dict['sample_images'].append(
preplot(output.detach().cpu().numpy()[0]))
loss_dict['mse_avg'] = np.average(loss_dict['mse']).squeeze()
loss_dict['psnr_avg'] = np.average(loss_dict['psnr']).squeeze()
loss_dict['lpips_avg'] = np.average(loss_dict['lpips']).squeeze()
loss_dict['data_loss_avg'] = np.average(
loss_dict['data_loss']).squeeze()
loss_dict['lpips_center_avg'] = np.average(
loss_dict['lpips_center']).squeeze()
loss_dict['mse_center_avg'] = np.average(
loss_dict['mse_center']).squeeze()
print('\r', 'avg mse:', loss_dict['mse_avg'], 'avg psnr:',
loss_dict['psnr_avg'], 'avg lpips:', loss_dict['lpips_avg'],
'avg lpips center:', loss_dict['lpips_center_avg'])
return loss_dict