#% 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)