latest stuff added

Author: akanimax

samples/.gitignore

@@ -8,3 +8,10 @@ GAN_GEN_SHADOW_8.pth
 interpolation.mp4
 video_2.gif
 video_3.gif
+
+# ignore the new latent_space interpolation video
+new_interp.mp4
+
+frames_pro/
+frames_mine/
+M_GAN_GEN_SHADOW_8.pth

samples/celebA-HQ.png

@@ -1,63 +1,124 @@
+""" live (realtime) latent space interpolations of trained models """
+
+import argparse
 import torch as th
+import numpy as np
 import matplotlib.pyplot as plt
 from matplotlib.animation import FuncAnimation
-from pro_gan_pytorch import PRO_GAN as pg
-
-# ==========================================================================
-# Tweakable parameters
-# ==========================================================================
-depth = 8
-num_points = 12
-transition_points = 30
-# ==========================================================================
+from pro_gan_pytorch.PRO_GAN import Generator
+from torchvision.utils import make_grid
+from math import ceil, sqrt
+from scipy.ndimage import gaussian_filter
 
 # create the device for running the demo:
 device = th.device("cuda" if th.cuda.is_available() else "cpu")
 
-# load the model for the demo
-gen = th.nn.DataParallel(pg.Generator(depth=9))
-gen.load_state_dict(th.load("GAN_GEN_SHADOW_8.pth", map_location=str(device)))
 
+def parse_arguments():
+    """
+    command line arguments parser
+    :return: args => parsed command line arguments
+    """
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--generator_file", action="store", type=str,
+                        default=None, help="path to the trained generator model")
+
+    parser.add_argument("--depth", action="store", type=int,
+                        default=9, help="Depth of the network")
+
+    parser.add_argument("--latent_size", action="store", type=int,
+                        default=512, help="Latent size for the network")
+
+    parser.add_argument("--num_points", action="store", type=int,
+                        default=12, help="Number of samples to be seen")
+
+    parser.add_argument("--transition_points", action="store", type=int,
+                        default=30,
+                        help="Number of transition samples for interpolation")
+
+    parser.add_argument("--smoothing", action="store", type=float,
+                        default=1.0,
+                        help="amount of transitional smoothing")
+
+    args = parser.parse_args()
+
+    return args
+
+
+def adjust_dynamic_range(data, drange_in=(-1, 1), drange_out=(0, 1)):
+    """
+    adjust the dynamic colour range of the given input data
+    :param data: input image data
+    :param drange_in: original range of input
+    :param drange_out: required range of output
+    :return: img => colour range adjusted images
+    """
+    if drange_in != drange_out:
+        scale = (np.float32(drange_out[1]) - np.float32(drange_out[0])) / (
+                np.float32(drange_in[1]) - np.float32(drange_in[0]))
+        bias = (np.float32(drange_out[0]) - np.float32(drange_in[0]) * scale)
+        data = data * scale + bias
+    return th.clamp(data, min=0, max=1)
+
+
+def get_image(gen, point, depth, alpha):
+    """
+    obtain an All-resolution grid of images from the given point
+    :param gen: the generator object
+    :param point: random latent point for generation
+    :param depth: value of depth for image generation (0 indexed)
+    :param alpha: value of alpha for fade-in (between 0 and 1)
+    :return: img => generated image
+    """
+    image = gen(point, depth, alpha).detach()
+    image = adjust_dynamic_range(image).squeeze(dim=0)
+    return image.cpu().numpy().transpose(1, 2, 0)
 
-# function to generate an image given a latent_point
-def get_image(point):
-    img = gen(point, depth=depth, alpha=1).detach().squeeze(0).permute(1, 2, 0)
-    img = (img - img.min()) / (img.max() - img.min())
-    return img.cpu().numpy()
 
+def main(args):
+    """
+    Main function for the script
+    :param args: parsed command line arguments
+    :return: None
+    """
 
-# generate the set of points:
-fixed_points = th.randn(num_points, 512).to(device)
-fixed_points = (fixed_points / fixed_points.norm(dim=1, keepdim=True)) * (512 ** 0.5)
-points = []  # start with an empty list
-for i in range(len(fixed_points) - 1):
-    pt_1 = fixed_points[i].view(1, -1)
-    pt_2 = fixed_points[i + 1].view(1, -1)
-    direction = pt_2 - pt_1
-    for j in range(transition_points):
-        pt = pt_1 + ((direction / transition_points) * j)
-        pt = (pt / pt.norm()) * (512 ** 0.5)
-        points.append(pt)
-    # also append the final point:
-    points.append(pt_2)
+    # load the model for the demo
+    gen = th.nn.DataParallel(
+        Generator(
+            depth=args.depth,
+            latent_size=args.latent_size))
+    gen.load_state_dict(th.load(args.generator_file, map_location=str(device)))
 
-start_point = points[0]
-points = points[1:]
+    # generate the set of points:
+    total_frames = args.num_points * args.transition_points
+    all_latents = th.randn(total_frames, args.latent_size).to(device)
+    all_latents = th.from_numpy(
+        gaussian_filter(
+            all_latents.cpu(),
+            [args.smoothing * args.transition_points, 0], mode="wrap"))
+    all_latents = (all_latents /
+                   all_latents.norm(dim=-1, keepdim=True)) * sqrt(args.latent_size)
 
-fig, ax = plt.subplots()
-plt.axis("off")
-shower = plt.imshow(get_image(start_point))
+    start_point = th.unsqueeze(all_latents[0], dim=0)
+    points = all_latents[1:]
 
+    fig, ax = plt.subplots()
+    plt.axis("off")
+    shower = plt.imshow(get_image(gen, start_point, args.depth - 1, 1))
 
-def init():
-    return shower,
+    def init():
+        return shower,
 
+    def update(point):
+        shower.set_data(get_image(gen, th.unsqueeze(point, dim=0), args.depth - 1, 1))
+        return shower,
 
-def update(point):
-    shower.set_data(get_image(point))
-    return shower,
+    # define the animation function
+    ani = FuncAnimation(fig, update, frames=points, 
+                        init_func=init)
+    plt.show(ani)
 
 
-ani = FuncAnimation(fig, update, frames=points,
-                    init_func=init, blit=False)
-plt.show()
+if __name__ == '__main__':
+    main(parse_arguments())

samples/demo.py

@@ -0,0 +1,114 @@
+""" Generate single image samples from a particular depth of a model """
+
+import argparse
+import torch as th
+import numpy as np
+import os
+from torch.backends import cudnn
+from pro_gan_pytorch.PRO_GAN import Generator
+from torch.nn.functional import interpolate
+from scipy.misc import imsave
+from tqdm import tqdm
+
+# turn on the fast GPU processing mode on
+cudnn.benchmark = True
+
+
+# set the manual seed
+# th.manual_seed(3)
+
+
+def parse_arguments():
+    """
+    default command line argument parser
+    :return: args => parsed command line arguments
+    """
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--generator_file", action="store", type=str,
+                        help="pretrained weights file for generator", required=True)
+
+    parser.add_argument("--latent_size", action="store", type=int,
+                        default=256,
+                        help="latent size for the generator")
+
+    parser.add_argument("--depth", action="store", type=int,
+                        default=9,
+                        help="depth of the network. **Starts from 1")
+
+    parser.add_argument("--out_depth", action="store", type=int,
+                        default=6,
+                        help="output depth of images. **Starts from 0")
+
+    parser.add_argument("--num_samples", action="store", type=int,
+                        default=300,
+                        help="number of synchronized grids to be generated")
+
+    parser.add_argument("--out_dir", action="store", type=str,
+                        default="interp_animation_frames/",
+                        help="path to the output directory for the frames")
+
+    args = parser.parse_args()
+
+    return args
+
+
+def adjust_dynamic_range(data, drange_in=(-1, 1), drange_out=(0, 1)):
+    """
+    adjust the dynamic colour range of the given input data
+    :param data: input image data
+    :param drange_in: original range of input
+    :param drange_out: required range of output
+    :return: img => colour range adjusted images
+    """
+    if drange_in != drange_out:
+        scale = (np.float32(drange_out[1]) - np.float32(drange_out[0])) / (
+                np.float32(drange_in[1]) - np.float32(drange_in[0]))
+        bias = (np.float32(drange_out[0]) - np.float32(drange_in[0]) * scale)
+        data = data * scale + bias
+    return th.clamp(data, min=0, max=1)
+
+
+def main(args):
+    """
+    Main function for the script
+    :param args: parsed command line arguments
+    :return: None
+    """
+
+    print("Creating generator object ...")
+    # create the generator object
+    gen = th.nn.DataParallel(Generator(
+        depth=args.depth,
+        latent_size=args.latent_size
+    ))
+
+    print("Loading the generator weights from:", args.generator_file)
+    # load the weights into it
+    gen.load_state_dict(
+        th.load(args.generator_file)
+    )
+
+    # path for saving the files:
+    save_path = args.out_dir
+
+    print("Generating scale synchronized images ...")
+    for img_num in tqdm(range(1, args.num_samples + 1)):
+        # generate the images:
+        with th.no_grad():
+            point = th.randn(1, args.latent_size)
+            point = (point / point.norm()) * (args.latent_size ** 0.5)
+            ss_image = gen(point, depth=args.out_depth, alpha=1)
+            # color adjust the generated image:
+            ss_image = adjust_dynamic_range(ss_image)
+
+        # save the ss_image in the directory
+        imsave(os.path.join(save_path, str(img_num) + ".png"),
+               ss_image.squeeze(0).permute(1, 2, 0).cpu())
+
+    print("Generated %d images at %s" % (args.num_samples, save_path))
+
+
+if __name__ == '__main__':
+    main(parse_arguments())