Adding more losses for perceptual model

README.md

@@ -1,5 +1,40 @@
-## StyleGAN2 — Official TensorFlow Implementation
+## StyleGAN2 — Encoder/Projector for Official TensorFlow Implementation
+![Python 3.6](https://img.shields.io/badge/python-3.6-green.svg?style=plastic)
+![TensorFlow 1.10](https://img.shields.io/badge/tensorflow-1.10-green.svg?style=plastic)
+![cuDNN 7.3.1](https://img.shields.io/badge/cudnn-7.3.1-green.svg?style=plastic)
+![License CC BY-NC](https://img.shields.io/badge/license-CC_BY--NC-green.svg?style=plastic)
 
+This is a port of [Puzer/stylegan-encoder](https://github.com/Puzer/stylegan-encoder) for [NVlabs/stylegan2](https://github.com/NVlabs/stylegan2), plus a modified StyleGAN2 projector.
+
+![Teaser image](./docs/stylegan2encoder-teaser-1024x256.png)
+
+### Generating latent representation of your images, using the original encoder
+`pip install tensorflow-gpu==1.14`
+
+`git clone https://github.com/rolux/stylegan2encoder.git`
+
+`cd stylegan2encoder`
+
+You can generate latent representations of your own images using two scripts:
+
+1) Extract and align faces from images
+
+`python align_images.py raw_images/ aligned_images/`
+
+2) Find latent representation of aligned images
+
+`python encode_images.py aligned_images/ generated_images/ latent_representations/`
+
+### Generating latent representation of your images, using the modified projector
+Replace step 2 with:
+
+`python project_images.py aligned_images/ generated_images/`
+
+This is usually preferable. It also allows you to render a video of the optimization process. To see all available options, type:
+
+`python project_images.py -h`
+
+## Original Readme
 ![Teaser image](./docs/stylegan2-teaser-1024x256.png)
 
 **Analyzing and Improving the Image Quality of StyleGAN**<br>

align_images.py

@@ -0,0 +1,37 @@
+import os
+import sys
+import bz2
+from keras.utils import get_file
+from ffhq_dataset.face_alignment import image_align
+from ffhq_dataset.landmarks_detector import LandmarksDetector
+
+LANDMARKS_MODEL_URL = 'http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2'
+
+
+def unpack_bz2(src_path):
+    data = bz2.BZ2File(src_path).read()
+    dst_path = src_path[:-4]
+    with open(dst_path, 'wb') as fp:
+        fp.write(data)
+    return dst_path
+
+
+if __name__ == "__main__":
+    """
+    Extracts and aligns all faces from images using DLib and a function from original FFHQ dataset preparation step
+    python align_images.py /raw_images /aligned_images
+    """
+
+    landmarks_model_path = unpack_bz2(get_file('shape_predictor_68_face_landmarks.dat.bz2',
+                                               LANDMARKS_MODEL_URL, cache_subdir='temp'))
+    RAW_IMAGES_DIR = sys.argv[1]
+    ALIGNED_IMAGES_DIR = sys.argv[2]
+
+    landmarks_detector = LandmarksDetector(landmarks_model_path)
+    for img_name in [f for f in os.listdir(RAW_IMAGES_DIR) if f[0] not in '._']:
+        raw_img_path = os.path.join(RAW_IMAGES_DIR, img_name)
+        for i, face_landmarks in enumerate(landmarks_detector.get_landmarks(raw_img_path), start=1):
+            face_img_name = '%s_%02d.png' % (os.path.splitext(img_name)[0], i)
+            aligned_face_path = os.path.join(ALIGNED_IMAGES_DIR, face_img_name)
+            os.makedirs(ALIGNED_IMAGES_DIR, exist_ok=True)
+            image_align(raw_img_path, aligned_face_path, face_landmarks)

dataset_tool.py

@@ -55,7 +55,7 @@ def close(self):
         self.tfr_writers = []
         if self.print_progress:
             print('%-40s\r' % '', end='', flush=True)
-            print('Added %d images.' % self.cur_images)
+            print('Added %d image%s.' % (self.cur_images, 's'[:self.cur_images > 1]))
 
     def choose_shuffled_order(self): # Note: Images and labels must be added in shuffled order.
         order = np.arange(self.expected_images)

dnnlib/tflib/network.py

@@ -359,6 +359,7 @@ def run(self,
             minibatch_size: int = None,
             num_gpus: int = 1,
             assume_frozen: bool = False,
+            custom_inputs: Any = None,
             **dynamic_kwargs) -> Union[np.ndarray, Tuple[np.ndarray, ...], List[np.ndarray]]:
         """Run this network for the given NumPy array(s), and return the output(s) as NumPy array(s).
 
@@ -374,6 +375,7 @@ def run(self,
             minibatch_size:     Maximum minibatch size to use, None = disable batching.
             num_gpus:           Number of GPUs to use.
             assume_frozen:      Improve multi-GPU performance by assuming that the trainable parameters will remain changed between calls.
+            custom_inputs:      Allow to use another tensor as input instead of default placeholders.
             dynamic_kwargs:     Additional keyword arguments to be passed into the network build function.
         """
         assert len(in_arrays) == self.num_inputs
@@ -398,9 +400,14 @@ def unwind_key(obj):
         # Build graph.
         if key not in self._run_cache:
             with tfutil.absolute_name_scope(self.scope + "/_Run"), tf.control_dependencies(None):
-                with tf.device("/cpu:0"):
-                    in_expr = [tf.placeholder(tf.float32, name=name) for name in self.input_names]
-                    in_split = list(zip(*[tf.split(x, num_gpus) for x in in_expr]))
+                if custom_inputs is not None:
+                    with tf.device("/gpu:0"):
+                        in_expr = [input_builder(name) for input_builder, name in zip(custom_inputs, self.input_names)]
+                        in_split = list(zip(*[tf.split(x, num_gpus) for x in in_expr]))
+                else:
+                    with tf.device("/cpu:0"):
+                        in_expr = [tf.placeholder(tf.float32, name=name) for name in self.input_names]
+                        in_split = list(zip(*[tf.split(x, num_gpus) for x in in_expr]))
 
                 out_split = []
                 for gpu in range(num_gpus):

dnnlib/tflib/tfutil.py

@@ -237,7 +237,7 @@ def convert_images_from_uint8(images, drange=[-1,1], nhwc_to_nchw=False):
     return images * ((drange[1] - drange[0]) / 255) + drange[0]
 
 
-def convert_images_to_uint8(images, drange=[-1,1], nchw_to_nhwc=False, shrink=1):
+def convert_images_to_uint8(images, drange=[-1,1], nchw_to_nhwc=False, shrink=1, uint8_cast=True):
     """Convert a minibatch of images from float32 to uint8 with configurable dynamic range.
     Can be used as an output transformation for Network.run().
     """
@@ -249,4 +249,6 @@ def convert_images_to_uint8(images, drange=[-1,1], nchw_to_nhwc=False, shrink=1)
         images = tf.transpose(images, [0, 2, 3, 1])
     scale = 255 / (drange[1] - drange[0])
     images = images * scale + (0.5 - drange[0] * scale)
-    return tf.saturate_cast(images, tf.uint8)
+    if uint8_cast:
+        images = tf.saturate_cast(images, tf.uint8)
+    return images

encode_images.py

@@ -0,0 +1,79 @@
+import os
+import argparse
+import pickle
+from tqdm import tqdm
+import PIL.Image
+import numpy as np
+import dnnlib
+import dnnlib.tflib as tflib
+import pretrained_networks
+from encoder.generator_model import Generator
+from encoder.perceptual_model import PerceptualModel
+
+
+def split_to_batches(l, n):
+    for i in range(0, len(l), n):
+        yield l[i:i + n]
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Find latent representation of reference images using perceptual loss')
+    parser.add_argument('src_dir', help='Directory with images for encoding')
+    parser.add_argument('generated_images_dir', help='Directory for storing generated images')
+    parser.add_argument('dlatent_dir', help='Directory for storing dlatent representations')
+
+    parser.add_argument('--network_pkl', default='gdrive:networks/stylegan2-ffhq-config-f.pkl', help='Path to local copy of stylegan2-ffhq-config-f.pkl')
+
+    # for now it's unclear if larger batch leads to better performance/quality
+    parser.add_argument('--batch_size', default=1, help='Batch size for generator and perceptual model', type=int)
+
+    # Perceptual model params
+    parser.add_argument('--image_size', default=256, help='Size of images for perceptual model', type=int)
+    parser.add_argument('--lr', default=1., help='Learning rate for perceptual model', type=float)
+    parser.add_argument('--iterations', default=1000, help='Number of optimization steps for each batch', type=int)
+
+    # Generator params
+    parser.add_argument('--randomize_noise', default=False, help='Add noise to dlatents during optimization', type=bool)
+    args, other_args = parser.parse_known_args()
+
+    ref_images = [os.path.join(args.src_dir, x) for x in os.listdir(args.src_dir)]
+    ref_images = list(filter(os.path.isfile, ref_images))
+
+    if len(ref_images) == 0:
+        raise Exception('%s is empty' % args.src_dir)
+
+    os.makedirs(args.generated_images_dir, exist_ok=True)
+    os.makedirs(args.dlatent_dir, exist_ok=True)
+
+    # Initialize generator and perceptual model
+    tflib.init_tf()
+    generator_network, discriminator_network, Gs_network = pretrained_networks.load_networks(args.network_pkl)
+
+    generator = Generator(Gs_network, args.batch_size, randomize_noise=args.randomize_noise)
+    perceptual_model = PerceptualModel(args.image_size, layer=9, batch_size=args.batch_size)
+    perceptual_model.build_perceptual_model(generator.generated_image)
+
+    # Optimize (only) dlatents by minimizing perceptual loss between reference and generated images in feature space
+    for images_batch in tqdm(split_to_batches(ref_images, args.batch_size), total=len(ref_images)//args.batch_size):
+        names = [os.path.splitext(os.path.basename(x))[0] for x in images_batch]
+
+        perceptual_model.set_reference_images(images_batch)
+        op = perceptual_model.optimize(generator.dlatent_variable, iterations=args.iterations, learning_rate=args.lr)
+        pbar = tqdm(op, leave=False, total=args.iterations)
+        for loss in pbar:
+            pbar.set_description(' '.join(names)+' Loss: %.2f' % loss)
+        print(' '.join(names), ' loss:', loss)
+
+        # Generate images from found dlatents and save them
+        generated_images = generator.generate_images()
+        generated_dlatents = generator.get_dlatents()
+        for img_array, dlatent, img_name in zip(generated_images, generated_dlatents, names):
+            img = PIL.Image.fromarray(img_array, 'RGB')
+            img.save(os.path.join(args.generated_images_dir, f'{img_name}.png'), 'PNG')
+            np.save(os.path.join(args.dlatent_dir, f'{img_name}.npy'), dlatent)
+
+        generator.reset_dlatents()
+
+
+if __name__ == "__main__":
+    main()

encoder/generator_model.py

@@ -0,0 +1,52 @@
+import tensorflow as tf
+import numpy as np
+import dnnlib.tflib as tflib
+from functools import partial
+
+
+def create_stub(name, batch_size):
+    return tf.constant(0, dtype='float32', shape=(batch_size, 0))
+
+
+def create_variable_for_generator(name, batch_size):
+    return tf.get_variable('learnable_dlatents',
+                           shape=(batch_size, 18, 512),
+                           dtype='float32',
+                           initializer=tf.initializers.random_normal())
+
+
+class Generator:
+    def __init__(self, model, batch_size, randomize_noise=False):
+        self.batch_size = batch_size
+
+        self.initial_dlatents = np.zeros((self.batch_size, 18, 512))
+        model.components.synthesis.run(self.initial_dlatents,
+                                       randomize_noise=randomize_noise, minibatch_size=self.batch_size,
+                                       custom_inputs=[partial(create_variable_for_generator, batch_size=batch_size),
+                                                      partial(create_stub, batch_size=batch_size)],
+                                       structure='fixed')
+
+        self.sess = tf.get_default_session()
+        self.graph = tf.get_default_graph()
+
+        self.dlatent_variable = next(v for v in tf.global_variables() if 'learnable_dlatents' in v.name)
+        self.set_dlatents(self.initial_dlatents)
+
+        self.generator_output = self.graph.get_tensor_by_name('G_synthesis_1/_Run/concat:0')
+        self.generated_image = tflib.convert_images_to_uint8(self.generator_output, nchw_to_nhwc=True, uint8_cast=False)
+        self.generated_image_uint8 = tf.saturate_cast(self.generated_image, tf.uint8)
+
+    def reset_dlatents(self):
+        self.set_dlatents(self.initial_dlatents)
+
+    def set_dlatents(self, dlatents):
+        assert (dlatents.shape == (self.batch_size, 18, 512))
+        self.sess.run(tf.assign(self.dlatent_variable, dlatents))
+
+    def get_dlatents(self):
+        return self.sess.run(self.dlatent_variable)
+
+    def generate_images(self, dlatents=None):
+        if dlatents:
+            self.set_dlatents(dlatents)
+        return self.sess.run(self.generated_image_uint8)

encoder/perceptual_model.py

@@ -0,0 +1,78 @@
+import numpy as np
+import tensorflow as tf
+from keras.models import Model
+from keras.applications.vgg16 import VGG16, preprocess_input
+from keras.preprocessing import image
+import keras.backend as K
+
+
+def load_images(images_list, img_size):
+    loaded_images = list()
+    for img_path in images_list:
+        img = image.load_img(img_path, target_size=(img_size, img_size))
+        img = np.expand_dims(img, 0)
+        loaded_images.append(img)
+    loaded_images = np.vstack(loaded_images)
+    preprocessed_images = preprocess_input(loaded_images)
+    return preprocessed_images
+
+
+class PerceptualModel:
+    def __init__(self, img_size, layer=9, batch_size=1, sess=None):
+        self.sess = tf.get_default_session() if sess is None else sess
+        K.set_session(self.sess)
+        self.img_size = img_size
+        self.layer = layer
+        self.batch_size = batch_size
+
+        self.perceptual_model = None
+        self.ref_img_features = None
+        self.features_weight = None
+        self.loss = None
+
+    def build_perceptual_model(self, generated_image_tensor):
+        vgg16 = VGG16(include_top=False, input_shape=(self.img_size, self.img_size, 3))
+        self.perceptual_model = Model(vgg16.input, vgg16.layers[self.layer].output)
+        generated_image = preprocess_input(tf.image.resize_images(generated_image_tensor,
+                                                                  (self.img_size, self.img_size), method=1))
+        generated_img_features = self.perceptual_model(generated_image)
+
+        self.ref_img_features = tf.get_variable('ref_img_features', shape=generated_img_features.shape,
+                                                dtype='float32', initializer=tf.initializers.zeros())
+        self.features_weight = tf.get_variable('features_weight', shape=generated_img_features.shape,
+                                               dtype='float32', initializer=tf.initializers.zeros())
+        self.sess.run([self.features_weight.initializer, self.features_weight.initializer])
+
+        self.loss = tf.losses.mean_squared_error(self.features_weight * self.ref_img_features,
+                                                 self.features_weight * generated_img_features) / 82890.0
+
+    def set_reference_images(self, images_list):
+        assert(len(images_list) != 0 and len(images_list) <= self.batch_size)
+        loaded_image = load_images(images_list, self.img_size)
+        image_features = self.perceptual_model.predict_on_batch(loaded_image)
+
+        # in case if number of images less than actual batch size
+        # can be optimized further
+        weight_mask = np.ones(self.features_weight.shape)
+        if len(images_list) != self.batch_size:
+            features_space = list(self.features_weight.shape[1:])
+            existing_features_shape = [len(images_list)] + features_space
+            empty_features_shape = [self.batch_size - len(images_list)] + features_space
+
+            existing_examples = np.ones(shape=existing_features_shape)
+            empty_examples = np.zeros(shape=empty_features_shape)
+            weight_mask = np.vstack([existing_examples, empty_examples])
+
+            image_features = np.vstack([image_features, np.zeros(empty_features_shape)])
+
+        self.sess.run(tf.assign(self.features_weight, weight_mask))
+        self.sess.run(tf.assign(self.ref_img_features, image_features))
+
+    def optimize(self, vars_to_optimize, iterations=500, learning_rate=1.):
+        vars_to_optimize = vars_to_optimize if isinstance(vars_to_optimize, list) else [vars_to_optimize]
+        optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate)
+        min_op = optimizer.minimize(self.loss, var_list=[vars_to_optimize])
+        for _ in range(iterations):
+            _, loss = self.sess.run([min_op, self.loss])
+            yield loss
+