add new files

app.py

@@ -9,43 +9,8 @@ from yarg import get
 from models.stylegan_generator import StyleGANGenerator
 from models.stylegan2_generator import StyleGAN2Generator
 
-VALID_CHOICES = [
-    "Bald",
-    "Young",
-    "Mustache",
-    "Eyeglasses",
-    "Hat",
-    "Smiling"
-]
-ENABLE_GPU = False
-MODEL_NAMES = [
-    'stylegan_ffhq',
-    'stylegan2_ffhq'
-]
-NB_IMG = 4
-OUTPUT_LIST = [
-    gr.outputs.Image(type="pil", label="Generated Image") for _ in range(NB_IMG)
-    ] + [
-    gr.outputs.Image(type="pil", label="Modified Image") for _ in range(NB_IMG)
-]
-description = """
-<p>
-<center>
-This is an interactive demo of the CVPR2020 InterfaceGAN paper, by adding other attributes such as Hat, Bald, etc. />
-</center>
-</p>
-"""
-
-def tensor_to_pil(input_object):
-    """Shows images in one figure."""
-    if isinstance(input_object, dict):
-        im_array = []
-        images = input_object['image']
-    else:
-        images = input_object
-    for _, image in enumerate(images):
-        im_array.append(PIL.Image.fromarray(image))
-    return im_array
+from utils.constants import VALID_CHOICES, ENABLE_GPU, MODEL_NAME, OUTPUT_LIST, description
+from utils.image_manip import tensor_to_pil, concat_images
 
 def get_generator(model_name):
     if model_name == 'stylegan_ffhq':
@@ -58,12 +23,16 @@ def get_generator(model_name):
         generator = generator.cuda()
     return generator
 
+generator = get_generator(MODEL_NAME)
+boundaries = {
+    boundary:np.squeeze(np.load(open(os.path.join('boundaries', MODEL_NAME, 'boundary_%s.npy' % boundary), 'rb'))) 
+    for boundary in VALID_CHOICES
+}
+
 @torch.no_grad()
-def inference(seed, choice, model_name, coef, nb_images=NB_IMG):
+def inference(seed, coef, nb_images, list_choices):
+    global generator, boundaries
     np.random.seed(seed)
-
-    boundary = np.squeeze(np.load(open(os.path.join('boundaries', model_name, 'boundary_%s.npy' % choice), 'rb')))
-    generator = get_generator(model_name)
     latent_codes = generator.easy_sample(nb_images)
     if ENABLE_GPU:
         latent_codes = latent_codes.cuda()
@@ -73,13 +42,15 @@ def inference(seed, choice, model_name, coef, nb_images=NB_IMG):
 
     new_latent_codes = latent_codes.copy()
     for i, _ in enumerate(generated_images):
-        new_latent_codes[i, :] +=  boundary*coef
+        for choice in list_choices:
+            new_latent_codes[i, :] +=  boundaries[choice]*coef
 
     modified_generated_images = generator.easy_synthesize(new_latent_codes)
     modified_generated_images = tensor_to_pil(modified_generated_images)
 
-    return generated_images + modified_generated_images
+    concatenated_output = concat_images(generated_images, modified_generated_images)
 
+    return concatenated_output
 
 iface = gr.Interface(
     fn=inference, 
@@ -88,19 +59,9 @@ iface = gr.Interface(
             minimum=0,
             maximum=1000,
             step=1,
-            default=264,
+            default=644,
             label="Random seed to use for the generation"
         ),
-        gr.inputs.Dropdown(
-            choices=VALID_CHOICES,
-            type="value",
-            label="Attribute to modify",
-        ),
-        gr.inputs.Dropdown(
-            choices=MODEL_NAMES,
-            type="value",
-            label="Model to use",
-        ),
         gr.inputs.Slider(
             minimum=-3,
             maximum=3,
@@ -108,6 +69,14 @@ iface = gr.Interface(
             default=0,
             label="Modification coefficient",
         ),
+        gr.inputs.Slider(
+            minimum=1,
+            maximum=10,
+            step=1,
+            default=2,
+            label="Number of images to generate",
+        ),
+        gr.inputs.CheckboxGroup(VALID_CHOICES, default=[], type="value", label=None, optional=False)
     ],
     outputs=OUTPUT_LIST,
     layout="horizontal",

utils/constants.py

@@ -0,0 +1,25 @@
+import gradio as gr
+
+VALID_CHOICES = [
+    "Bald",
+    "Young",
+    "Mustache",
+    "Eyeglasses",
+    "Hat",
+    "Smiling",
+    "Gender",
+    "Beard"
+]
+ENABLE_GPU = False
+MODEL_NAME = "stylegan_ffhq"
+OUTPUT_LIST = [
+    gr.outputs.Image(type="pil", label="Generated Images"),
+    gr.outputs.Image(type="pil", label="Modified Images"),
+]
+description = """
+<p>
+<center>
+This is an interactive demo of the CVPR2020 InterfaceGAN paper, by adding other attributes such as Hat, Bald, etc. 
+</center>
+</p>
+"""

utils/image_manip.py

@@ -0,0 +1,19 @@
+import numpy as np
+import PIL
+
+def concat_images(generated_images, modified_generated_images):
+    """Shows images in one figure."""
+    concatenated_array_genenerated_images = np.concatenate([np.array(image) for image in generated_images], axis=1)
+    concatenated_array_modified_generated_images = np.concatenate([np.array(image) for image in modified_generated_images], axis=1)
+
+    return [PIL.Image.fromarray(concatenated_array_genenerated_images), PIL.Image.fromarray(concatenated_array_modified_generated_images)]
+
+def tensor_to_pil(input_object):
+    if isinstance(input_object, dict):
+        im_array = []
+        images = input_object['image']
+    else:
+        images = input_object
+    for _, image in enumerate(images):
+        im_array.append(PIL.Image.fromarray(image))
+    return im_array