|
|
import os |
|
|
import numpy as np |
|
|
import tensorflow as tf |
|
|
|
|
|
from PIL import Image |
|
|
|
|
|
from model.vae import VAE |
|
|
from model import tidev2 |
|
|
|
|
|
|
|
|
def init_vae_model(model_name, latent_dim, input_shape): |
|
|
if model_name == 'tidev2': |
|
|
vae_model = VAE(tidev2.ConvNeXtEncoderTiny(latent_dim=latent_dim), |
|
|
tidev2.ConvNeXtDecoderTiny(latent_dim=latent_dim, image_dims=input_shape[:2], out_channels=input_shape[-1]) |
|
|
) |
|
|
vae_model.build((None, *input_shape)) |
|
|
return vae_model |
|
|
|
|
|
|
|
|
def load_weights(vae, weights_path): |
|
|
print("Loading weights from {}".format(weights_path)) |
|
|
if "ckpt-" in weights_path: |
|
|
ckpt = tf.train.Checkpoint(vae=vae) |
|
|
ckpt.restore(weights_path).expect_partial() |
|
|
return vae |
|
|
if ".TF" in weights_path: |
|
|
vae.load_weights(weights_path, by_name=True) |
|
|
return vae |
|
|
|
|
|
|
|
|
def get_noise_seeded(noise_shape, seed=0): |
|
|
np.random.seed(seed) |
|
|
random_z = np.random.normal(0, 1, noise_shape) |
|
|
return random_z |
|
|
|
|
|
def decode_noise(trained_vae, noise, return_list=False): |
|
|
print("Generating synthetic images ...") |
|
|
pred = trained_vae.decoder.predict(noise, batch_size=1) |
|
|
|
|
|
pred *= 255.0 |
|
|
|
|
|
if return_list: |
|
|
return [img for img in pred] |
|
|
return pred |
|
|
|
|
|
|
|
|
def save_images(save_folder, images, seed=None): |
|
|
print(f"Saving synthetic images into {save_folder}") |
|
|
if isinstance(images, list): |
|
|
for i, image in enumerate(images): |
|
|
image = image.astype(np.uint8) |
|
|
if image.shape[-1] == 1: |
|
|
image = np.squeeze(image, axis=-1) |
|
|
save_filename = f"image-{i}.jpg" if seed is None else f"image-{seed}.jpg" |
|
|
Image.fromarray(image).save(os.path.join(save_folder, save_filename)) |
