Yoder/rosie stage

dataset_tool.py

@@ -18,6 +18,7 @@
 import tensorflow as tf
 import PIL.Image
 import dnnlib.tflib as tflib
+import pickle 
 
 from training import dataset
 
@@ -526,6 +527,36 @@ def create_from_images(tfrecord_dir, image_dir, shuffle):
             tfr.add_image(img)
 
 #----------------------------------------------------------------------------
+
+def create_from_pickle(tfrecord_dir, pickle_path, shuffle):
+    tf.enable_eager_execution()
+    print('Loading images from "%s"' % pickle_path)
+    data = None
+    with open(pickle_path, 'rb') as f:
+        data = pickle.load(f)
+    image_filenames = []
+    image_filenames += data['train']['norms']
+    image_filenames += data['train']['highs']
+    image_filenames += data['train']['lows']
+    image_filenames += data['train']['unlabl']
+
+    if len(image_filenames) == 0:
+        error('Failed to load pickle data :(')
+    else:
+        print("Loaded file count: ", len(image_filenames))
+
+    with TFRecordExporter(tfrecord_dir, len(image_filenames)) as tfr:
+        for idx in range(len(image_filenames)):
+            img = PIL.Image.open(image_filenames[idx])
+            img = img.resize((256,256))
+            img = np.asarray(img)
+            img = img.transpose([2, 0, 1]) # HWC => CHW
+            tfr.add_image(img)
+
+#            img = np.asarray(PIL.Image.open(image_filenames[idx]))
+#            img = tf.image.random_crop(img, (256,256,3)).numpy()
+
+#----------------------------------------------------------------------------
 
 def create_from_hdf5(tfrecord_dir, hdf5_filename, shuffle):
     print('Loading HDF5 archive from "%s"' % hdf5_filename)
@@ -624,6 +655,12 @@ def add_command(cmd, desc, example=None):
     p.add_argument(     'tfrecord_dir',     help='New dataset directory to be created')
     p.add_argument(     'image_dir',        help='Directory containing the images')
     p.add_argument(     '--shuffle',        help='Randomize image order (default: 1)', type=int, default=1)
+
+    p = add_command(    'create_from_pickle', 'Create dataset from a directory full of images.',
+                                            'create_from_pickle datasets/mydataset pickle')
+    p.add_argument(     'tfrecord_dir',     help='New dataset directory to be created')
+    p.add_argument(     'pickle_path',        help='Pickle containing the images')
+    p.add_argument(     '--shuffle',        help='Randomize image order (default: 1)', type=int, default=1)
 
     p = add_command(    'create_from_hdf5', 'Create dataset from legacy HDF5 archive.',
                                             'create_from_hdf5 datasets/celebahq ~/downloads/celeba-hq-1024x1024.h5')

dnnlib/tflib/network.py

@@ -547,7 +547,7 @@ def setup_weight_histograms(self, title: str = None) -> None:
 
                 tf.summary.histogram(name, var)
 
-#----------------------------------------------------------------------------
+# ----------------------------------------------------------------------------
 # Backwards-compatible emulation of legacy output transformation in Network.run().
 
 _print_legacy_warning = True

projector.py

@@ -14,16 +14,16 @@
 #----------------------------------------------------------------------------
 
 class Projector:
-    def __init__(self):
-        self.num_steps                  = 1000
-        self.dlatent_avg_samples        = 10000
-        self.initial_learning_rate      = 0.1
-        self.initial_noise_factor       = 0.05
+    def __init__(self,num_steps=1000,initial_learning_rate=0.1,initial_noise_factor=0.05):
+        self.num_steps                  = num_steps
+        self.dlatent_avg_samples        = 1000
+        self.initial_learning_rate      = initial_learning_rate
+        self.initial_noise_factor       = initial_noise_factor
         self.lr_rampdown_length         = 0.25
         self.lr_rampup_length           = 0.05
         self.noise_ramp_length          = 0.75
         self.regularize_noise_weight    = 1e5
-        self.verbose                    = False
+        self.verbose                    = True
         self.clone_net                  = True
 
         self._Gs                    = None
@@ -104,7 +104,7 @@ def set_network(self, Gs, minibatch_size=1):
 
         # Loss graph.
         self._info('Building loss graph...')
-        self._target_images_var = tf.Variable(tf.zeros(proc_images_expr.shape), name='target_images_var')
+        self._target_images_var = tf.Variable(tf.zeros(self._images_expr.shape), name='target_images_var')
         if self._lpips is None:
             self._lpips = misc.load_pkl('https://nvlabs-fi-cdn.nvidia.com/stylegan/networks/metrics/vgg16_zhang_perceptual.pkl')
         self._dist = self._lpips.get_output_for(proc_images_expr, self._target_images_var)
@@ -150,12 +150,13 @@ def start(self, target_images):
         # Prepare target images.
         self._info('Preparing target images...')
         target_images = np.asarray(target_images, dtype='float32')
-        target_images = (target_images + 1) * (255 / 2)
+        # target_images = (target_images + 1) * (255 / 2)
         sh = target_images.shape
         assert sh[0] == self._minibatch_size
         if sh[2] > self._target_images_var.shape[2]:
             factor = sh[2] // self._target_images_var.shape[2]
             target_images = np.reshape(target_images, [-1, sh[1], sh[2] // factor, factor, sh[3] // factor, factor]).mean((3, 5))
+            print('gagan note: resizing something that shouldnt be resized?\n')
 
         # Initialize optimization state.
         self._info('Initializing optimization state...')

run_generator.py

@@ -12,6 +12,9 @@
 import re
 import sys
 
+import os
+import numpy as np
+
 import pretrained_networks
 
 #----------------------------------------------------------------------------
@@ -37,6 +40,50 @@ def generate_images(network_pkl, seeds, truncation_psi):
 
 #----------------------------------------------------------------------------
 
+def generate_images_of_mean(network_pkl, count, truncation_psi=0.7, fixed_noise=True):
+    count = int(count)
+    print('Loading networks from "%s"...' % network_pkl)
+    _G, _D, Gs = pretrained_networks.load_networks(network_pkl)
+    noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]
+
+    Gs_kwargs = dnnlib.EasyDict()
+    Gs_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
+    Gs_kwargs.randomize_noise = False
+    if truncation_psi is not None:
+        Gs_kwargs.truncation_psi = truncation_psi
+
+    rnd = np.random.RandomState(42)
+    w_avg = Gs.get_var('dlatent_avg') # [component]
+
+    Gs_syn_kwargs = dnnlib.EasyDict()
+    Gs_syn_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
+    Gs_syn_kwargs.randomize_noise = False
+    Gs_syn_kwargs.minibatch_size = 1
+
+    # w code means
+    latents_dir = '/data/yoder_lab/means_highprob'
+    for mean_code in os.listdir(latents_dir):
+        if mean_code.endswith('.npy') and mean_code[-9] == 'w':
+            w = np.load(os.path.join(latents_dir, mean_code)) # [minibatch, component]
+            print('shape: ', w.shape)
+            w = w_avg + (w - w_avg) * truncation_psi # [minibatch, layer, component]
+
+            for idx in range(0,count):
+                tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width]
+                images = Gs.components.synthesis.run(w, **Gs_syn_kwargs)
+                PIL.Image.fromarray(images[0], 'RGB').save('/data/yoder_lab/means_highprob/cat_%s_%04d.png' % (mean_code[-7],idx))
+
+    # z code means
+    # for mean_code in os.listdir(latents_dir):
+    #     if mean_code.endswith('.npy') and mean_code[5] == 'z':
+    #         z = np.load(os.path.join(latents_dir, mean_code))
+    #         for idx in range(0,count):
+    #             tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width]
+    #             images = Gs.run(z, None, **Gs_kwargs) # [minibatch, height, width, channel]
+    #             PIL.Image.fromarray(images[0], 'RGB').save('/data/yoder_lab/means/latent_means_w_fixed/cat_%s_%04d.png' % (mean_code[7],idx))
+
+#----------------------------------------------------------------------------
+
 def style_mixing_example(network_pkl, row_seeds, col_seeds, truncation_psi, col_styles, minibatch_size=4):
     print('Loading networks from "%s"...' % network_pkl)
     _G, _D, Gs = pretrained_networks.load_networks(network_pkl)
@@ -120,7 +167,7 @@ def main():
     parser = argparse.ArgumentParser(
         description='''StyleGAN2 generator.
 
-Run 'python %(prog)s <subcommand> --help' for subcommand help.''',
+        Run 'python %(prog)s <subcommand> --help' for subcommand help.''',
         epilog=_examples,
         formatter_class=argparse.RawDescriptionHelpFormatter
     )
@@ -141,6 +188,12 @@ def main():
     parser_style_mixing_example.add_argument('--truncation-psi', type=float, help='Truncation psi (default: %(default)s)', default=0.5)
     parser_style_mixing_example.add_argument('--result-dir', help='Root directory for run results (default: %(default)s)', default='results', metavar='DIR')
 
+    parser_style_mixing_example = subparsers.add_parser('generate-means', help='Generate style mixing video')
+    parser_style_mixing_example.add_argument('--network', help='Network pickle filename', dest='network_pkl', required=True)
+    parser_style_mixing_example.add_argument('--count', help='num to make', dest='count', default=100)
+    parser_style_mixing_example.add_argument('--truncation-psi', type=float, help='Truncation psi (default: %(default)s)', default=0.5)
+    parser_style_mixing_example.add_argument('--result-dir', help='Root directory for run results (default: %(default)s)', default='results', metavar='DIR')
+
     args = parser.parse_args()
     kwargs = vars(args)
     subcmd = kwargs.pop('command')
@@ -158,7 +211,8 @@ def main():
 
     func_name_map = {
         'generate-images': 'run_generator.generate_images',
-        'style-mixing-example': 'run_generator.style_mixing_example'
+        'style-mixing-example': 'run_generator.style_mixing_example',
+        'generate-means': 'run_generator.generate_images_of_mean'
     }
     dnnlib.submit_run(sc, func_name_map[subcmd], **kwargs)
 

run_training_rosie.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+# MSOE ROSIE run StyleGAN2
+# author: gagany <[email protected]>
+
+set MCW_RESEARCH=/srv/data/mcw_research
+set TRAINING_DATA=$MCW_RESEARCH/tfrecord/0.5x
+set OUTPUT=$MCW_RESEARCH/stylegan2/training_output
+
+python run_training.py --num-gpus=8 --data-dir=$OUTPUT --config=config-f --dataset=$TRAINING_DATA --mirror-augment=false