def call(self, img1, img2):
#psnr part
psnr = psnr_cal((img1.numpy()+1.)/2.,(img2.numpy()+1.)/2.)
#lpips part
loss_fn_alex = lpips.LPIPS(net='alex') # best forward scores
loss_fn_vgg = lpips.LPIPS(net='vgg') # closer to "traditional" perceptual loss, when used for optimization
if torch.cuda.is_available():
loss_fn_alex = loss_fn_alex.cuda()
loss_fn_vgg = loss_fn_vgg.cuda()
img1 = img1.cuda()
img2 = img2.cuda()
lpips_value_alex = loss_fn_alex(img1, img2)
lpips_value_vgg = loss_fn_vgg(img1, img2)
# img1_01 = (img1.numpy()+1.)/2.
# img2_01 = (img2.numpy()+1.)/2.
#ssim part
#if torch.cuda.is_available():
# img1 = img1.cuda()
# img2 = img2.cuda()
#print(pytorch_ssim.ssim(img1, img2))
ssim_loss = pytorch_ssim.SSIM(window_size = 11)
ssim_value = ssim_loss(img1, img2)
return lpips_value_alex,lpips_value_vgg, psnr, ssim_value
def __init__(self,device,metrics = 'ssim,msssim,mse,psnr') -> None:
self.metrics = metrics
self.device = device
if(('lpipsVGG' in metrics)):
self.lpipsVGG = lpips.LPIPS(net='vgg').to(device)
if(('lpipsAlex' in metrics)):
self.lpipsAlex = lpips.LPIPS(net='alex').to(device)
def setUp(self):
# initialize evaluation dataset/dataloader
torch.manual_seed(42)
stack = contextlib.ExitStack()
dataset_root, path_manager = stack.enter_context(get_skateboard_data())
self.addCleanup(stack.close)
category = "skateboard"
frame_file = os.path.join(dataset_root, category, "frame_annotations.jgz")
sequence_file = os.path.join(dataset_root, category, "sequence_annotations.jgz")
self.image_size = 256
self.dataset = ImplicitronDataset(
frame_annotations_file=frame_file,
sequence_annotations_file=sequence_file,
dataset_root=dataset_root,
image_height=self.image_size,
image_width=self.image_size,
box_crop=True,
path_manager=path_manager,
)
self.bg_color = 0.0
# init the lpips model for eval
provide_lpips_vgg()
self.lpips_model = lpips.LPIPS(net="vgg")
def __init__(self) -> None:
super(LPIPSLoss, self).__init__()
self.feature_extract = lpips.LPIPS(net="vgg", verbose=False).eval()
# Freeze model all parameters. Don't train.
for name, parameters in self.feature_extract.named_parameters():
parameters.requires_grad = False
def __init__(self,
model='net-lin',
net='vgg',
use_gpu=True,
spatial=False):
super(LPIPS_Loss, self).__init__()
self.model = lpips.LPIPS(net=net, spatial=spatial).eval()
def image_quality_evaluation(sr_filename: str, hr_filename: str, device: torch.device = "cpu"):
"""Image quality evaluation function.
Args:
sr_filename (str): Image file name after super resolution.
hr_filename (str): Original high resolution image file name.
device (optional, torch.device): Selection of data processing equipment in PyTorch. (Default: ``cpu``).
Returns:
If the `simple` variable is set to ``False`` return `mse, rmse, psnr, ssim, msssim, niqe, sam, vifp, lpips`,
else return `psnr, ssim`.
"""
# Reference sources from `https://github.com/richzhang/PerceptualSimilarity`
lpips_loss = lpips.LPIPS(net="vgg", verbose=False).to(device)
# Evaluate performance
sr = cv2.imread(sr_filename)
hr = cv2.imread(hr_filename)
# For LPIPS evaluation
sr_tensor = opencv2tensor(sr, device)
hr_tensor = opencv2tensor(hr, device)
# Complete estimate.
mse_value = mse(sr, hr)
rmse_value = rmse(sr, hr)
psnr_value = psnr(sr, hr)
ssim_value = ssim(sr, hr)
msssim_value = msssim(sr, hr)
niqe_value = niqe(sr_filename)
sam_value = sam(sr, hr)
vifp_value = vifp(sr, hr)
lpips_value = lpips_loss(sr_tensor, hr_tensor)
return mse_value, rmse_value, psnr_value, ssim_value, msssim_value, niqe_value, sam_value, vifp_value, lpips_value
def __init__(self, options):
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.options = options
self.test_ds = get_dataset_from_options(self.options, is_train=False)
self.model = get_model_from_options(self.options)
self.models_dict = {'model':self.model}
self.models_dict = {k:v.to(self.device)
for k,v in self.models_dict.items()}
self.optimizers_dict = {}
# Load the latest checkpoints
self.saver = CheckpointSaver(save_dir=options.checkpoint_dir)
self.checkpoint = None
if self.saver.exists_checkpoint():
self.checkpoint = self.saver.load_checkpoint(self.models_dict,
self.optimizers_dict,
checkpoint_file=self.options.checkpoint)
else:
print("No checkpoint found. The program will exit")
for model in self.models_dict:
self.models_dict[model].eval()
self.metrics = {'lpips': lpips.LPIPS(net='alex').cuda(),
'mse': torch.nn.MSELoss()
}
self.images_to_save = ['target_frames', 'pred', 'generator/object_pixels']
self.image_path = "out_images"
def __init__(self,
net='vgg',
eval_mode=True,
pooling=True,
normalize=False,
layer_res=None):
super(LPIPF, self).__init__()
# Using composition instead of inheritance because of path issues...
self.ps = lpips.LPIPS(net=net, eval_mode=eval_mode)
assert self.ps.lpips
assert pooling or not layer_res
self.pooling = pooling
self.normalize = normalize
if layer_res:
self.layer_res = {net: layer_res}
else:
self.layer_res = {
'vgg': [8, 4, 2, 2, 1],
# 'vgg': [16, 8, 4, 4, 2],
# 'vgg': [32, 16, 8, 8, 4],
# 'vgg': [64, 32, 16, 16, 8],
# 'vgg': [128, 64, 32, 16, 8], # No Pooling
'alex': [7, 3, 3, 1, 1],
# 'alex': [15, 7, 7, 7, 7],
# 'alex': [31, 15, 7, 7, 7], # No pooling
}
def __init__(self, loss_lambda=0.5):
super(LPIPSLoss_L1, self).__init__()
self.similarity = lpips.LPIPS(net='vgg').cuda()
self.l1_loss = nn.L1Loss()
# self.similarity = LPIPS(network='vgg').cuda()
# self.l2_loss = nn.MSELoss()
self.loss_lambda = loss_lambda
def __init__(self,
data_set: EvalDataSet,
ssim_psnr=True,
fft_root=None,
layer_map_root=None,
tf_backend=False,
mse=False,
mae=False,
lpips_flag=False,
sample_num=None):
"""
:param data_set:
:param fft_root: None or 'full path', None to skip
:param layer_map_root: same as fft_root
"""
self.data_set = data_set
self.fft_root = fft_root
self.layer_map_root = layer_map_root
self.tf_backend = tf_backend
self.ssim_psnr_flag = ssim_psnr
self.mse = mse
self.mae = mae
self.lpips_flag = lpips_flag
self.check_path(fft_root)
self.check_path(layer_map_root)
self.sample_num = len(data_set) if sample_num is None else sample_num
self.lpips_alex = lpips.LPIPS(net='alex').cuda()
def __init__(self, net, layers=None, use_dropout=True):
super(LayerLoss, self).__init__()
assert net in ['alex', 'vgg', 'squeeze']
self.net = lpips.LPIPS(pretrained=True,
net=net,
lpips=False,
use_dropout=use_dropout).net
self.layers = layers
def criterions(name):
if name == "mse":
return nn.MSELoss()
elif name == "l1":
return nn.L1Loss()
elif name == "lpips":
return lpips.LPIPS(net="vgg").cuda()
elif name == "ms_ssim":
return MS_SSIM(data_range=1.0)
def __init__(self, dtype, weight_pb=None):
"""
Args:
dtype (str or numpy.dtype): Data type, from which maximum allowed
will be derived.
weight_pb (str, optional): Path to the network weight protobuf.
Defaults to the bundled ``net-lin_alex_v0.1.pb``.
"""
self.lpips = lpips.LPIPS(net='alex')
def get_close_far_members(args):
with torch.no_grad():
lpips_fn = lpips.LPIPS(net='vgg').to(device)
preprocess = transforms.Compose([
transforms.Resize([256, 256]),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])
cluster_size = 50
base_path = os.path.join("cluster_centers", "%s" % (args.dataset),
"%s" % (args.baseline))
avg_dist = torch.zeros([
10,
])
for k in range(10):
curr_path = os.path.join(base_path, "c%d" % (k))
files_list = os.listdir(curr_path)
files_list.remove('center.png')
random.shuffle(files_list)
files_list = files_list[:cluster_size]
dists = []
min_dist, max_dist = 1000.0, 0.0
min_ind, max_ind = 0, 0
# center image
input1_path = os.path.join(curr_path, 'center.png')
input_image1 = Image.open(input1_path)
input_tensor1 = preprocess(input_image1)
input_tensor1 = input_tensor1.to(device)
for i in range(len(files_list)):
input2_path = os.path.join(curr_path, files_list[i])
input_image2 = Image.open(input2_path)
input_tensor2 = preprocess(input_image2)
input_tensor2 = input_tensor2.to(device)
dist = lpips_fn(input_tensor1, input_tensor2)
if dist <= min_dist:
min_ind = i
min_dist = dist
if dist >= max_dist:
max_ind = i
max_dist = dist
print(min_ind, max_ind)
if len(files_list) > 0:
# saving the closest member
path_closest = os.path.join(curr_path, files_list[min_ind])
new_closest = os.path.join(curr_path, 'closest.png')
os.system("cp %s %s" % (path_closest, new_closest))
# saving the farthest member
path_farthest = os.path.join(curr_path, files_list[max_ind])
new_farthest = os.path.join(curr_path, 'farthest.png')
os.system("cp %s %s" % (path_farthest, new_farthest))
else:
print("no members in cluster %d" % (k))
def __init__(self, net='alex', if_spatial=False, if_cuda=True):
super().__init__()
import lpips
self.lpips_fn = lpips.LPIPS(net=net, spatial=if_spatial)
if if_cuda:
self.lpips_fn.cuda()
self.lsb = True
def eval_octree(t, dataset, args, want_lpips=True, want_frames=False):
import svox
w, h, focal = dataset.w, dataset.h, dataset.focal
if 'llff' in args.config and (not args.spherify):
ndc_config = svox.NDCConfig(width=w, height=h, focal=focal)
else:
ndc_config = None
r = svox.VolumeRenderer(t,
step_size=args.renderer_step_size,
ndc=ndc_config)
print('Evaluating octree')
device = t.data.device
if want_lpips:
import lpips
lpips_vgg = lpips.LPIPS(net="vgg").eval().to(device)
avg_psnr = 0.0
avg_ssim = 0.0
avg_lpips = 0.0
out_frames = []
for idx in tqdm(range(dataset.size)):
c2w = torch.from_numpy(dataset.camtoworlds[idx]).float().to(device)
im_gt_ten = torch.from_numpy(dataset.images[idx]).float().to(device)
im = r.render_persp(c2w,
width=w,
height=h,
fx=focal,
fast=not args.no_early_stop)
im.clamp_(0.0, 1.0)
mse = ((im - im_gt_ten)**2).mean()
psnr = compute_psnr(mse).mean()
ssim = compute_ssim(im, im_gt_ten, max_val=1.0).mean()
avg_psnr += psnr.item()
avg_ssim += ssim.item()
if want_lpips:
lpips_i = lpips_vgg(im_gt_ten.permute([2, 0, 1]).contiguous(),
im.permute([2, 0, 1]).contiguous(),
normalize=True)
avg_lpips += lpips_i.item()
if want_frames:
im = im.cpu()
# vis = np.hstack((im_gt_ten.cpu().numpy(), im.cpu().numpy()))
vis = im.cpu().numpy() # for lpips calculation
vis = (vis * 255).astype(np.uint8)
out_frames.append(vis)
avg_psnr /= dataset.size
avg_ssim /= dataset.size
avg_lpips /= dataset.size
return avg_psnr, avg_ssim, avg_lpips, out_frames
def init():
global transform, loss_fn_alex, device
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
loss_fn_alex = lpips_ori.LPIPS(net='alex')
if torch.cuda.is_available():
loss_fn_alex.cuda()
device = next(loss_fn_alex.parameters()).device
def evalNetQuant(netWeights, batchSz=4, isPM=False):
"""Evaluate the network quantitatively over the entire validation set"""
## Load Model for training:
device = "cuda" if torch.cuda.is_available() else "cpu"
model = centerEsti()
model.load_state_dict(netWeights)
model.to(device)
valSet = getSetLoader(SetType.test, batchSz)
totPsnr = 0
totSsim = 0
percSimAlex = 0
percSimVgg = 0
lossAlex = lpips.LPIPS(net='alex').to(device) # best forward scores
lossVgg = lpips.LPIPS(net='vgg').to(device) # closer to "traditional" perceptual loss, when used for optimization
transRes = lambda x: x.to("cpu").numpy() / len(valSet)
model.eval()
valProg = tqdm(valSet, desc='Test', leave=False, ncols=100)
with torch.no_grad():
for y, x in valProg:
if not isPM:
y = x.mean(axis=1) # TODO: Verify if correct!
nFrames = x.shape[1]
midFrame = int(np.floor(nFrames / 2))
inputs, targets = y.to(device), x[:, midFrame, :, :].to(device)
outputs = model(inputs)
## Loss Calculation:
curPsnr, curSsim = evalImgQuant(targets, outputs)
totPsnr += curPsnr
totSsim += curSsim
percSimAlex += percSim(targets[0], outputs[0], lossAlex)
percSimVgg += percSim(targets[0], outputs[0], lossVgg)
return transRes(totPsnr), transRes(totSsim), transRes(percSimAlex), transRes(percSimVgg)
def __init__(self,
segmenter,
dynamics_model,
generator,
accumulate_flows=False,
detach_extractor=False,
recon_from_first=False,
coord_loss_scale=1.0,
patch_loss_scale=1.0,
mask_patch_loss=False,
seg_loss_scale=1.0,
pred_attn_loss_scale=1.0,
pred_loss_scale=1.0,
dataset="robonet",
feature_extractor="precropped"):
super().__init__()
self._segmenter = segmenter
self._dynamics_model = dynamics_model
self._generator = generator
self._accumulate_flows = accumulate_flows
self._coord_loss_scale = coord_loss_scale
self._patch_loss_scale = patch_loss_scale
self._detach_extractor = detach_extractor
self._recon_from_first = recon_from_first
self._mask_patch_loss = mask_patch_loss
self._seg_loss_scale = seg_loss_scale
self._pred_attn_loss_scale = pred_attn_loss_scale
self._pred_loss_scale = pred_loss_scale
self._lpips = lpips.LPIPS(net='alex') #.cuda()
self.dataset = dataset
self.feature_extractor = feature_extractor
# attn loss
self.l1_loss = nn.L1Loss(
reduction="sum") # manually normalizing based on size of mask
self.mse_loss = nn.MSELoss()
self._mask_patch_loss = mask_patch_loss
self._lpips = lpips.LPIPS(net='alex').cuda()
def main():
args = get_args()
img_size = args.eval_image_size
src_dir = args.src_dir
num_samples = args.num_samples
l1_loss_fn = nn.L1Loss()
loss_fn_vgg = lpips.LPIPS(net='vgg')
l1_results = np.empty(num_samples)
ssim_results = np.empty(num_samples)
lpips_results = np.empty(num_samples)
transforms = trafo.Compose(
[trafo.ToTensor(),
trafo.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
for idx in range(num_samples):
idx_str = "{:05d}".format(idx)
# reads images as (height,width,channel)
im_pred = imageio.imread("{}/{}_pred.png".format(src_dir, idx_str))
im_true = imageio.imread("{}/{}_tgt.png".format(src_dir, idx_str))
# normalize the values to be [-1,1]
im_pred_torch = transforms(im_pred).unsqueeze(0)
im_true_torch = transforms(im_true).unsqueeze(0)
# resize images to be at the desired resolution if not already
if not args.no_resize and (im_pred_torch.shape[-1] != img_size
or im_true_torch.shape[-1] != img_size):
im_pred_torch = F.interpolate(im_pred_torch, (img_size, img_size))
im_true_torch = F.interpolate(im_true_torch, (img_size, img_size))
# compute l1 error
l1_loss = l1_loss_fn(im_pred_torch, im_true_torch)
l1_results[idx] = l1_loss
# compute ssim error
ssim_loss = ssim(im_pred_torch, im_true_torch)
ssim_results[idx] = ssim_loss
# Compute lpips score
lp = loss_fn_vgg(im_pred_torch, im_true_torch)
lpips_results[idx] = lp
if idx % 1000 == 0:
print(
f"{idx}, im dim: {im_pred_torch.shape[-1]} -- l1, ssim loss, lpips: {l1_loss, ssim_loss, lp}"
)
print(f"L1 loss mean: {l1_results.mean()}, std: {l1_results.std()}")
print(f"SSIM loss mean: {ssim_results.mean()}, std: {ssim_results.std()}")
print(f"LPIPS score: {lpips_results.mean()}, std: {lpips_results.std()}")
def __init__(self, generator, z_size, lr, device, target_image, logProbWeight=0.0, shiftLossWeight=1.0, loss_type="default"):
self.logProbWeight = logProbWeight
self.loss_type = loss_type
self.generator = generator
self.z_size = z_size
self.z_pdf = torch.distributions.Normal(0, 1)
# Latent Vector to optimize over
self.z_vec = torch.randn(1, z_size, 1, 1, device=device).clone().detach()
self.z_vec.requires_grad = True
# Affine transformation variables to optimize over
self.dy = torch.tensor(0.0).clone().detach() # Shift image in y direction. (+) -> down, (-) -> up, range [-1, 1]
self.dy.requires_grad = True
self.dx = torch.tensor(0.0).clone().detach() # Shift image in x direction. (+) -> right, (-) -> left, range [-1, 1]
self.dx.requires_grad = True
self.scale = torch.tensor(0.0).clone().detach() # Zoom in or out. (+) -> zoom in, (-) -> zoom out, range [-1, 1]
self.scale.requires_grad = True
self.rot = torch.tensor(0.0).clone().detach() # Rotate image. (+) -> Clockwise, (-) -> Counterclockwise
self.rot.requires_grad = True
self.shiftLossWeight = shiftLossWeight
# Define optimizer and params to optimize over
self.params = [self.z_vec]
# Add affine parameters if affine loss is specified
if "affine" in self.loss_type:
self.params.extend([self.dx, self.dy])
if "scale" in self.loss_type:
self.params.extend([self.scale])
if "rot" in self.loss_type:
self.params.extend([self.rot])
self.optim = torch.optim.Adam(self.params, lr=lr, betas=(0.9, 0.99))
#self.optim = torch.optim.SGD(self.params, lr=lr, momentum=0.9)
#self.scheduler = CosineAnnealingWarmRestarts(self.optim, 100, 1)
#self.scheduler = CyclicLR(self.optim, base_lr=1e-2, max_lr=0.1, step_size_up=500)
self.target_image = target_image
# Gradient Convolution Kernels
self.kernel_x = torch.FloatTensor([[-1, 0, 1],
[-1, 0, 1],
[-1, 0, 1]]).unsqueeze(0).unsqueeze(0).to(device)
self.kernel_y = torch.FloatTensor([[-1, -1, -1],
[0, 0, 0],
[1, 1, 1]]).unsqueeze(0).unsqueeze(0).to(device)
self.lapacian_operator = torch.FloatTensor([[0, 1, 0],
[1, -4, 1],
[0, 1, 0]]).unsqueeze(0).unsqueeze(0).to(device)
if 'lpips' in self.loss_type:
self.lpips_loss = lpips.LPIPS(net='vgg').to(device)
def __init__(self, net, type="mae"):
super(SelfSupDisReconLoss, self).__init__()
assert type in ["mse", "mae", "perceptual"]
if type == "mse":
self.loss = nn.MSELoss(reduction="mean")
elif type == "mae":
self.loss = nn.L1Loss(reduction="mean")
# TODO: Check device casting
elif type == "perceptual":
self.loss = lpips.LPIPS(net='vgg', verbose=False)
self.net = net
def __init__(self, net='vgg', use_gpu=True):
""" LPIPS loss with spatial weighting """
super(PerceptualLoss, self).__init__()
self.loss_fn = lpips.LPIPS(net=net, spatial=True)
if use_gpu:
self.loss_fn = self.loss_fn.cuda()
# current pip version does not support DataParallel
# self.loss_fn = nn.DataParallel(self.loss_fn)
return
def calculate_lpips(img1, img2, crop_border, input_order='HWC'):
"""Calculate LPIPS metric.
We use the official params estimated from the pristine dataset.
We use the recommended block size (96, 96) without overlaps.
Args:
img (ndarray): Input image whose quality needs to be computed.
The input image must be in range [0, 255] with float/int type.
The input_order of image can be 'HW' or 'HWC' or 'CHW'. (BGR order)
If the input order is 'HWC' or 'CHW', it will be converted to gray
or Y (of YCbCr) image according to the ``convert_to`` argument.
crop_border (int): Cropped pixels in each edge of an image. These
pixels are not involved in the metric calculation.
input_order (str): Whether the input order is 'HW', 'HWC' or 'CHW'.
Default: 'HWC'.
Returns:
float: LPIPS result.
"""
assert img1.shape == img2.shape, (
f'Image shapes are differnet: {img1.shape}, {img2.shape}.')
if input_order not in ['HWC', 'CHW']:
raise ValueError(
f'Wrong input_order {input_order}. Supported input_orders are '
'"HWC" and "CHW"')
img1 = reorder_image(img1, input_order=input_order)
img2 = reorder_image(img2, input_order=input_order)
img1 = img1.astype(np.float64)
img2 = img2.astype(np.float64)
if crop_border != 0:
img1 = img1[crop_border:-crop_border, crop_border:-crop_border, ...]
img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]
img1 = img1.astype(np.float32)
img2 = img2.astype(np.float32)
img1, img2 = totensor([img1, img2], bgr2rgb=False, float32=True)
img1 = img1.unsqueeze(0)
img2 = img2.unsqueeze(0)
# image should be RGB, IMPORTANT: normalized to [-1,1]
img1 = (img1 / 255. - 0.5) * 2
img2 = (img2 / 255. - 0.5) * 2
loss_fn_alex = lpips.LPIPS(net='alex',
verbose=False) # best forward scores
metric = loss_fn_alex(img1, img2).squeeze(0).float().detach().cpu().numpy()
return metric.mean()
def calculate_lpips(img1, img2):
# img1 and img2 have range [0, 255]
loss_fn_vgg = lpips.LPIPS(net='vgg')
img1 = img1[np.newaxis, :]
img2 = img2[np.newaxis, :]
img1 = img1.transpose(0, 3, 2, 1)
img2 = img2.transpose(0, 3, 2, 1)
img1 = torch.from_numpy(2 * (img1 / 255) - 1).float()
img2 = torch.from_numpy(2 * (img2 / 255) - 1).float()
return loss_fn_vgg(img1, img2).detach().numpy().flatten()[0]
def __init__(self, imgs_limit=50, sample_size=5, num_simulations=1):
self.lpips = 0
self.psnr = 0
self.ssim = 0
self.imgs_limit = imgs_limit
self.sample_size = sample_size
self.num_simulations = num_simulations
self.base_path = '/Volumes/LaCie SSD/cassava-leaf-disease-classification'
self.imgpack = LoadImageDataset(self.base_path, 'train_images',
'test_images', self.imgs_limit)
self.loss_fn = lpips.LPIPS(net='alex', version='0.1')
self.run()
def lpips_calculate(x, y, net='alex', gpu=True):
# input range is 0~255
# image should be RGB, and normalized to [-1,1]
x = im2tensor(x[:, :, ::-1])
y = im2tensor(y[:, :, ::-1])
loss_fn = lpips.LPIPS(net=net, verbose=False)
if gpu:
x = x.cuda()
y = y.cuda()
loss_fn = loss_fn.cuda()
lpips_value = loss_fn(x, y)
return lpips_value.item()
def __init__(self, device, G_mapping, G_synthesis, E):
super().__init__()
self.device = device
self.G_mapping = G_mapping
self.G_synthesis = G_synthesis
self.E = E
self.lpips_loss = lpips.LPIPS(net='vgg').to(device)
self.mse_loss = nn.MSELoss().to(device)
#self.l1_loss = nn.SmoothL1Loss().to(device)
self.lambda1 = 0.02
self.lambda2 = 0.3
def __init__(self, net='vgg', use_gpu=True, precision='float'):
""" LPIPS loss with spatial weighting """
super(Masked_LPIPS_Loss, self).__init__()
self.lpips = lpips.LPIPS(net=net, spatial=True).eval()
if use_gpu:
self.lpips = nn.DataParallel(self.lpips).cuda()
if precision == 'half':
self.lpips.half()
elif precision == 'float':
self.lpips.float()
elif precision == 'double':
self.lpips.double()
return
def __init__(self, net='vgg', device='cuda'):
'''
net: ['alex', 'vgg']
References:
https://github.com/richzhang/PerceptualSimilarity/blob/master/lpips_2imgs.py
'''
import lpips
self.device = device
## Initializing the model
self.loss_fn = lpips.LPIPS(net=net)
self.loss_fn = self.loss_fn.to(device).eval()
pass
