#%
prev_axes = Wmat[:, list(annot_dict.values())].T.cpu().numpy()
np.savez(join(savedir, "prev_axes.npz"),
basis=Wmat.cpu().numpy(),
annot_dict=annot_dict,
prev_axes=prev_axes)
#%%
refvec = torch.randn(1, 512).cuda()
movvecs = refvec + torch.linspace(-9, 9, 5).cuda().unsqueeze(1) @ Wmat[:,
20:21].T
imgs = G((movvecs).unsqueeze(2).unsqueeze(3)) #
show_imgrid([torch.clamp((imgs + 1) / 2, 0, 1)])
#%%
from lpips import LPIPS
ImDist = LPIPS(net="squeeze").cuda()
ImDist.requires_grad_(False)
#%%
class PGGAN_wrapper2(): # nn.Module
"""
model = loadPGGAN(onlyG=False)
G = PGGAN_wrapper(model.avgG)
model = loadPGGAN()
G = PGGAN_wrapper(model)
"""
def __init__(
self,
PGGAN,
):
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")
G = StyleGAN2_wrapper(SGAN)