GAN Image Generator

""" This module contains the implementation of a Generative Adversarial Network (GAN) using TensorFlow and Keras. The GAN is composed of a generator and a discriminator, both implemented as Keras models. The generator takes a random noise vector as input and produces an image, while the discriminator takes an image as input and classifies it as real or fake. The module provides functions to build the generator, the discriminator, and the complete GAN, as well as to train the GAN on the MNIST dataset and generate and save images. Functions: build_generator(latent_dim: int) -> models.Model build_discriminator(img_shape: tuple) -> models.Model build_gan(generator: models.Model, discriminator: models.Model) -> models. Model train_gan(generator: models.Model, discriminator: models.Model, gan: models.Model, mnist_data: np.array, latent_dim: int, epochs: int, batch_size: int, sample_interval: int) generate_and_save_images(generator: models.Model, epoch: int, test_input: np.array) """ import matplotlib.pyplot as plt import numpy as np import tensorflow as tf from tensorflow.keras import layers, models def build_generator(latent_dim): """ Build the generator part of the GAN. Args: latent_dim (int): The size of the random noise vector used as input for the generator. Returns: model: A Keras model representing the generator. """ model = models.Sequential() model.add(layers.Dense(7 * 7 * 128, input_dim=latent_dim)) model.add(layers.LeakyReLU(alpha=0.2)) model.add(layers.Reshape((7, 7, 128))) model.add(layers.Conv2DTranspose( 128, (4, 4), strides=(2, 2), padding='same')) model.add(layers.LeakyReLU(alpha=0.2)) model.add(layers.Conv2DTranspose( 128, (4, 4), strides=(2, 2), padding='same')) model.add(layers.LeakyReLU(alpha=0.2)) model.add(layers.Conv2D(1, (7, 7), activation='tanh', padding='same')) return model def build_discriminator(img_shape): """ Build the discriminator part of the GAN. Args: img_shape (tuple): The shape of the input images. Returns: model: A Keras model representing the discriminator. """ model = models.Sequential() model.add(layers.Conv2D(64, (3, 3), strides=(2, 2), padding='same', input_shape=img_shape)) model.add(layers.LeakyReLU(alpha=0.2)) model.add(layers.Conv2D(128, (3, 3), strides=(2, 2), padding='same')) model.add(layers.LeakyReLU(alpha=0.2)) model.add(layers.Flatten()) model.add(layers.Dense(1, activation='sigmoid')) return model def build_gan(generator, discriminator): """ Build the complete GAN by stacking the generator and discriminator. Args: generator (Keras model): The generator model. discriminator (Keras model): The discriminator model. Returns: model: A Keras model representing the complete GAN. """ discriminator.trainable = False model = models.Sequential() model.add(generator) model.add(discriminator) return model def load_mnist(): """ Load and preprocess the MNIST dataset. Returns: X_train (numpy array): The preprocessed training images. """ (X_train, _), (_, _) = tf.keras.datasets.mnist.load_data() X_train = (X_train.astype(np.float32) - 127.5) / 127.5 X_train = np.expand_dims(X_train, axis=-1) return X_train def train_gan(generator, discriminator, gan, mnist_data, latent_dim, epochs=10000, batch_size=128, sample_interval=1000): """ Train the GAN on the MNIST data. Args: generator (Keras model): The generator model. discriminator (Keras model): The discriminator model. gan (Keras model): The complete GAN model. mnist_data (numpy array): The preprocessed training images. latent_dim (int): The size of the random noise vector used as input for the generator. epochs (int, optional): The number of epochs to train for. Defaults to 10000. batch_size (int, optional): The batch size. Defaults to 128. sample_interval (int, optional): The interval at which to save generated images. Defaults to 1000. """ half_batch = batch_size // 2 for epoch in range(epochs): # Train Discriminator  idx = np.random.randint(0, mnist_data.shape[0], half_batch) imgs = mnist_data[idx] noise = np.random.normal(0, 1, (half_batch, latent_dim)) gen_imgs = generator.predict(noise) d_loss_real = discriminator.train_on_batch( imgs, np.ones((half_batch, 1))) d_loss_fake = discriminator.train_on_batch( gen_imgs, np.zeros((half_batch, 1))) d_loss = 0.5 * np.add(d_loss_real, d_loss_fake) # Train Generator  noise = np.random.normal(0, 1, (batch_size, latent_dim)) valid_labels = np.ones((batch_size, 1)) g_loss = gan.train_on_batch(noise, valid_labels) # Print progress and save generated images at specified intervals  if epoch % sample_interval == 0: print(f"Epoch {epoch}, D Loss: {d_loss[0]}, G Loss: {g_loss}") save_generated_images(epoch, generator) def generate_and_save_images(generator, epoch, test_input): """ Generate and save images using the generator and display them using matplotlib. Args: generator (Keras model): The generator model. epoch (int): The current epoch. test_input (numpy array): A random noise vector used as input for the generator. """ predictions = generator(test_input, training=False) fig = plt.figure(figsize=(4, 4)) for i in range(predictions.shape[0]): plt.subplot(4, 4, i + 1) plt.imshow(predictions[i, :, :, 0] * 0.5 + 0.5, cmap='gray') plt.axis('off') plt.savefig(f"gan_generated_image_epoch_{epoch}.png") plt.show() def save_generated_images(epoch, generator, latent_dim=100, examples=100, dim=(10, 10), figsize=(10, 10)): """ Save the images generated by the generator. Args: epoch (int): The current epoch. generator (Keras model): The generator model. latent_dim (int, optional): The size of the random noise vector used as input for the generator. Defaults to 100. examples (int, optional): The number of examples to generate. Defaults to 100. dim (tuple, optional): The dimensions of the grid of images. Defaults to (10, 10). figsize (tuple, optional): The size of the figure. Defaults to (10, 10). """ noise = np.random.normal(0, 1, size=[examples, latent_dim]) generated_images = generator.predict(noise) generated_images = generated_images.reshape(examples, 28, 28) plt.figure(figsize=figsize) for i in range(generated_images.shape[0]): plt.subplot(dim[0], dim[1], i+1) plt.imshow(generated_images[i], interpolation='nearest', cmap='gray_r') plt.axis('off') plt.tight_layout() plt.savefig(f'gan_generated_image_epoch_{epoch}.png') if __name__ == "__main__": # Define the size of the random vector used as input for the generator  latent_dim = 100 # Build the generator  generator = build_generator(latent_dim) # Build the discriminator  img_shape = (28, 28, 1)