def test_render_text(self, tensor):
if len(tensor.shape) == 3:
image = make_image(tensor)
else:
image = make_image(tensor)[0]
image_with_text = render_text_on_image("test", Image.fromarray(image))
text_array = numpy.array(image_with_text)
assert not numpy.allclose(image, text_array)
assert numpy.allclose(image[:128, :128, :], text_array[:128, :128, :])
def create_gif(self, latent_paths: LatentPaths, image_name: str, destination_dir: Path) -> None:
destination_dir = destination_dir / 'gifs'
destination_dir.mkdir(parents=True, exist_ok=True)
with tempfile.TemporaryDirectory() as temp_dir, tempfile.NamedTemporaryFile(mode='w') as temp_file:
latent_paths = latent_paths.to(self.device)
temp_dir = Path(temp_dir)
progress_bar = tqdm(latent_paths)
progress_bar.set_description("creating gif")
for i, latent in enumerate(progress_bar):
img_gen, _ = self.generator([latent.latent], input_is_latent=True, noise=latent.noise)
image_index = i * self.debug_step
temp_dest_name = temp_dir / f"{image_index}.png"
img_ar = make_image(img_gen)[0]
image = Image.fromarray(img_ar)
image = render_text_on_image(f"{image_index:06}", image)
image.save(temp_dest_name)
print(temp_dest_name, file=temp_file)
process_args = [
'convert',
'-delay 10',
'-loop 0',
f'@{temp_file.name}',
str(destination_dir / f"{image_name}.gif")
]
temp_file.flush()
subprocess.run(' '.join(process_args), shell=True, check=True)
def run(self, trainer: Trainer):
torch.cuda.empty_cache()
try:
for network in self.networks:
network.eval()
with torch.no_grad():
predictions = self.get_predictions()
finally:
for network in self.networks:
network.train()
display_images = torch.cat(predictions, dim=0)
image_grid = torchvision.utils.make_grid(
display_images, nrow=self.input_images.shape[0])
dest_file_name = os.path.join(self.image_dir,
f"{trainer.updater.iteration:08d}.png")
dest_image = make_image(image_grid)
Image.fromarray(dest_image).save(dest_file_name)
if self.log_to_logger:
with get_current_reporter() as reporter:
reporter.add_image({"image_plotter": dest_image},
trainer.updater.iteration)
del display_images
torch.cuda.empty_cache()
def test_make_image(self, tensor):
normalized_tensor = tensor * 2 - 1
image = make_image(normalized_tensor)
assert isinstance(image, numpy.ndarray)
assert image.max() == 255
assert image.min() == 0
def save_images(images: List[torch.Tensor], save_dir: Path, index: int):
dest_file_name = save_dir / f"{index}.png"
images = [Image.fromarray(make_image(image, normalize_func=lambda x: x)) for image in images]
dest_image = Image.new((im := images[0]).mode, (im.width * len(images), im.height))
for i, image in enumerate(images):
dest_image.paste(image, (image.width * i, 0))
dest_image.save(str(dest_file_name))
def save_stylized_images(self, stylized_images: Dict[int, Tuple[torch.Tensor, Optional[LatentPaths]]], content_path: Path, style_path: Path, destination_dir: Path, create_gif: bool = False):
for index, (image_array, optimization_path) in stylized_images.items():
content_name = content_path.stem
style_name = style_path.stem
image_name = f"{content_name}_{style_name}_{index}"
destination_name = destination_dir / f"{image_name}.png"
Image.fromarray(make_image(image_array)[0]).save(destination_name)
if optimization_path is not None and create_gif:
self.projector.create_gif(optimization_path, image_name, destination_dir)
def main(args: Namespace):
transferer = StyleTransferer(args)
content_latents, style_latents = transferer.get_latents(
args.content_path, args.style_path)
if args.mixing_index < 0:
stylized_images = {
i: transferer.do_style_transfer(content_latents, style_latents, i)
for i in range(content_latents.latent.shape[1])
}
else:
stylized_images = {
args.mixing_index:
transferer.do_style_transfer(content_latents, style_latents,
args.mixing_index)
}
destination_dir = Path(
args.content_path
).parent / "simple_style_transfer" / args.destination_dir
destination_dir.mkdir(parents=True, exist_ok=True)
for index, (image_array, optimization_path) in stylized_images.items():
content_base_name = args.content_path
style_base_name = args.style_path
content_name = Path(content_base_name).stem
style_name = Path(style_base_name).stem
image_name = f"{content_name}_{style_name}_{index}"
destination_name = destination_dir / f"{image_name}.png"
Image.fromarray(make_image(image_array)[0]).save(destination_name)
if optimization_path is not None and args.gif:
transferer.projector.create_gif(optimization_path.latent,
optimization_path.noise,
image_name, destination_dir)
def main(args):
projector = Projector(args)
transform = projector.get_transforms()
imgs = []
image_names = []
for file_name in os.listdir(args.files):
if os.path.splitext(file_name)[-1] not in Image.EXTENSION.keys():
continue
image_name = os.path.join(args.files, file_name)
img = transform(Image.open(image_name).convert('RGB'))
image_names.append(image_name)
imgs.append(img)
imgs = torch.stack(imgs, 0).to(args.device)
n_mean_latent = 10000
latent_mean, latent_std = projector.get_mean_latent(n_mean_latent)
for idx in trange(0, len(imgs), args.batch_size):
images = imgs[idx:idx + args.batch_size]
base_noises = projector.generator.make_noise()
base_noises = [
noise.repeat(len(images), 1, 1, 1) for noise in base_noises
]
noises = [noise.detach().clone() for noise in base_noises]
if args.no_mean_latent:
latent_in = torch.normal(0,
latent_std.item(),
size=(len(images),
projector.config['latent_size']),
device=args.device)
else:
latent_in = latent_mean.detach().clone().unsqueeze(0).repeat(
len(images), 1)
if args.w_plus:
latent_in = latent_in.unsqueeze(1).repeat(
1, projector.generator.n_latent, 1)
# optimize latent vector
paths, best_latent = run_image_reconstruction(
args,
projector,
Latents(latent_in, noises),
images,
do_optimize_noise=args.optimize_noise)
# result_file = {'noises': noises}
img_gen, _ = projector.generator(
[best_latent.latent.cuda()],
input_is_latent=True,
noise=[noise.cuda() for noise in best_latent.noise])
img_ar = make_image(img_gen)
destination_dir = Path(args.files) / 'projected' / args.destination
destination_dir.mkdir(parents=True, exist_ok=True)
path_per_image = paths.split()
for i in range(len(images)):
image_name = image_names[idx + i]
image_latent = best_latent[i]
result_file = {
'noise': image_latent.noise,
'latent': image_latent.latent,
}
image_base_name = os.path.splitext(os.path.basename(image_name))[0]
img_name = image_base_name + '-project.png'
pil_img = Image.fromarray(img_ar[i])
pil_img.save(destination_dir / img_name)
torch.save(result_file,
destination_dir / f'results_{image_base_name}.pth')
if args.create_gif:
projector.create_gif(path_per_image[i].to(args.device),
image_base_name, destination_dir)
projector.render_log(destination_dir, image_base_name)
# cleanup
del paths
del best_latent
torch.cuda.empty_cache()
projector.reset()