def get_image_from_latant_code(latent_code, file_name):
tflib.init_tf()
with dnnlib.util.open_url(STYLEGAN_MODEL_URL,
cache_dir=config.cache_dir) as f:
_G, _D, Gs = pickle.load(f)
print(latent_code.shape)
latent_code = tf.constant(latent_code)
generated_img = Gs.components.synthesis.get_output_for(
latent_code, randomize_noise=False)
generated_img = tf.transpose(generated_img, [0, 2, 3, 1])
generated_img = ((generated_img + 1) / 2) * 255
generated_img_resized_to_original = tf.image.resize_images(
generated_img,
tuple(args.input_img_size),
method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
generated_img_for_display = tf.saturate_cast(
generated_img_resized_to_original, tf.uint8)
sess = tf.get_default_session()
#tf.global_variables_initializer()
sess.run(tf.variables_initializer([]))
reconstructed_imgs = sess.run(fetches=generated_img_for_display)
print("final image shape is:", reconstructed_imgs.shape)
imageio.imwrite(os.path.join(args.restorations_dir, file_name),
reconstructed_imgs[0])
def get_features(imgs, c):
all_images = []
for image_input in imgs:
tflib.init_tf()
image = cv2.imread(str(image_input))
image = cv2.normalize(image,
None,
alpha=0,
beta=1,
norm_type=cv2.NORM_MINMAX,
dtype=cv2.CV_32F)
image = cv2.resize(image, (256, 256), interpolation=cv2.INTER_AREA)
image = np.array(image)
image = np.expand_dims(image.T, axis=0)
all_images.append(image)
concat_imgs = tf.concat(all_images, 0, name='concat')
logits = c.get_output_for(concat_imgs,
None,
is_validation=True,
randomize_noise=True)
predictions = [tf.nn.softmax(tf.concat([logits, -logits], axis=1))]
result = tflib.run(predictions)[0].tolist()
# return logits
return result[0]
def get_image_from_latant_code(latent_code):
tflib.init_tf()
with dnnlib.util.open_url(STYLEGAN_MODEL_URL,
cache_dir=config.cache_dir) as f:
_G, _D, Gs = pickle.load(f)
generated_img = Gs.components.synthesis.get_output_for(
latent_code, randomize_noise=False)
generated_img = tf.transpose(generated_img, [0, 2, 3, 1])
generated_img = ((generated_img + 1) / 2) * 255
return generated_img
def __init__(self):
self.tf_session = tflib.init_tf()
self.tf_graph = tf.get_default_graph()
# with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
_G, D, Gs = pickle.load(open("stylegan/network-final.pkl", "rb"))
self.cover_gen = Gs
self.device = "cuda"
self.title_gen = TitleGenerator(220, 2048, 1024, use_images=True).to(self.device)
self.title_gen.load_state_dict(torch.load("titlegen/models/title-gen.pt"))
self.title_gen.use_images = True
self.int2char = load_int2char("titlegen")
self.rnd = np.random.RandomState(None)
self.transform = transforms.Compose([
# transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406),
(0.229, 0.224, 0.225))]
)
!git clone https://github.com/NVlabs/stylegan.git
# Commented out IPython magic to ensure Python compatibility.
# %tensorflow_version 1.x
import os
import pickle
import numpy as np
import PIL.Image
import stylegan
from stylegan import config
from stylegan.dnnlib import tflib
from tensorflow.python.util import module_wrapper
module_wrapper._PER_MODULE_WARNING_LIMIT = 0
# Initialize TensorFlow
tflib.init_tf()
# Go into that cloned directory
path = 'stylegan/'
if "stylegan" not in os.getcwd():
os.chdir(path)
# Load pre-trained network
# url = 'https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ' # Downloads the pickled model file: karras2019stylegan-ffhq-1024x1024.pkl
url = 'https://bitbucket.org/ezelikman/gans/downloads/karras2019stylegan-ffhq-1024x1024.pkl'
with stylegan.dnnlib.util.open_url(url, cache_dir=config.cache_dir) as f:
print(f)
_G, _D, Gs = pickle.load(f)
# Gs.print_layers() # Print network details
"""Now you can use this separate script to play with the model without needing to reload it everytime (that would take a while).
def optimize_latent_codes_source_target(args, source_latent_code):
#tf.enable_eager_execution()
tflib.init_tf()
with dnnlib.util.open_url(STYLEGAN_MODEL_URL,
cache_dir=config.cache_dir) as f:
_G, _D, Gs = pickle.load(f)
latent_code_2 = tf.get_variable(
name='latent_code_2',
shape=(1, 18, 512),
dtype='float32',
initializer=tf.initializers.constant(source_latent_code))
generated_img = Gs.components.synthesis.get_output_for(
latent_code_2, randomize_noise=False)
generated_img = tf.transpose(generated_img, [0, 2, 3, 1])
generated_img = ((generated_img + 1) / 2) * 255
original_img = tf.placeholder(
tf.float32, [None, args.input_img_size[0], args.input_img_size[1], 3])
blur_kernel = tf.placeholder(
tf.float32,
[None, args.blur_parameters[0], args.blur_parameters[0], 3])
degraded_img_resized_for_perceptual = tf.image.resize_images(
add_motion_blur(original_img, blur_kernel),
tuple(args.perceptual_img_size),
method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
generated_img_resized_to_original = tf.image.resize_images(
generated_img,
tuple(args.input_img_size),
method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
generated_img_resized_for_perceptual = tf.image.resize_images(
add_motion_blur(generated_img_resized_to_original, blur_kernel),
tuple(args.perceptual_img_size),
method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
generated_img_for_display = tf.saturate_cast(
generated_img_resized_to_original, tf.uint8)
perceptual_model = PerceptualModel(img_size=args.perceptual_img_size)
generated_img_features = perceptual_model(
generated_img_resized_for_perceptual)
target_img_features = perceptual_model(degraded_img_resized_for_perceptual)
loss_op = tf.reduce_mean(
tf.abs(generated_img_features - target_img_features))
optimizer = tf.train.GradientDescentOptimizer(
learning_rate=args.learning_rate)
train_op = optimizer.minimize(loss_op, var_list=[latent_code_2])
sess = tf.get_default_session()
img_names = sorted(os.listdir(args.imgs_dir_second))
for img_name in img_names:
img = imageio.imread(os.path.join(args.imgs_dir_second, img_name))
img = cv2.resize(img, dsize=tuple(args.input_img_size))
blur_kernel_3d = motion_blur_kernel(args.blur_parameters[0],
args.blur_parameters[1])
corrupted_img = add_motion_blur_single_image(img,
args.blur_parameters[0],
args.blur_parameters[1])
imageio.imwrite(os.path.join(args.corruptions_dir, img_name),
corrupted_img)
#imageio.imwrite(os.path.join(args.masks_dir, img_name), mask * 255)
sess.run(
tf.variables_initializer([latent_code_2] + optimizer.variables()))
progress_bar_iterator = tqdm(
iterable=range(args.total_iterations),
bar_format=
'{desc}: {percentage:3.0f}% |{bar}| {n_fmt}/{total_fmt}{postfix}',
desc=img_name)
for i in progress_bar_iterator:
loss, _ = sess.run(fetches=[loss_op, train_op],
feed_dict={
original_img: img[np.newaxis, ...],
blur_kernel: blur_kernel_3d[np.newaxis, ...]
})
progress_bar_iterator.set_postfix_str('loss=%.2f' % loss)
reconstructed_imgs, latent_codes = sess.run(
fetches=[generated_img_for_display, latent_code_2],
feed_dict={
original_img: img[np.newaxis, ...],
blur_kernel: blur_kernel_3d[np.newaxis, ...]
})
imageio.imwrite(os.path.join(args.restorations_dir, img_name),
reconstructed_imgs[0])
np.savez(file=os.path.join(args.latents_dir, img_name + '.npz'),
latent_code=latent_codes[0])
return latent_codes[0]
def optimize_latent_codes(args):
tflib.init_tf()
with dnnlib.util.open_url(STYLEGAN_MODEL_URL,
cache_dir=config.cache_dir) as f:
_G, _D, Gs = pickle.load(f)
latent_code = tf.get_variable(name='latent_code',
shape=(1, 18, 512),
dtype='float32',
initializer=tf.initializers.zeros())
generated_img = Gs.components.synthesis.get_output_for(
latent_code, randomize_noise=False)
generated_img = tf.transpose(generated_img, [0, 2, 3, 1])
generated_img = ((generated_img + 1) / 2) * 255
original_img = tf.placeholder(
tf.float32, [None, args.input_img_size[0], args.input_img_size[1], 3])
degraded_img_resized_for_perceptual = tf.image.grayscale_to_rgb(
tf.image.resize_images(convert_to_greyscale(original_img),
tuple(args.perceptual_img_size),
method=tf.image.ResizeMethod.NEAREST_NEIGHBOR))
print("degraded_img_resized_for_perceptual shape is: ",
degraded_img_resized_for_perceptual.shape)
generated_img_resized_to_original = tf.image.resize_images(
generated_img,
tuple(args.input_img_size),
method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
print("generated_img_resized_to_original shape is: ",
generated_img_resized_to_original.shape)
generated_img_resized_for_perceptual = tf.image.grayscale_to_rgb(
tf.image.resize_images(
convert_to_greyscale(generated_img_resized_to_original),
tuple(args.perceptual_img_size),
method=tf.image.ResizeMethod.NEAREST_NEIGHBOR))
print("generated_img_resized_for_perceptual shape is: ",
generated_img_resized_for_perceptual.shape)
generated_img_for_display = tf.saturate_cast(
generated_img_resized_to_original, tf.uint8)
print("generated image shape is: ", generated_img_for_display.shape)
perceptual_model = PerceptualModel(img_size=args.perceptual_img_size)
generated_img_features = perceptual_model(
generated_img_resized_for_perceptual)
target_img_features = perceptual_model(degraded_img_resized_for_perceptual)
loss_op = tf.reduce_mean(
tf.abs(generated_img_features - target_img_features))
optimizer = tf.train.GradientDescentOptimizer(
learning_rate=args.learning_rate)
train_op = optimizer.minimize(loss_op, var_list=[latent_code])
sess = tf.get_default_session()
img_names = sorted(os.listdir(args.imgs_dir))
for img_name in img_names:
img = imageio.imread(os.path.join(args.imgs_dir, img_name))
img = cv2.resize(img, dsize=tuple(args.input_img_size))
# img = tf.image.convert_image_dtype(img, dtype=tf.float32, saturate=False)
greyscale_img = np.dot(img, [0.2126, 0.7152, 0.0722])
imageio.imwrite(os.path.join(args.corruptions_dir, img_name),
greyscale_img)
sess.run(
tf.variables_initializer([latent_code] + optimizer.variables()))
progress_bar_iterator = tqdm(
iterable=range(args.total_iterations),
bar_format=
'{desc}: {percentage:3.0f}% |{bar}| {n_fmt}/{total_fmt}{postfix}',
desc=img_name)
for i in progress_bar_iterator:
loss, _ = sess.run(fetches=[loss_op, train_op],
feed_dict={original_img: img[np.newaxis, ...]})
progress_bar_iterator.set_postfix_str('loss=%.2f' % loss)
reconstructed_imgs, latent_codes = sess.run(
fetches=[generated_img_for_display, latent_code],
feed_dict={original_img: img[np.newaxis, ...]})
imageio.imwrite(os.path.join(args.restorations_dir, img_name),
reconstructed_imgs[0])
np.savez(file=os.path.join(args.latents_dir, img_name + '.npz'),
latent_code=latent_codes[0])
