[ { "path": "README.md", "content": "# ProGAN\n\nImplementation of Progressive Generative Adversarial Network based on research done by Tero Karras\n\n\nThe model was trained on landscape images collected from Reddit.\n\nhttp://research.nvidia.com/sites/default/files/pubs/2017-10_Progressive-Growing-of/karras2018iclr-paper.pdf\n\n![generated images](https://github.com/perplexingpegasus/ProGAN/blob/master/example_images.png?raw=true)\n" }, { "path": "feed_dict.py", "content": "import os\nimport pickle\nimport numpy as np\nfrom itertools import cycle\n\n''' \nFeedDict handles several numpy mem_map arrays of image data saved within the directory. The arrays \nshould be named in the format \"n1_n2.npy\" where n1 x n1 is the resolution of the image data in the \narray, and n2 is its number used for indexing purposes. Data should be of type np.float32 and scaled \nbetween -1.0 and 1.0. In order to avoid loading unnecessary data into memory, only one mem_map is \nloaded at a time.\n'''\n\nclass FeedDict:\n\n pickle_filename = 'fd_log.pkl'\n\n def __init__(self, logdir, imgdir, z_length, n_examples, shuffle=True, min_size=4, max_size=1024):\n\n self.logdir = logdir\n self.shuffle = shuffle\n self.z_length = z_length\n\n self.sizes = [2 ** i for i in range(\n int(np.log2(min_size)),\n int(np.log2(max_size)) + 1\n )]\n\n files = os.listdir(imgdir)\n self.arrays = dict()\n\n for s in [2 ** i for i in range(2, 11)]:\n path_list = []\n for f in files:\n\n if f.startswith('{}_'.format(s)):\n path_list.append(os.path.join(imgdir, f))\n\n if shuffle: np.random.shuffle(path_list)\n self.arrays.update({s: cycle(path_list)})\n\n self.z_fixed = self.z_batch(n_examples, z_length)\n\n self.cur_res = None\n self.cur_path = None\n self.cur_array = None\n self.cur_array_len = 0\n self.idx = 0\n\n @property\n def n_sizes(self): return len(self.sizes)\n\n def __change_res(self, res):\n assert res in self.arrays.keys()\n self.cur_res = res\n self.__change_array()\n\n def __change_array(self):\n new_path = next(self.arrays[self.cur_res])\n print('Loaded new memmap array: {}'.format(new_path))\n if new_path != self.cur_path:\n self.cur_path = new_path\n self.cur_array = np.load(new_path)\n self.cur_array_len = self.cur_array.shape[0]\n if self.shuffle: np.random.shuffle(self.cur_array)\n self.idx = 0\n\n def z_batch(self, batch_size, random_state=None):\n if random_state is not None:\n np.random.seed(random_state)\n return np.random.normal(0.0, 1.0, size=[batch_size, self.z_length])\n\n def x_batch(self, batch_size, res):\n if res != self.cur_res:\n self.__change_res(res)\n\n remaining = self.cur_array_len - self.idx\n start = self.idx\n\n if remaining >= batch_size:\n stop = start + batch_size\n batch = self.cur_array[start:stop]\n\n else:\n stop = batch_size - remaining\n batch = self.cur_array[start:]\n self.__change_array()\n batch = np.concatenate((batch, self.cur_array[:stop]))\n\n self.idx = stop\n return batch\n\n @classmethod\n def load(cls, logdir, **kwargs):\n path = os.path.join(logdir, cls.pickle_filename)\n if os.path.exists(path):\n with open(path, 'rb') as f:\n fd = pickle.load(f)\n if type(fd) == cls:\n print('Restored feed_dict -------\\n')\n return fd\n return cls(logdir, **kwargs)\n\n def save(self):\n path = os.path.join(self.logdir, self.pickle_filename)\n with open(path, 'wb') as f:\n pickle.dump(self, f, pickle.HIGHEST_PROTOCOL)" }, { "path": "make_video.py", "content": "from progan_v15 import ProGAN\n\nimport librosa\nimport numpy as np\nfrom moviepy.video.VideoClip import VideoClip\nfrom moviepy.editor import AudioFileClip\nfrom sklearn.preprocessing import StandardScaler\n\n\ndef get_z_from_audio(audio, z_length, n_bins=60, hop_length=512, random_state=50):\n np.random.seed(random_state)\n if type(audio) == str:\n audio, sr = librosa.load(audio)\n\n y = librosa.core.cqt(audio, n_bins=n_bins, hop_length=hop_length)\n mag, phase = librosa.core.magphase(y)\n mag = mag.T\n mag = StandardScaler().fit_transform(mag)\n\n s0, s1 = mag.shape\n static = np.random.normal(size=[z_length - s1])\n static = np.tile(static, (s0, 1))\n\n z = np.concatenate((mag, static), 1)\n z = z.T\n np.random.shuffle(z)\n z = z.T\n return z\n\ndef make_video(audio, filename, progan, n_bins=60, random_state=0, imgs_per_batch=20):\n y, sr = librosa.load(audio)\n song_length = len(y) / sr\n z_audio = get_z_from_audio(y, z_length=progan.z_length, n_bins=n_bins, random_state=random_state)\n fps = z_audio.shape[0] / song_length\n res = progan.get_cur_res()\n shape = (res, res * 16 // 9, 3)\n\n imgs = np.zeros(shape=[imgs_per_batch, *shape], dtype=np.float32)\n\n def make_frame(t):\n global imgs\n cur_frame_idx = int(t * fps)\n\n if cur_frame_idx >= len(z_audio):\n return np.zeros(shape=shape, dtype=np.uint8)\n\n if cur_frame_idx % imgs_per_batch == 0:\n imgs = progan.generate(z_audio[cur_frame_idx:cur_frame_idx + imgs_per_batch])\n imgs = imgs[:, :, :res * 8 // 9, :]\n imgs_rev = np.flip(imgs, 2)\n imgs = np.concatenate((imgs, imgs_rev), 2)\n\n return imgs[cur_frame_idx % imgs_per_batch]\n\n video_clip = VideoClip(make_frame=make_frame, duration=song_length)\n audio_clip = AudioFileClip(audio)\n video_clip = video_clip.set_audio(audio_clip)\n video_clip.write_videofile(filename, fps=fps)\n\nif __name__ == '__main__':\n progan = ProGAN(\n logdir='logdir_v2',\n imgdir='img_arrays',\n )\n make_video('videos\\\\eco_zones.mp3', 'eco_zones.mp4', progan, random_state=768)" }, { "path": "ops.py", "content": "import tensorflow as tf\n\n\nweight_init = tf.random_normal_initializer()\nbias_init = tf.constant_initializer(0)\n\n\ndef conv(input, out_channels, filter_size=3, k=1, padding='SAME', mode=None, output_shape=None):\n\n in_shape = tf.shape(input)\n input_channels = int(input.get_shape()[1])\n\n if mode == 'upscale' or mode == 'transpose':\n filter_shape = [filter_size, filter_size, out_channels, input_channels]\n else:\n filter_shape = [filter_size, filter_size, input_channels, out_channels]\n\n filter = tf.get_variable('filter', filter_shape, initializer=weight_init)\n fan_in = float(filter_size ** 2 * input_channels)\n filter = filter * tf.sqrt(2.0 / fan_in)\n\n b = tf.get_variable('bias', [1, out_channels, 1, 1], initializer=bias_init)\n\n if mode == 'upscale':\n filter = tf.pad(filter, [[1, 1], [1, 1], [0, 0], [0, 0]], mode='CONSTANT')\n filter = tf.add_n([filter[1:, 1:], filter[:-1, 1:], filter[1:, :-1], filter[:-1, :-1]])\n output_shape = [in_shape[0], out_channels, in_shape[2] * 2, in_shape[3] * 2]\n output = tf.nn.conv2d_transpose(input, filter, output_shape, [1, 1, 2, 2],\n padding=padding, data_format='NCHW')\n\n elif mode == 'downscale':\n filter = tf.pad(filter, [[1, 1], [1, 1], [0, 0], [0, 0]], mode='CONSTANT')\n filter = tf.add_n([filter[1:, 1:], filter[:-1, 1:], filter[1:, :-1], filter[:-1, :-1]])\n filter *= 0.25\n output = tf.nn.conv2d(input, filter, [1, 1, 2, 2], padding=padding, data_format='NCHW')\n\n elif mode == 'transpose':\n output = tf.nn.conv2d_transpose(input, filter, output_shape, [1, 1, k, k],\n padding=padding, data_format='NCHW')\n\n else:\n output = tf.nn.conv2d(input, filter, [1, 1, k, k], padding=padding, data_format='NCHW')\n\n output += b\n\n if out_channels == 1:\n output = tf.squeeze(output, 3)\n\n return output\n\n\ndef dense(input, output_size):\n fan_in = int(input.get_shape()[1])\n W = tf.get_variable('W', [fan_in, output_size], initializer=weight_init)\n W = W * tf.sqrt(2.0 / float(fan_in))\n b = tf.get_variable('b', [1, output_size, 1, 1], initializer=bias_init)\n return tf.matmul(input, W) + b\n\n\ndef leaky_relu(input, alpha=0.2):\n return tf.nn.leaky_relu(input, alpha=alpha)\n\n\ndef pixelwise_norm(input):\n pixel_var = tf.reduce_mean(tf.square(input), 1, keepdims=True)\n return input / tf.sqrt(pixel_var + 1e-8)\n\n\ndef g_conv_layer(input, out_channels, **kwargs):\n return pixelwise_norm(leaky_relu(conv(input, out_channels, **kwargs)))\n\n\ndef d_conv_layer(input, out_channels, **kwargs):\n return leaky_relu(conv(input, out_channels, **kwargs))\n\n\ndef minibatch_stddev(input):\n shape = tf.shape(input)\n group_size = tf.minimum(4, shape[0])\n x = tf.reshape(input, [group_size, -1, shape[1], shape[2], shape[3]])\n\n mu = tf.reduce_mean(x, axis=0, keepdims=True)\n sigma = tf.sqrt(tf.reduce_mean(tf.square(x - mu), axis=0) + 1e-8)\n\n sigma_avg = tf.reduce_mean(sigma, axis=[1, 2, 3], keepdims=True)\n sigma_avg = tf.tile(sigma_avg, [group_size, 1, shape[2], shape[3]])\n return tf.concat((input, sigma_avg), axis=1)\n\n\ndef upscale(input):\n shape = tf.shape(input)\n channels = input.get_shape()[1]\n output = tf.reshape(input, [-1, channels, shape[2], 1, shape[3], 1])\n output = tf.tile(output, [1, 1, 1, 2, 1, 2])\n return tf.reshape(output, [-1, channels, shape[2] * 2, shape[3] * 2])\n\n\ndef downscale(input):\n return tf.nn.avg_pool(input, ksize=[1, 1, 2, 2], strides=[1, 1, 2, 2],\n padding='SAME', data_format='NCHW')\n\n\ndef resize_images(input, dims=None):\n if dims is None:\n dims = tf.shape(input)[2] * 2, tf.shape(input)[3] * 2\n return tf.image.resize_nearest_neighbor(input, dims)\n\n\ndef scale_uint8(input):\n input = tf.to_float(input)\n return (input / 127.5) - 1\n\n\ndef tensor_to_imgs(input, switch_dims=True):\n if switch_dims: input = tf.transpose(input, (0, 2, 3, 1))\n imgs = tf.minimum(tf.maximum(input, -tf.ones_like(input)), tf.ones_like(input))\n imgs = (imgs + 1) * 127.5\n return tf.cast(imgs, tf.uint8)" }, { "path": "progan_v15.py", "content": "import os\nimport datetime as dt\n\n# Operations used in building the network. Many are not used in the current model\nfrom ops import *\n# FeedDict object used to continuously provide new training data\nfrom feed_dict import FeedDict\n\n\n# TODO: add argparser and flags\n# TODO: refactor training function\n# TODO: train next version of model using reset_optimizer=True\n\n\nclass ProGAN:\n def __init__(self,\n logdir, # directory of stored models\n imgdir, # directory of images for FeedDict\n learning_rate=0.001, # Adam optimizer learning rate\n beta1=0, # Adam optimizer beta1\n beta2=0.99, # Adam optimizer beta2\n w_lambda=10.0, # WGAN-GP/LP lambda\n w_gamma=1.0, # WGAN-GP/LP gamma\n epsilon=0.001, # WGAN-GP/LP lambda\n z_length=512, # latent variable size\n n_imgs=800000, # number of images to show in each growth step\n batch_repeats=1, # number of times to repeat minibatch\n n_examples=24, # number of example images to generate\n lipschitz_penalty=True, # if True, use WGAN-LP instead of WGAN-GP\n big_image=True, # Generate a single large preview image, only works if n_examples = 24\n scaling_factor=None, # factor to scale down number of trainable parameters\n reset_optimizer=False, # reset optimizer variables with each new layer\n ):\n\n # Scale down the number of factors if scaling_factor is provided\n self.channels = [512, 512, 512, 512, 256, 128, 64, 32, 16, 8]\n if scaling_factor:\n assert scaling_factor > 1\n self.channels = [max(4, c // scaling_factor) for c in self.channels]\n\n self.batch_size = [16, 16, 16, 16, 16, 16, 8, 4, 3]\n self.z_length = z_length\n self.n_examples = n_examples\n self.batch_repeats = batch_repeats if batch_repeats else 1\n self.n_imgs = n_imgs\n self.logdir = logdir\n self.big_image = big_image\n self.w_lambda = w_lambda\n self.w_gamma = w_gamma\n self.epsilon = epsilon\n self.reset_optimizer=reset_optimizer\n self.lipschitz_penalty = lipschitz_penalty\n self.start = True\n\n # Generate fized latent variables for image previews\n np.random.seed(0)\n self.z_fixed = np.random.normal(size=[self.n_examples, self.z_length])\n\n # Initialize placeholders\n self.x_placeholder = tf.placeholder(tf.float32, [None, None, None, 3])\n self.z_placeholder = tf.placeholder(tf.float32, [None, self.z_length])\n\n # Global step\n with tf.variable_scope('global_step'):\n self.global_step = tf.Variable(0, name='global_step', trainable=False)\n self.global_step_op = tf.assign(self.global_step, tf.add(self.global_step, 1))\n\n # Non-trainable variables for counting to next layer and incrementing value of alpha\n with tf.variable_scope('image_count'):\n self.total_imgs = tf.Variable(0.0, name='image_step', trainable=False)\n self.img_count_placeholder = tf.placeholder(tf.float32)\n self.img_step_op = tf.assign(self.total_imgs,\n tf.add(self.total_imgs, self.img_count_placeholder))\n\n self.img_step = tf.mod(tf.add(self.total_imgs, self.n_imgs), self.n_imgs * 2)\n self.alpha = tf.minimum(1.0, tf.div(self.img_step, self.n_imgs))\n self.layer = tf.floor_div(tf.add(self.total_imgs, self.n_imgs), self.n_imgs * 2)\n\n # Initialize optimizer as member variable if not rest_optimizer, otherwise generate new\n # optimizer for each layer\n if self.reset_optimizer:\n self.lr = learning_rate\n self.beta1 = beta1\n self.beta2 = beta2\n else:\n self.g_optimizer = tf.train.AdamOptimizer(learning_rate, beta1, beta2)\n self.d_optimizer = tf.train.AdamOptimizer(learning_rate, beta1, beta2)\n\n # Initialize FeedDict\n self.feed = FeedDict.load(imgdir, logdir)\n self.n_layers = int(np.log2(1024)) - 1\n self.networks = [self._create_network(i + 1) for i in range(self.n_layers)]\n\n # Initialize Session, FileWriter and Saver\n self.sess = tf.Session()\n self.sess.run(tf.global_variables_initializer())\n self.writer = tf.summary.FileWriter(self.logdir, graph=self.sess.graph)\n self.saver = tf.train.Saver()\n\n # Look in logdir to see if a saved model already exists. If so, load it\n try:\n self.saver.restore(self.sess, tf.train.latest_checkpoint(self.logdir))\n print('Restored ----------------\\n')\n except Exception:\n pass\n\n # Function for fading input of current layer into previous layer based on current value of alpha\n def _reparameterize(self, x0, x1):\n return tf.add(\n tf.scalar_mul(tf.subtract(1.0, self.alpha), x0),\n tf.scalar_mul(self.alpha, x1)\n )\n\n # Function for creating network layout at each layer\n def _create_network(self, layers):\n\n # Build the generator for this layer\n def generator(z):\n with tf.variable_scope('Generator'):\n with tf.variable_scope('latent_vector'):\n z = tf.expand_dims(z, 1)\n g1 = tf.expand_dims(z, 2)\n for i in range(layers):\n with tf.variable_scope('layer_{}'.format(i)):\n if i > 0:\n g1 = resize(g1)\n if i == layers - 1 and layers > 1:\n g0 = g1\n with tf.variable_scope('1'):\n if i == 0:\n g1 = pixelwise_norm(leaky_relu(conv2d_transpose(\n g1, [tf.shape(g1)[0], 4, 4, self.channels[0]])))\n else:\n g1 = pixelwise_norm(leaky_relu(conv2d(g1, self.channels[i])))\n with tf.variable_scope('2'):\n g1 = pixelwise_norm(leaky_relu(conv2d(g1, self.channels[i])))\n with tf.variable_scope('rgb_layer_{}'.format(layers - 1)):\n g1 = conv2d(g1, 3, 1, weight_norm=False)\n if layers > 1:\n with tf.variable_scope('rgb_layer_{}'.format(layers - 2)):\n g0 = conv2d(g0, 3, 1, weight_norm=False)\n g = self._reparameterize(g0, g1)\n else:\n g = g1\n return g\n\n # Build the discriminator for this layer\n def discriminator(x):\n with tf.variable_scope('Discriminator'):\n if layers > 1:\n with tf.variable_scope('rgb_layer_{}'.format(layers - 2)):\n d0 = avg_pool(x)\n d0 = leaky_relu(conv2d(d0, self.channels[layers - 1], 1))\n with tf.variable_scope('rgb_layer_{}'.format(layers - 1)):\n d1 = leaky_relu(conv2d(x, self.channels[layers], 1))\n for i in reversed(range(layers)):\n with tf.variable_scope('layer_{}'.format(i)):\n if i == 0:\n d1 = minibatch_stddev(d1)\n with tf.variable_scope('1'):\n d1 = leaky_relu(conv2d(d1, self.channels[i]))\n with tf.variable_scope('2'):\n if i == 0:\n d1 = leaky_relu(conv2d(d1, self.channels[i], 4, padding='VALID'))\n else:\n d1 = leaky_relu(conv2d(d1, self.channels[i]))\n if i != 0:\n d1 = avg_pool(d1)\n if i == layers - 1 and layers > 1:\n d1 = self._reparameterize(d0, d1)\n with tf.variable_scope('dense'):\n d = tf.reshape(d1, [-1, self.channels[0]])\n d = dense_layer(d, 1)\n return d\n\n # image dimensions\n dim = 2 ** (layers + 1)\n\n # Build the current network\n with tf.variable_scope('Network', reuse=tf.AUTO_REUSE):\n Gz = generator(self.z_placeholder)\n Dz = discriminator(Gz)\n\n # Mix different resolutions of input images according to value of alpha\n with tf.variable_scope('reshape'):\n if layers > 1:\n x0 = resize(self.x_placeholder, (dim // 2, dim // 2))\n x0 = resize(x0, (dim, dim))\n x1 = resize(self.x_placeholder, (dim, dim))\n x = self._reparameterize(x0, x1)\n else:\n x = resize(self.x_placeholder, (dim, dim))\n Dx = discriminator(x)\n\n # Fake and real image mixing for WGAN-GP loss function\n interp = tf.random_uniform(shape=[tf.shape(Dz)[0], 1, 1, 1], minval=0., maxval=1.)\n x_hat = interp * x + (1 - interp) * Gz\n Dx_hat = discriminator(x_hat)\n\n # Loss function and scalar summaries\n with tf.variable_scope('Loss_Function'):\n\n # Wasserstein Distance\n wd = Dz - Dx\n\n # Gradient/Lipschitz Penalty\n grads = tf.gradients(Dx_hat, [x_hat])[0]\n slopes = tf.sqrt(tf.reduce_sum(tf.square(grads), [1, 2, 3]))\n if self.lipschitz_penalty:\n gp = tf.square(tf.maximum((slopes - self.w_gamma) / self.w_gamma, 0))\n else:\n gp = tf.square((slopes - self.w_gamma) / self.w_gamma)\n gp_scaled = self.w_lambda * gp\n\n # Epsilon penalty keeps discriminator output for drifting too far away from zero\n epsilon_cost = self.epsilon * tf.square(Dx)\n\n # Cost and summary scalars\n g_cost = tf.reduce_mean(-Dz)\n d_cost = tf.reduce_mean(wd + gp_scaled + epsilon_cost)\n wd = tf.abs(tf.reduce_mean(wd))\n gp = tf.reduce_mean(gp)\n\n # Summaries\n wd_sum = tf.summary.scalar('Wasserstein_distance_{}x{}'.format(dim, dim), wd)\n gp_sum = tf.summary.scalar('gradient_penalty_{}x{}'.format(dim, dim), gp)\n\n # Collecting variables to be trained by optimizers\n g_vars, d_vars = [], []\n var_scopes = ['layer_{}'.format(i) for i in range(layers)]\n var_scopes.extend(['dense', 'rgb_layer_{}'.format(layers - 1), 'rgb_layer_{}'.format(layers - 2)])\n for scope in var_scopes:\n g_vars.extend(tf.get_collection(\n tf.GraphKeys.GLOBAL_VARIABLES,\n scope='Network/Generator/{}'.format(scope)))\n d_vars.extend(tf.get_collection(\n tf.GraphKeys.GLOBAL_VARIABLES,\n scope='Network/Discriminator/{}'.format(scope)))\n\n # Generate optimizer operations\n # if self.reset_optimizer is True then initialize a new optimizer for each layer\n with tf.variable_scope('Optimize'):\n if self.reset_optimizer:\n g_train = tf.train.AdamOptimizer(\n self.lr, self.beta1, self.beta2, name='G_optimizer_{}'.format(layers - 1)).minimize(\n g_cost, var_list=g_vars)\n d_train = tf.train.AdamOptimizer(\n self.lr, self.beta1, self.beta2, name='D_optimizer_{}'.format(layers - 1)).minimize(\n d_cost, var_list=d_vars)\n else:\n g_train = self.g_optimizer.minimize(g_cost, var_list=g_vars)\n d_train = self.d_optimizer.minimize(d_cost, var_list=d_vars)\n\n # Print variable names to before running model\n print([var.name for var in g_vars])\n print([var.name for var in d_vars])\n\n # Generate preview images\n with tf.variable_scope('image_preview'):\n fake_imgs = tf.minimum(tf.maximum(Gz, -tf.ones_like(Gz)), tf.ones_like(Gz))\n real_imgs = x[:min(self.batch_size[layers - 1], 4), :, :, :]\n\n # Upsize images to normal visibility\n if dim < 256:\n fake_imgs = resize(fake_imgs, (256, 256))\n real_imgs = resize(real_imgs, (256, 256))\n\n # Concatenate images into one large image for preview, only used if 24 preview images are requested\n if self.big_image and self.n_examples == 24:\n fake_img_list = tf.unstack(fake_imgs, num=24)\n fake_img_list = [tf.concat(fake_img_list[6 * i:6 * (i + 1)], 1) for i in range(4)]\n fake_imgs = tf.concat(fake_img_list, 0)\n fake_imgs = tf.expand_dims(fake_imgs, 0)\n\n real_img_list = tf.unstack(real_imgs, num=min(self.batch_size[layers - 1], 4))\n real_imgs = tf.concat(real_img_list, 1)\n real_imgs = tf.expand_dims(real_imgs, 0)\n\n # images summaries\n fake_img_sum = tf.summary.image('fake{}x{}'.format(dim, dim), fake_imgs, self.n_examples)\n real_img_sum = tf.summary.image('real{}x{}'.format(dim, dim), real_imgs, 4)\n\n return (dim, wd, gp, wd_sum, gp_sum, g_train, d_train,\n fake_img_sum, real_img_sum, Gz, discriminator)\n\n # Summary adding function\n def _add_summary(self, string, gs):\n self.writer.add_summary(string, gs)\n\n # Latent variable 'z' generator\n def _z(self, batch_size):\n return np.random.normal(0.0, 1.0, [batch_size, self.z_length])\n\n # Main training function\n def train(self):\n prev_layer = None\n start_time = dt.datetime.now()\n total_imgs = self.sess.run(self.total_imgs)\n\n while total_imgs < (self.n_layers - 0.5) * self.n_imgs * 2:\n\n # Get current layer, global step, alpha and total number of images used so far\n layer, gs, img_step, alpha, total_imgs = self.sess.run([\n self.layer, self.global_step, self.img_step, self.alpha, self.total_imgs])\n layer = int(layer)\n\n # Global step interval to save model and generate image previews\n save_interval = max(1000, 10000 // 2 ** layer)\n\n # Get network operations and loss functions for current layer\n (dim, wd, gp, wd_sum, gp_sum, g_train, d_train,\n fake_img_sum, real_img_sum, Gz, discriminator) = self.networks[layer]\n\n # Get training data and latent variables to store in feed_dict\n feed_dict = {self.x_placeholder: self.feed.next_batch(self.batch_size[layer], dim),\n self.z_placeholder: self._z(self.batch_size[layer])}\n\n # Reset start times if a new layer has begun training\n if layer != prev_layer:\n start_time = dt.datetime.now()\n\n # Here's where we actually train the model\n for _ in range(self.batch_repeats):\n self.sess.run(g_train, feed_dict)\n self.sess.run(d_train, feed_dict)\n\n # Get loss values and summaries\n wd_, gp_, wd_sum_str, gp_sum_str = self.sess.run([wd, gp, wd_sum, gp_sum], feed_dict)\n\n # Print current status, loss functions, etc.\n percent_done = np.round(img_step * 50 / self.n_imgs, 4)\n imgs_done = int(img_step)\n cur_layer_imgs = self.n_imgs * 2\n if dim == 4:\n percent_done = np.round((percent_done - 50) * 2, 4)\n imgs_done -= self.n_imgs\n cur_layer_imgs //= 2\n print('dimensions: {}x{} ---- {}% ---- images: {}/{} ---- alpha: {} ---- global step: {}'\n '\\nWasserstein distance: {}\\ngradient penalty: {}\\n'.format(\n dim, dim, percent_done, imgs_done, cur_layer_imgs, alpha, gs, wd_, gp_))\n\n # Log scalar data every 20 global steps\n if gs % 20 == 0:\n self._add_summary(wd_sum_str, gs)\n self._add_summary(gp_sum_str, gs)\n\n # Operations to run every save interval\n if gs % save_interval == 0:\n\n # Do not save the model or generate images immediately after loading/preloading\n if self.start:\n self.start = False\n\n # Save the model and generate image previews\n else:\n print('saving and making images...\\n')\n self.feed.save()\n self.saver.save(\n self.sess, os.path.join(self.logdir, \"model.ckpt\"),\n global_step=self.global_step)\n real_img_sum_str = self.sess.run(real_img_sum, feed_dict)\n img_preview_feed_dict = {\n self.x_placeholder: feed_dict[self.x_placeholder][:4],\n self.z_placeholder: self.z_fixed}\n fake_img_sum_str = self.sess.run(fake_img_sum, img_preview_feed_dict)\n self._add_summary(fake_img_sum_str, gs)\n self._add_summary(real_img_sum_str, gs)\n\n # Increment image count and global step variables\n img_count = self.batch_repeats * self.batch_size[layer]\n self.sess.run(self.global_step_op)\n self.sess.run(self.img_step_op, {self.img_count_placeholder: img_count})\n\n # Calculate and print estimated time remaining\n prev_layer = layer\n avg_time = (dt.datetime.now() - start_time) / (imgs_done + self.batch_size[layer])\n steps_remaining = cur_layer_imgs - imgs_done\n time_reamining = avg_time * steps_remaining\n print('est. time remaining on current layer: {}'.format(time_reamining))\n\n def get_cur_res(self):\n cur_layer = int(self.sess.run(self.layer))\n return 2 ** (2 + cur_layer)\n\n # Function for generating images from a 1D or 2D array of latent vectors\n def generate(self, z):\n if len(z.shape) == 1:\n z = np.expand_dims(z, 0)\n\n cur_layer = int(self.sess.run(self.layer))\n G = self.networks[cur_layer][9]\n imgs = self.sess.run(G, {self.z_placeholder: z})\n\n imgs = np.minimum(imgs, 1.0)\n imgs = np.maximum(imgs, -1.0)\n imgs = (imgs + 1) * 255 / 2\n imgs = np.uint8(imgs)\n\n if imgs.shape[0] == 1:\n imgs = np.squeeze(imgs, 0)\n return imgs\n\n\n def transform(self, input_img, n_iter=100000):\n with tf.variable_scope('transform'):\n global_step = tf.Variable(0, name='transform_global_step', trainable=False)\n transform_img = tf.Variable(input_img, name='transform_img', dtype=tf.float32)\n\n cur_layer = int(self.sess.run(self.layer))\n (dim, wd, gp, wd_sum, gp_sum, g_train, d_train,\n ake_img_sum, real_img_sum, Gz, discriminator) = self.networks[cur_layer]\n\n with tf.variable_scope('Network', reuse=tf.AUTO_REUSE):\n with tf.variable_scope('resize'):\n jitter = tf.random_uniform([2], -10, 10, tf.int32)\n img = tf.manip.roll(transform_img, jitter, [1, 2])\n img = resize(img, (dim, dim))\n Dt = discriminator(img)\n\n t_cost = tf.reduce_mean(-Dt)\n tc_sum = tf.summary.scalar('transform_cost_{}x{}'.format(dim, dim), t_cost)\n t_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='transform/transform_img')\n t_train = tf.train.AdamOptimizer(0.0001).minimize(\n t_cost, var_list=t_vars, global_step=global_step)\n transform_img_sum = tf.summary.image('transform', transform_img)\n\n self.sess.run(tf.global_variables_initializer())\n\n for i in range(n_iter):\n gs, t_cost_, tc_sum_str, _ = self.sess.run([global_step, t_cost, tc_sum, t_train])\n print('Global step: {}, cost: {}\\n\\n'.format(gs, t_cost_))\n if i % 20 == 0:\n self._add_summary(tc_sum_str, gs)\n if i % 1000 == 0:\n img_sum_str = self.sess.run(transform_img_sum)\n self._add_summary(img_sum_str, gs)\n\n\nif __name__ == '__main__':\n\n progan = ProGAN(\n logdir='logdir_v2',\n imgdir='img_arrays',\n )\n # progan = ProGAN(\n # logdir='logdir_v3',\n # imgdir='img_arrays_botanical',\n # reset_optimizer=True\n # )\n progan.train()" }, { "path": "progan_v16.py", "content": "import datetime as dt\nimport os\n\nimport numpy as np\n\n# Operations used in building the network. Many are not used in the current model\nfrom ops import *\n# FeedDict object used to continuously provide new training data\nfrom feed_dict import FeedDict\n\n\n# TODO: add argparser and flags\n\n\nclass ProGAN:\n def __init__(self,\n logdir, # directory of stored models\n imgdir, # directory of images for FeedDict\n learning_rate=0.001, # Adam optimizer learning rate\n beta1=0, # Adam optimizer beta1\n beta2=0.99, # Adam optimizer beta2\n w_lambda=10.0, # WGAN-GP/LP lambda\n w_gamma=1.0, # WGAN-GP/LP gamma\n epsilon=0.001, # WGAN-GP/LP lambda\n z_length=512, # latent variable size\n n_imgs=800000, # number of images to show in each growth step\n batch_repeats=1, # number of times to repeat minibatch\n n_examples=24, # number of example images to generate\n lipschitz_penalty=True, # if True, use WGAN-LP instead of WGAN-GP\n big_image=True, # Generate a single large preview image, only works if n_examples = 24\n reset_optimizer=True, # reset optimizer variables with each new layer\n batch_sizes=None,\n channels=None,\n ):\n\n # Scale down the number of factors if scaling_factor is provided\n self.channels = channels if channels else [512, 512, 512, 512, 256, 128, 64, 32, 16, 16]\n self.batch_sizes = batch_sizes if batch_sizes else [16, 16, 16, 16, 16, 16, 12, 4, 3]\n\n self.z_length = z_length\n self.n_examples = n_examples\n self.batch_repeats = batch_repeats if batch_repeats else 1\n self.n_imgs = n_imgs\n self.logdir = logdir\n self.big_image = big_image\n self.w_lambda = w_lambda\n self.w_gamma = w_gamma\n self.epsilon = epsilon\n self.reset_optimizer=reset_optimizer\n self.lipschitz_penalty = lipschitz_penalty\n\n # Initialize FeedDict\n self.feed = FeedDict.load(logdir, imgdir=imgdir, z_length=z_length, n_examples=n_examples)\n self.n_layers = self.feed.n_sizes\n self.max_imgs = (self.n_layers - 0.5) * self.n_imgs * 2\n\n # Initialize placeholders\n self.x_placeholder = tf.placeholder(tf.uint8, [None, 3, None, None])\n self.z_placeholder = tf.placeholder(tf.float32, [None, self.z_length])\n\n # Global step\n with tf.variable_scope('global_step'):\n self.global_step = tf.Variable(0, name='global_step', trainable=False, dtype=tf.int32)\n\n # Non-trainable variables for counting to next layer and incrementing value of alpha\n with tf.variable_scope('image_count'):\n self.total_imgs = tf.Variable(0, name='total_images', trainable=False, dtype=tf.int32)\n\n img_offset = tf.add(self.total_imgs, self.n_imgs)\n imgs_per_layer = self.n_imgs * 2\n\n self.img_step = tf.mod(img_offset, imgs_per_layer)\n self.layer = tf.minimum(tf.floor_div(img_offset, imgs_per_layer), self.n_layers - 1)\n\n fade_in = tf.to_float(self.img_step) / float(self.n_imgs)\n self.alpha = tf.minimum(1.0, tf.maximum(0.0, fade_in))\n\n # Initialize optimizer as member variable if not rest_optimizer, otherwise generate new\n # optimizer for each layer\n if self.reset_optimizer:\n self.lr = learning_rate\n self.beta1 = beta1\n self.beta2 = beta2\n else:\n self.g_optimizer = tf.train.AdamOptimizer(learning_rate, beta1, beta2)\n self.d_optimizer = tf.train.AdamOptimizer(learning_rate, beta1, beta2)\n self.networks = [self.create_network(i + 1) for i in range(self.n_layers)]\n\n # Initialize Session, FileWriter and Saver\n self.sess = tf.Session()\n self.sess.run(tf.global_variables_initializer())\n self.writer = tf.summary.FileWriter(self.logdir, graph=self.sess.graph)\n self.saver = tf.train.Saver()\n\n # Look in logdir to see if a saved model already exists. If so, load it\n try:\n self.saver.restore(self.sess, tf.train.latest_checkpoint(self.logdir))\n print('Restored model -----------\\n')\n except Exception:\n pass\n\n\n # Function for fading input of current layer into previous layer based on current value of alpha\n def reparameterize(self, x0, x1):\n return tf.add(\n tf.scalar_mul(tf.subtract(1.0, self.alpha), x0),\n tf.scalar_mul(self.alpha, x1)\n )\n\n\n # Build a generator for n layers\n def generator(self, z, n_layers):\n with tf.variable_scope('Generator'):\n\n with tf.variable_scope('latent_vector'):\n z = tf.expand_dims(z, 2)\n g1 = tf.expand_dims(z, 3)\n\n for i in range(n_layers):\n with tf.variable_scope('layer_{}'.format(i)):\n\n if i == n_layers - 1:\n g0 = g1\n\n with tf.variable_scope('1'):\n if i == 0:\n g1 = g_conv_layer(g1, self.channels[i],\n filter_size=4, padding='VALID', mode='transpose',\n output_shape=[tf.shape(g1)[0], self.channels[i], 4, 4])\n else:\n g1 = g_conv_layer(g1, self.channels[i], mode='upscale')\n\n with tf.variable_scope('2'):\n g1 = g_conv_layer(g1, self.channels[i])\n\n with tf.variable_scope('rgb_layer_{}'.format(n_layers - 1)):\n g1 = conv(g1, 3, filter_size=1)\n\n if n_layers > 1:\n with tf.variable_scope('rgb_layer_{}'.format(n_layers - 2)):\n g0 = conv(g0, 3, filter_size=1)\n g0 = upscale(g0)\n g = self.reparameterize(g0, g1)\n else:\n g = g1\n\n return g\n\n\n # Build a discriminator n layers\n def discriminator(self, x, n_layers):\n with tf.variable_scope('Discriminator'):\n\n if n_layers > 1:\n with tf.variable_scope('rgb_layer_{}'.format(n_layers - 2)):\n d0 = downscale(x)\n d0 = d_conv_layer(d0, self.channels[n_layers - 1], filter_size=1)\n\n with tf.variable_scope('rgb_layer_{}'.format(n_layers - 1)):\n d1 = d_conv_layer(x, self.channels[n_layers], filter_size=1)\n\n for i in reversed(range(n_layers)):\n with tf.variable_scope('layer_{}'.format(i)):\n\n if i == 0:\n d1 = minibatch_stddev(d1)\n\n with tf.variable_scope('1'):\n d1 = d_conv_layer(d1, self.channels[i])\n\n with tf.variable_scope('2'):\n if i == 0:\n d1 = d_conv_layer(d1, self.channels[0],\n filter_size=4, padding='VALID')\n else:\n d1 = d_conv_layer(d1, self.channels[i], mode='downscale')\n\n if i == n_layers - 1 and n_layers > 1:\n d1 = self.reparameterize(d0, d1)\n\n with tf.variable_scope('dense'):\n d = tf.reshape(d1, [-1, self.channels[0]])\n d = dense(d, 1)\n\n return d\n\n\n # Function for creating network layout at each layer\n def create_network(self, n_layers):\n\n # image dimensions\n dim = 2 ** (n_layers + 1)\n\n # Build the current network\n with tf.variable_scope('Network', reuse=tf.AUTO_REUSE):\n Gz = self.generator(self.z_placeholder, n_layers)\n Dz = self.discriminator(Gz, n_layers)\n\n # Mix different resolutions of input images according to value of alpha\n with tf.variable_scope('training_images'):\n x = scale_uint8(self.x_placeholder)\n if n_layers > 1:\n x0 = upscale(downscale(x))\n x1 = x\n x = self.reparameterize(x0, x1)\n\n Dx = self.discriminator(x, n_layers)\n\n # Fake and real image mixing for WGAN-GP loss function\n interp = tf.random_uniform(shape=[tf.shape(Dz)[0], 1, 1, 1], minval=0.0, maxval=1.0)\n x_hat = interp * x + (1 - interp) * Gz\n Dx_hat = self.discriminator(x_hat, n_layers)\n\n # Loss function and scalar summaries\n with tf.variable_scope('Loss_Function'):\n\n # Wasserstein Distance\n wd = Dz - Dx\n\n # Gradient/Lipschitz Penalty\n grads = tf.gradients(Dx_hat, [x_hat])[0]\n slopes = tf.sqrt(tf.reduce_sum(tf.square(grads), [1, 2, 3]))\n\n if self.lipschitz_penalty:\n gp = tf.square(tf.maximum((slopes - self.w_gamma) / self.w_gamma, 0))\n else:\n gp = tf.square((slopes - self.w_gamma) / self.w_gamma)\n\n gp_scaled = self.w_lambda * gp\n\n # Epsilon penalty keeps discriminator output for drifting too far away from zero\n epsilon_cost = self.epsilon * tf.square(Dx)\n\n # Cost and summary scalars\n g_cost = tf.reduce_mean(-Dz)\n d_cost = tf.reduce_mean(wd + gp_scaled + epsilon_cost)\n wd = tf.abs(tf.reduce_mean(wd))\n gp = tf.reduce_mean(gp)\n\n # Summaries\n wd_sum = tf.summary.scalar('Wasserstein_distance_{}_({}x{})'.format(\n n_layers - 1, dim, dim), wd)\n gp_sum = tf.summary.scalar('gradient_penalty_{}_({}x{})'.format(\n n_layers - 1, dim, dim), gp)\n\n # Collecting variables to be trained by optimizers\n g_vars, d_vars = [], []\n var_scopes = ['layer_{}'.format(i) for i in range(n_layers)]\n var_scopes.extend([\n 'dense',\n 'rgb_layer_{}'.format(n_layers - 2),\n 'rgb_layer_{}'.format(n_layers - 1)\n ])\n\n for scope in var_scopes:\n g_vars.extend(tf.get_collection(\n tf.GraphKeys.GLOBAL_VARIABLES, scope='Network/Generator/{}'.format(scope)\n ))\n d_vars.extend(tf.get_collection(\n tf.GraphKeys.GLOBAL_VARIABLES, scope='Network/Discriminator/{}'.format(scope)\n ))\n\n # Generate optimizer operations\n # if self.reset_optimizer is True then initialize a new optimizer for each layer\n with tf.variable_scope('Optimize'):\n if self.reset_optimizer:\n g_train = tf.train.AdamOptimizer(\n self.lr, self.beta1, self.beta2, name='G_optimizer_{}'.format(n_layers - 1)\n ).minimize(\n g_cost, var_list=g_vars)\n d_train = tf.train.AdamOptimizer(\n self.lr, self.beta1, self.beta2, name='D_optimizer_{}'.format(n_layers - 1)\n ).minimize(\n d_cost, var_list=d_vars, global_step=self.global_step)\n\n else:\n g_train = self.g_optimizer.minimize(g_cost, var_list=g_vars)\n d_train = self.d_optimizer.minimize(d_cost, var_list=d_vars, global_step=self.global_step)\n\n # Increment image count\n n_imgs = tf.shape(x)[0]\n new_image_count = tf.add(self.total_imgs, n_imgs)\n img_step_op = tf.assign(self.total_imgs, new_image_count)\n d_train = tf.group(d_train, img_step_op)\n\n # Print variable names to before running model\n print('\\nGenerator variables for layer {} ({} x {}):'.format(n_layers - 1, dim, dim))\n print([var.name for var in g_vars])\n print('\\nDiscriminator variables for layer {} ({} x {}):'.format(n_layers - 1, dim, dim))\n print([var.name for var in d_vars])\n\n # Generate preview images\n with tf.variable_scope('image_preview'):\n n_real_imgs = min(self.batch_sizes[n_layers - 1], 4)\n fake_imgs = tensor_to_imgs(Gz)\n real_imgs = tensor_to_imgs(x[:n_real_imgs])\n\n # Upsize images to normal visibility\n if dim < 256:\n fake_imgs = resize_images(fake_imgs, (256, 256))\n real_imgs = resize_images(real_imgs, (256, 256))\n\n # Concatenate images into one large image for preview, only used if 24 preview images are requested\n if self.big_image and self.n_examples == 24:\n fake_img_list = tf.unstack(fake_imgs, num=24)\n fake_img_list = [tf.concat(fake_img_list[6 * i:6 * (i + 1)], 1) for i in range(4)]\n fake_imgs = tf.concat(fake_img_list, 0)\n fake_imgs = tf.expand_dims(fake_imgs, 0)\n\n real_img_list = tf.unstack(real_imgs, num=n_real_imgs)\n real_imgs = tf.concat(real_img_list, 1)\n real_imgs = tf.expand_dims(real_imgs, 0)\n\n # images summaries\n fake_img_sum = tf.summary.image('fake{}x{}'.format(dim, dim), fake_imgs, self.n_examples)\n real_img_sum = tf.summary.image('real{}x{}'.format(dim, dim), real_imgs, 4)\n\n return dict(\n wd=wd, gp=gp, wd_sum=wd_sum, gp_sum=gp_sum, g_train=g_train, d_train=d_train,\n fake_img_sum=fake_img_sum, real_img_sum=real_img_sum, Gz=Gz\n )\n\n\n # Get current layer, global step, alpha and total number of images used so far\n def get_global_vars(self):\n gs, layer, img_step, alpha, total_imgs = self.sess.run([\n self.global_step, self.layer, self.img_step, self.alpha, self.total_imgs\n ])\n if layer == 0: img_step -= self.n_imgs\n return gs, layer, img_step, alpha, total_imgs\n\n\n def get_layer_ops(self, layer):\n dim = 2 ** (layer + 2)\n batch_size = self.batch_sizes[layer]\n n_imgs = self.n_imgs\n if layer > 0: n_imgs *= 2\n\n layer_ops = self.networks[layer]\n g_train = layer_ops.get('g_train')\n d_train = layer_ops.get('d_train')\n get_ops = lambda *op_names: [layer_ops.get(name) for name in op_names]\n scalar_sum_ops = get_ops('wd', 'gp', 'wd_sum', 'gp_sum')\n img_sum_ops = get_ops('fake_img_sum', 'real_img_sum')\n\n return dim, batch_size, n_imgs, g_train, d_train, scalar_sum_ops, img_sum_ops\n\n\n # Main training function\n def train(self, save_interval=80000):\n\n def get_loop_progress(layer, img_step):\n percent_done = img_step / self.n_imgs\n if layer > 0: percent_done /= 2\n time = dt.datetime.now()\n return time, percent_done\n\n gs, prev_layer, img_step, alpha, total_imgs = self.get_global_vars()\n start_time, start_percent_done = get_loop_progress(prev_layer, img_step)\n dim, batch_size, n_imgs, g_train, d_train, scalar_sum_ops, img_sum_ops = self.get_layer_ops(prev_layer)\n\n save_step = (total_imgs // save_interval + 1) * save_interval\n\n while total_imgs < self.max_imgs:\n gs, layer, img_step, alpha, total_imgs = self.get_global_vars()\n\n # Get network operations and loss functions for current layer\n if layer != prev_layer:\n start_time, start_percent_done = get_loop_progress(prev_layer, img_step)\n dim, batch_size, n_imgs, g_train, d_train, scalar_sum_ops, img_sum_ops = self.get_layer_ops(layer)\n\n # Get training data and latent variables to store in feed_dict\n feed_dict = {\n self.x_placeholder: self.feed.x_batch(batch_size, dim),\n self.z_placeholder: self.feed.z_batch(batch_size)\n }\n\n # Here's where we actually train the model\n for _ in range(self.batch_repeats):\n self.sess.run(d_train, feed_dict)\n self.sess.run(g_train, feed_dict)\n\n if gs % 20 == 0:\n\n # Get loss values and summaries\n wd_value, gp_value, wd_sum_str, gp_sum_str = self.sess.run(scalar_sum_ops, feed_dict)\n\n # Print current status, loss functions, etc.\n time, percent_done = get_loop_progress(layer, img_step)\n print(\n 'dimensions: ({} x {}) ---- {}% ---- images: {}/{} ---- alpha: {} ---- global step: {}'\n '\\nWasserstein distance: {}\\ngradient penalty: {}'.format(\n dim, dim, np.round(percent_done * 100, 4), img_step, n_imgs,\n np.round(alpha, 4), gs, wd_value, gp_value\n ))\n\n # Calculate and print estimated time remaining\n delta_t = time - start_time\n time_remaining = delta_t * (1 / (percent_done - start_percent_done + 1e-8) - 1)\n print('est. time remaining on layer {}: {}\\n'.format(layer, time_remaining))\n\n # Log scalar data every 20 global steps\n self.writer.add_summary(wd_sum_str, gs)\n self.writer.add_summary(gp_sum_str, gs)\n\n # Operations to run every save interval\n if total_imgs > save_step:\n save_step += save_interval\n\n # Save the model and generate image previews\n print('\\nsaving and making images...\\n')\n self.saver.save(\n self.sess, os.path.join(self.logdir, \"model.ckpt\"),\n global_step=self.global_step\n )\n self.feed.save()\n\n img_preview_feed_dict = {\n self.x_placeholder: feed_dict[self.x_placeholder][:4],\n self.z_placeholder: self.feed.z_fixed\n }\n\n fake_img_sum_str, real_img_sum_str = self.sess.run(\n img_sum_ops, img_preview_feed_dict\n )\n self.writer.add_summary(fake_img_sum_str, gs)\n self.writer.add_summary(real_img_sum_str, gs)\n\n prev_layer = layer\n\n\n def get_cur_res(self):\n cur_layer = self.sess.run(self.layer)\n return 2 ** (2 + cur_layer)\n\n\n def generate(self, z):\n solo = z.ndim == 1\n if solo:\n z = np.expand_dims(z, 0)\n\n cur_layer = int(self.sess.run(self.layer))\n imgs = self.networks[cur_layer][9]\n imgs = self.sess.run(imgs, {self.z_placeholder: z})\n\n if solo:\n imgs = np.squeeze(imgs, 0)\n return imgs\n\n\nif __name__ == '__main__':\n # progan = ProGAN(logdir='logdir_v5', imgdir='memmaps')\n\n # progan = ProGAN(logdir='logdir_v6', imgdir='memmaps', batch_repeats=4)\n\n progan = ProGAN(logdir='logdir_v8', imgdir='memmaps', batch_repeats=4)\n # progan = ProGAN(logdir='logdir_v9', imgdir='memmaps', batch_repeats=4, batch_sizes=[128, 128, 128, 64, 32, 16, 12, 8, 4])\n\n progan.train()" }, { "path": "scripts/downloader.py", "content": "import os\nimport requests\n\nfrom selenium import webdriver\nfrom selenium.webdriver.common.by import By\nfrom selenium.webdriver.support.ui import WebDriverWait\nfrom selenium.webdriver.support import expected_conditions as EC\n\n\nsubreddit = input('Enter subreddit name: ')\nsave_dir = input('Enter name of folder to save images in: ')\n\nif not os.path.isdir(save_dir):\n os.makedirs(save_dir)\n\npages = 100\nimg_n = 0\nbrowser = webdriver.Firefox()\nbrowser.get('https://old.reddit.com/r/{}'.format(subreddit))\n\nfor i in range(pages):\n icons = WebDriverWait(browser, 300).until(\n EC.presence_of_all_elements_located(\n (By.CLASS_NAME, \"expando-button\")\n )\n )\n\n for icon in icons:\n icon.click()\n\n links = WebDriverWait(browser, 300).until(\n EC.presence_of_all_elements_located((By.CLASS_NAME, \"may-blank\"))\n )\n links = list(set([a.get_attribute('href') for a in links if a.get_attribute('href').endswith('.jpg')]))\n\n for link in links:\n image = requests.get(link)\n with open('{}/img_{}.jpg'.format(save_dir, img_n), 'wb') as f:\n f.write(image.content)\n img_n += 1\n\n if i != pages - 1:\n next_button = WebDriverWait(browser, 300).until(\n EC.presence_of_element_located((By.CLASS_NAME, \"next-button\"))\n )\n next_button.click()\n\n print('page: {}, images: {}'.format(i, len(links)))" }, { "path": "scripts/image_reshape.py", "content": "import os\nimport numpy as np\nfrom PIL import Image\n\n\ndef generate_square_crops(imgdir, savedir, crops_per_img=10, max_size=1024, filter=Image.BICUBIC):\n\n img_files = [os.path.join(imgdir, f) for f in os.listdir(imgdir)]\n savedir = os.path.join(savedir, '_temp')\n if not os.path.exists(savedir): os.makedirs(savedir)\n\n for i, f in enumerate(img_files):\n\n with Image.open(f) as img:\n width, height = img.size\n\n if width < max_size or height < max_size: continue\n\n landscape = width > height\n if landscape:\n new_height = max_size\n new_width = int(width * (max_size / height))\n offset = int(max_size * (width / height - 1) + 1)\n else:\n new_width = max_size\n new_height = int(height * (max_size / width))\n offset = int(max_size * (height / width - 1) + 1)\n\n n_crops = min(offset, crops_per_img)\n window_slide_len = offset / n_crops\n\n try:\n img = img.convert('RGB')\n img = img.resize((new_width, new_height), filter)\n\n for j in range(n_crops):\n shift = int(j * window_slide_len)\n\n if landscape: window = (shift, 0, max_size + shift, max_size)\n else: window = (0, shift, max_size, max_size + shift)\n\n cropped_img = img.crop(window)\n mirror_img = cropped_img.transpose(Image.FLIP_LEFT_RIGHT)\n\n path = os.path.join(savedir, 'img_{}_{}.jpg'.format(i, j))\n mirror_path = os.path.join(savedir, 'img_{}_{}_mirror.jpg'.format(i, j))\n cropped_img.save(path, \"JPEG\")\n mirror_img.save(mirror_path, \"JPEG\")\n\n print('Processed {}\\n'.format(f))\n\n except OSError:\n continue\n\n\ndef resize(savedir, NCHW=True, min_size=4, max_size=1024, max_mem=0.8,\n use_uint8=True, filter=Image.BICUBIC):\n\n resized_img_dir = os.path.join(savedir, '_temp')\n img_files = [os.path.join(resized_img_dir, f) for f in os.listdir(resized_img_dir)]\n np.random.shuffle(img_files)\n savedir = os.path.join(savedir, 'memmaps')\n if not os.path.exists(savedir): os.makedirs(savedir)\n\n sizes = [\n 2 ** i for i in range(\n int(np.log2(min_size)),\n int(np.log2(max_size)) + 1\n )]\n\n pixel_bytes = 3 if use_uint8 else 12\n max_bytes = max_mem * 1e9\n\n for s in sizes:\n max_imgs = int(max_bytes / (pixel_bytes * s ** 2))\n batch_shape = (max_imgs, 3, s, s) if NCHW else (max_imgs, s, s, 3)\n batch = np.zeros(batch_shape, np.uint8)\n img_count = 0\n batch_count = 0\n\n for f in img_files:\n\n with Image.open(f) as img:\n width, height = img.size\n\n if width != s and height != s:\n img = img.resize((s, s), filter)\n img = np.asarray(img, np.uint8)\n if NCHW:\n img = np.transpose(img, (2, 0, 1))\n batch[img_count] = img\n\n if img_count < max_imgs - 1:\n img_count += 1\n else:\n path = os.path.join(savedir, '{}_{}.npy'.format(s, batch_count))\n np.save(path, batch)\n print('Saved {}'.format(path))\n img_count = 0\n batch_count += 1\n\n if img_count != 0:\n path = os.path.join(savedir, '{}_{}.npy'.format(s, batch_count))\n np.save(path, batch[:img_count])\n print('Saved {}'.format(path))\n\n\nif __name__ == '__main__':\n imgdir = input('Image directory: ')\n savedir = input('Memmap directory: ')\n\n #generate_square_crops(imgdir, savedir)\n resize(savedir)" } ]