def image_list2latent_old(img_list, iterations=750, learning_rate=1.
, init_dlatents: Optional[np.ndarray] = None
, return_image: bool = False) -> Tuple[Optional[np.ndarray], np.ndarray]:
"""
:return: sizes of (batch_size, img_height, img_width, 3), (batch_size, 18, 512)
"""
batch_size = len(img_list)
config.init_generator(batch_size=batch_size)
# generator = config.generator # TODO: messes with parallel
generator = Generator(config.Gs_network, batch_size=batch_size)
generator.reset_dlatents()
perceptual_model = PerceptualModelOld(256, batch_size=batch_size)
perceptual_model.build_perceptual_model(generator.generated_image)
perceptual_model.set_reference_images_from_image(np.array([np.array(im) for im in img_list]))
if init_dlatents is not None:
generator.set_dlatents(init_dlatents)
op = perceptual_model.optimize(generator.dlatent_variable, iterations=iterations, learning_rate=learning_rate)
best_loss = None
best_dlatent = None
with tqdm(total=iterations) as pbar:
for iteration, loss in enumerate(op):
pbar.set_description('Loss: %.2f' % loss)
if best_loss is None or loss < best_loss:
best_loss = loss
best_dlatent = generator.get_dlatents()
generator.stochastic_clip_dlatents()
pbar.update()
print(f"final loss {loss}")
generator.set_dlatents(best_dlatent)
if return_image:
return generator.generate_images(), generator.get_dlatents()
return None, generator.get_dlatents()
def latent2image(self, latent_vector):
config.init_generator()
latent_vector = latent_vector.reshape((1, 18, 512))
config.generator.set_dlatents(latent_vector)
img_array = config.generator.generate_images()[0]
img = Image.fromarray(img_array, 'RGB')
return img.resize((256, 256))
def image2latent(img_list, iterations=750, init_dlatents: Optional[np.ndarray] = None, args=None
, return_image: bool = False, is_aligned:bool=False) \
-> Tuple[List[Optional[np.ndarray]], List[np.ndarray]]:
"""
:return: sizes of (batch_size, img_height, img_width, 3), (batch_size, 18, 512)
"""
batch_size = len(img_list)
args = config.DEFAULT_PERC_PARAM if args is None else args
config.init_generator(batch_size=batch_size)
# generator = config.generator # TODO: messes with parallel
generator = Generator(config.Gs_network, batch_size=batch_size)
generator.reset_dlatents()
perc_model = None
if (args.use_lpips_loss > 0.00000001):
with dnnlib.util.open_url(URL_VGG_16,
cache_dir=stylgan_config.cache_dir) as f:
perc_model = pickle.load(f)
perceptual_model = PerceptualModel(args,
perc_model=perc_model,
batch_size=batch_size)
perceptual_model.build_perceptual_model(generator)
generated_images_list, generated_dlatents_list = [], []
for images_batch in tqdm(split_to_batches(img_list, batch_size),
total=len(img_list) // batch_size):
names = [f"image_{n}" for n in range(len(images_batch))]
# names = [os.path.splitext(os.path.basename(x))[0] for x in images_batch]
perceptual_model.set_reference_images(images_batch,
is_aligned=is_aligned)
# FIXME: split also init_dlatents to batches
if init_dlatents is not None:
generator.set_dlatents(init_dlatents)
op = perceptual_model.optimize(generator.dlatent_variable,
iterations=iterations)
pbar = tqdm(op, position=0, leave=True, total=iterations)
best_loss = None
best_dlatent = None
for loss_dict in pbar:
pbar.set_description(" ".join(names) + ": " + "; ".join(
["{} {:.4f}".format(k, v) for k, v in loss_dict.items()]))
if best_loss is None or loss_dict["loss"] < best_loss:
best_loss = loss_dict["loss"]
best_dlatent = generator.get_dlatents()
generator.stochastic_clip_dlatents()
print(" ".join(names), " Loss {:.4f}".format(best_loss))
generator.set_dlatents(best_dlatent)
generated_images_list = [
generator.generate_images() if return_image else None
]
generated_dlatents_list += [generator.get_dlatents()]
return generated_images_list, generated_dlatents_list