import torch.nn.functional as F
import numpy as np
import matplotlib.pylab as plt
from matplotlib import cm
from tqdm import tqdm
from time import time
from os.path import join
import os
import sys
import lpips
from GAN_hessian_compute import hessian_compute, get_full_hessian
from torchvision.transforms import ToPILImage
from torchvision.utils import make_grid
from lpips import LPIPS
ImDist = LPIPS(net="squeeze")
ImDist.cuda()
rootdir = r"E:\OneDrive - Washington University in St. Louis\Hessian_summary\App\Inversion"
os.makedirs(rootdir, exist_ok=True)
#%%
data = np.load(
r"E:\OneDrive - Washington University in St. Louis\Hessian_summary\StyleGAN\H_avg_StyleGAN.npz"
)
evc, eva = data["evc_avg"], data["eva_avg"],
evc_tsr = torch.from_numpy(evc).cuda().float()
eva_tsr = torch.from_numpy(eva).cuda().float()
#%%
from torch.optim import Adam, SGD
SGAN = loadStyleGAN()
G = StyleGAN_wrapper(SGAN)
#%%
refvec = torch.randn(1, 512).cuda()
import matplotlib.pylab as plt
from GAN_utils import StyleGAN2_wrapper, loadStyleGAN2
from lpips import LPIPS
from load_hessian_data import load_Haverage
from torch_utils import show_imgrid, save_imgrid
def MSE(im1, im2, mask=None):
# mask is size [sampn, H, W]
if mask is None:
return (im1 - im2).pow(2).mean(dim=[1,2,3])
else:
valnum = mask.sum([1, 2])
diffsum = ((im1 - im2).pow(2).mean(1) * mask).sum([1, 2])
return diffsum / valnum
#%
D = LPIPS(net="vgg")
D.cuda()
D.requires_grad_(False)
D.spatial = True
def mask_LPIPS(im1, im2, mask=None):
diffmap = D(im1, im2) # note there is a singleton channel dimension
if mask is None:
return diffmap.mean([1, 2, 3])
else:
diffsum = (diffmap[:, 0, :, :] * mask).sum([1, 2])
valnum = mask.sum([1, 2])
return diffsum / valnum
#%%
imroot = r"E:\OneDrive - Washington University in St. Louis\GAN_photoedit\src"
resdir = r"E:\OneDrive - Washington University in St. Louis\GAN_photoedit\results"
SGAN = loadStyleGAN2("ffhq-512-avg-tpurun1.pt")