|
|
|
|
|
|
|
import numpy as np |
|
from PIL import Image |
|
import wandb |
|
from pti.pti_configs import global_config |
|
import torch |
|
import matplotlib.pyplot as plt |
|
|
|
|
|
def log_image_from_w(w, G, name): |
|
img = get_image_from_w(w, G) |
|
pillow_image = Image.fromarray(img) |
|
wandb.log( |
|
{f"{name}": [ |
|
wandb.Image(pillow_image, caption=f"current inversion {name}")]}, |
|
step=global_config.training_step) |
|
|
|
|
|
def log_images_from_w(ws, G, names): |
|
for name, w in zip(names, ws): |
|
w = w.to(global_config.device) |
|
log_image_from_w(w, G, name) |
|
|
|
|
|
def plot_image_from_w(w, G): |
|
img = get_image_from_w(w, G) |
|
pillow_image = Image.fromarray(img) |
|
plt.imshow(pillow_image) |
|
plt.show() |
|
|
|
|
|
def plot_image(img): |
|
img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).detach().cpu().numpy() |
|
pillow_image = Image.fromarray(img[0]) |
|
plt.imshow(pillow_image) |
|
plt.show() |
|
|
|
|
|
def save_image(name, method_type, results_dir, image, run_id): |
|
image.save(f'{results_dir}/{method_type}_{name}_{run_id}.jpg') |
|
|
|
|
|
def save_w(w, G, name, method_type, results_dir): |
|
im = get_image_from_w(w, G) |
|
im = Image.fromarray(im, mode='RGB') |
|
save_image(name, method_type, results_dir, im) |
|
|
|
|
|
def save_concat_image(base_dir, image_latents, new_inv_image_latent, new_G, |
|
old_G, |
|
file_name, |
|
extra_image=None): |
|
images_to_save = [] |
|
if extra_image is not None: |
|
images_to_save.append(extra_image) |
|
for latent in image_latents: |
|
images_to_save.append(get_image_from_w(latent, old_G)) |
|
images_to_save.append(get_image_from_w(new_inv_image_latent, new_G)) |
|
result_image = create_alongside_images(images_to_save) |
|
result_image.save(f'{base_dir}/{file_name}.jpg') |
|
|
|
|
|
def save_single_image(base_dir, image_latent, G, file_name): |
|
image_to_save = get_image_from_w(image_latent, G) |
|
image_to_save = Image.fromarray(image_to_save, mode='RGB') |
|
image_to_save.save(f'{base_dir}/{file_name}.jpg') |
|
|
|
|
|
def create_alongside_images(images): |
|
res = np.concatenate([np.array(image) for image in images], axis=1) |
|
return Image.fromarray(res, mode='RGB') |
|
|
|
|
|
def get_image_from_w(w, G): |
|
if len(w.size()) <= 2: |
|
w = w.unsqueeze(0) |
|
with torch.no_grad(): |
|
img = G.synthesis(w, noise_mode='const') |
|
img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).detach().cpu().numpy() |
|
return img[0] |
|
|
