more demo stuff added

Author: akanimax

samples/.gitignore

@@ -2,4 +2,9 @@
 pro-gan_training_video_smaller.mp4
 
 # also ignore the trained model weights
-GAN_GEN_SHADOW_8.pth
+GAN_GEN_SHADOW_8.pth
+
+# ignore some huge videos:
+interpolation.mp4
+video_2.gif
+video_3.gif

samples/demo.gif

@@ -3,8 +3,6 @@
 from matplotlib.animation import FuncAnimation
 from pro_gan_pytorch import PRO_GAN as pg
 
-th.manual_seed(1)
-
 # ==========================================================================
 # Tweakable parameters
 # ==========================================================================
@@ -30,13 +28,15 @@ def get_image(point):
 
 # 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)

samples/demo.py

@@ -0,0 +1,135 @@
+""" script for generating samples from a trained model """
+
+import argparse
+import os
+
+import torch as th
+
+# define the device for the training script
+device = th.device("cuda" if th.cuda.is_available() else "cpu")
+
+# set manual seed to 3
+th.manual_seed(3)
+
+
+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,
+                        help="pretrained weights file for generator", required=True)
+
+    parser.add_argument("--latent_size", action="store", type=int,
+                        default=512,
+                        help="latent size for the generator")
+
+    parser.add_argument("--depth", action="store", type=int,
+                        default=9, help="Value of depth for image generation (Resolution)")
+
+    parser.add_argument("--alpha", action="store", type=float,
+                        default=1, help="Value of alpha (Fade-in factor)")
+
+    parser.add_argument("--num_samples", action="store", type=int,
+                        default=64,
+                        help="number of samples in the sheet (preferably a square number)")
+
+    parser.add_argument("--time", action="store", type=float,
+                        default=1,
+                        help="Number of minutes for the video to make")
+
+    parser.add_argument("--traversal_time", action="store", type=float,
+                        default=3,
+                        help="Number of seconds to go from one point to another")
+
+    parser.add_argument("--static_time", action="store", type=float,
+                        default=1,
+                        help="Number of seconds to display a sample")
+
+    parser.add_argument("--fps", action="store", type=int,
+                        default=30, help="Frames per second in the video")
+
+    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 main(args):
+    """
+    Main function of the script
+    :param args: parsed commandline arguments
+    :return: None
+    """
+    from pro_gan_pytorch.PRO_GAN import Generator, ProGAN
+
+    # create generator object:
+    print("Creating a generator object ...")
+    generator = th.nn.DataParallel(
+        Generator(depth=args.depth,
+                  latent_size=args.latent_size).to(device))
+
+    # load the trained generator weights
+    print("loading the trained generator weights ...")
+    generator.load_state_dict(th.load(args.generator_file))
+
+    # total_frames in the video:
+    total_time_for_one_transition = args.traversal_time + args.static_time
+    total_frames_for_one_transition = (total_time_for_one_transition * args.fps)
+    number_of_transitions = int((args.time * 60) / total_time_for_one_transition)
+    total_frames = int(number_of_transitions * total_frames_for_one_transition)
+
+    # Let's create the animation video from the latent space interpolation
+    # I save the frames required for making the video here
+    points_1 = th.randn(args.num_samples, args.latent_size).to(device)
+    points_1 = (points_1 / points_1.norm(dim=1, keepdim=True)) * (args.latent_size ** 0.5)
+
+    # create output directory
+    os.makedirs(args.out_dir, exist_ok=True)
+
+    # Run the main loop for the interpolation:
+    global_frame_counter = 1  # counts number of frames
+    while global_frame_counter <= total_frames:
+        points_2 = th.randn(args.num_samples, args.latent_size).to(device)
+        points_2 = (points_2 / points_2.norm(dim=1, keepdim=True)) * (args.latent_size ** 0.5)
+        direction = points_2 - points_1
+
+        # create the points for images in this space:
+        number_of_points = int(args.traversal_time * args.fps)
+        for i in range(number_of_points):
+            points = points_1 + ((direction / number_of_points) * i)
+            points = (points / points.norm(dim=1, keepdim=True)) * (args.latent_size ** 0.5)
+
+            # generate the image for this point:
+            img = generator(points, depth=4, alpha=args.alpha)
+
+            # save the image:
+            ProGAN.create_grid(img, 1, os.path.join(args.out_dir, str(global_frame_counter) + ".png"))
+
+            # increment the counter:
+            global_frame_counter += 1
+
+        # at point_2, now add static frames:
+        img = generator(points_2, depth=4, alpha=args.alpha)
+
+        # now save the same image a number of times:
+        for _ in range(int(args.static_time * args.fps)):
+            ProGAN.create_grid(img, 1, os.path.join(args.out_dir, str(global_frame_counter) + ".png"))
+            global_frame_counter += 1
+
+        # set the point_1 := point_2
+        points_1 = points_2
+
+        print("Generated %d frames ..." % global_frame_counter)
+
+    # video frames have been generated
+    print("Video frames have been generated at:", args.out_dir)
+
+
+if __name__ == "__main__":
+    main(parse_arguments())