def main():
t0 = time.time()
print('t0:', t0)
# Initialize TensorFlow.
tflib.init_tf() # 0.82s
print('t1:', time.time() - t0)
# Load pre-trained network.
with open('./models/stylegan2-ffhq-config-f.pkl', 'rb') as f:
print('t2:', time.time() - t0)
_G, _D, Gs = pickle.load(f) # 13.09s
print('t3:', time.time() - t0)
with open('./models/vgg16_zhang_perceptual.pkl', 'rb') as f:
lpips = pickle.load(f)
print('t4:', time.time() - t0)
proj = Projector()
proj.set_network(Gs, lpips)
image = PIL.Image.open('./images/example.png')
#image = image.resize((Di.input_shape[2], Di.input_shape[3]), PIL.Image.ANTIALIAS)
image_array = np.array(image).swapaxes(0, 2).swapaxes(1, 2)
image_array = misc.adjust_dynamic_range(image_array, [0, 255], [-1, 1])
print('t5:', time.time() - t0)
proj.start([image_array])
for step in proj.runSteps(1000):
print('\rstep: %d' % step, end='', flush=True)
if step % 10 == 0:
results = proj.get_images()
pilImage = misc.convert_to_pil_image(
misc.create_image_grid(results), drange=[-1, 1])
pilImage.save('./images/project-%d.png' % step)
print('t6:', time.time() - t0)
dlatents = proj.get_dlatents()
noises = proj.get_noises()
print('dlatents:', dlatents.shape)
print('noises:', len(noises), noises[0].shape, noises[-1].shape)
# Load model
g = Generator(1024, 512, 8, pretrained=True).to(device).train()
for param in g.parameters():
param.requires_grad = False
proj = Projector(g)
for i, file in tqdm(enumerate(sorted(image_files))):
print('Projecting {}'.format(file))
# Load image
target_image = Image.open(file)
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
target_image = transform(target_image).to(device)
# Run projector
proj.run(target_image)
# Collect results
generated = proj.get_images()
latents = proj.get_latents()
# Save results
save_str = target_dir + file.split('/')[-1].split('.')[0]
print('Saving {}'.format(save_str + '_p.png'))
save_image(generated, save_str + '_p.png', normalize=True)
torch.save(latents.detach().cpu(), save_str + '.pt')
