def create_gif(anim_file,path,class_label):
with imageio.get_writer(os.path.join(path,anim_file), mode='I') as writer:
filenames = glob.glob(path+'/image_class_'+str(class_label)+'*.png')
filenames = sorted(filenames)
for filename in filenames:
image = imageio.imread(filename)
writer.append_data(image)
image = imageio.imread(filename)
writer.append_data(image)
embed.embed_file(os.path.join(path,anim_file))
checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
train(train_dataset, 1)
generator.save_model("digit_generator")
discriminator.save_model("digit_discriminator")
# Display a single image using the epoch number
def display_image(epoch_no):
return PIL.Image.open('image_at_epoch_{:04d}.png'.format(epoch_no))
display_image(EPOCHS)
anim_file = 'dcgan.gif'
with imageio.get_writer(anim_file, mode='I') as writer:
filenames = glob.glob('image*.png')
filenames = sorted(filenames)
for filename in filenames:
image = imageio.imread(filename)
writer.append_data(image)
image = imageio.imread(filename)
writer.append_data(image)
import tensorflow_docs.vis.embed as embed
embed.embed_file(anim_file)
def main():
# Constants and hyperparameters.
batch_size = 64
num_channels = 1
num_classes = 10
image_size = 28
latent_dim = 128
# Loading the MNIST dataset and preprocessing it.
# Use all the available examples from both the training and test
# sets.
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
all_digits = np.concatenate([x_train, x_test])
all_labels = np.concatenate([y_train, y_test])
# Scale the pixel values to [0, 1] range, add a channel dimension
# to the images, and one-hot encode the labels.
all_digits = all_digits.astype("float32") / 255.0
all_digits = np.reshape(all_digits, (-1, 28, 28, 1))
all_labels = keras.utils.to_categorical(all_labels, 10)
# Create tf.data.Dataset.
dataset = tf.data.Dataset.from_tensor_slices((all_digits, all_labels))
dataset = dataset.shuffle(buffer_size=1024).batch(batch_size)
print(f"Shape of training images: {all_digits.shape}")
print(f"Shape of training labels: {all_labels.shape}")
# Calculating the number of input channels for the generator and
# discriminator.
# In a regular (unconditional) GAN, start by sampling noise (of
# some fixed dimension) from a normal distribution. In this case,
# there also needs to account for the class labels. Add the number
# of classes to the input channels of the generator (noise input)
# as well as the discriminator (generated image input).
generator_in_channels = latent_dim + num_classes
discriminator_in_channels = num_channels + num_classes
print(generator_in_channels, discriminator_in_channels)
# Creating the discriminator and generator.
# The model definitions (discriminator, generator, and
# ConditionalGAN) have been adapted from this example
# (https://keras.io/guides/customizing_what_happens_in_fit/).
# Create the discriminator.
discriminator = keras.Sequential(
[
keras.layers.InputLayer((28, 28, discriminator_in_channels)),
layers.Conv2D(64, (3, 3), strides=(2, 2), padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.Conv2D(128, (3, 3), strides=(2, 2), padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.GlobalMaxPooling2D(),
layers.Dense(1),
],
name="discriminator",
)
# Create the generator.
generator = keras.Sequential(
[
keras.layers.InputLayer((generator_in_channels, )),
# Want to generate 128 + num_classes coefficients to
# reshape into a 7 x 7 x (128 + num_classes) map.
layers.Dense(7 * 7 * generator_in_channels),
layers.LeakyReLU(alpha=0.2),
layers.Reshape((7, 7, generator_in_channels)),
layers.Conv2DTranspose(128,
(4, 4), strides=(2, 2), padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.Conv2DTranspose(128,
(4, 4), strides=(2, 2), padding="same"),
layers.LeakyReLU(alpha=0.2),
layers.Conv2D(1, (7, 7), padding="same", activation="sigmoid"),
],
name="generator",
)
# Creating a ConditionalGAN model.
class ConditionalGAN(keras.Model):
def __init__(self, discriminator, generator, latent_dim):
super(ConditionalGAN, self).__init__()
self.discriminator = discriminator
self.generator = generator
self.latent_dim = latent_dim
self.gen_loss_tracker = keras.metrics.Mean(name="generator_loss")
self.disc_loss_tracker = keras.metrics.Mean(
name="discriminator_loss")
@property
def metrics(self):
return [self.gen_loss_tracker, self.disc_loss_tracker]
def compile(self, d_optimizer, g_optimizer, loss_fn):
super(ConditionalGAN, self).compile()
self.d_optimizer = d_optimizer
self.g_optimizer = g_optimizer
self.loss_fn = loss_fn
def train_step(self, data):
# Unpack the data.
real_images, one_hot_labels = data
# Add dummy dimensions to the labels so that they can be
# concatenated with the images. This is for the
# discriminator.
image_one_hot_labels = one_hot_labels[:, :, None, None]
image_one_hot_labels = tf.repeat(image_one_hot_labels,
repeats=[image_size * image_size])
image_one_hot_labels = tf.reshape(
image_one_hot_labels,
(-1, image_size, image_size, num_classes))
# Sample random points in the latent space and concatenate
# the labels. This is for the generator.
batch_size = tf.shape(real_images)[0]
random_latent_vectors = tf.random.normal(shape=(batch_size,
self.latent_dim))
random_vector_labels = tf.concat(
[random_latent_vectors, one_hot_labels], axis=1)
# Decode the noise (guided by labels) to fake images.
generated_images = self.generator(random_vector_labels)
# Combine them with real images. Note that the labels are
# being concatenated with the images here.
fake_image_and_labels = tf.concat(
[generated_images, image_one_hot_labels], -1)
real_image_and_labels = tf.concat(
[real_images, image_one_hot_labels], -1)
combined_images = tf.concat(
[fake_image_and_labels, real_image_and_labels], axis=0)
# Assemble labels discriminating real from fake images.
labels = tf.concat(
[tf.ones((batch_size, 1)),
tf.zeros((batch_size, 1))], axis=0)
# Train the discriminator.
with tf.GradientTape() as tape:
predictions = self.discriminator(combined_images)
d_loss = self.loss_fn(labels, predictions)
grads = tape.gradient(d_loss, self.discriminator.trainable_weights)
self.d_optimizer.apply_gradients(
zip(grads, self.discriminator.trainable_weights))
# Sample random points in the latent space.
random_latent_vectors = tf.random.normal(shape=(batch_size,
self.latent_dim))
random_vector_labels = tf.concat(
[random_latent_vectors, one_hot_labels], axis=1)
# Assemble labels that say "all real images".
misleading_labels = tf.zeros((batch_size, 1))
# Train the generator (not that we should *not* update the
# weights of the discriminator)!
with tf.GradientTape() as tape:
fake_images = self.generator(random_vector_labels)
fake_image_and_labels = tf.concat(
[fake_images, image_one_hot_labels], -1)
predictions = self.discriminator(fake_image_and_labels)
g_loss = self.loss_fn(misleading_labels, predictions)
grads = tape.gradient(g_loss, self.generator.trainable_weights)
self.g_optimizer.apply_gradients(
zip(grads, self.generator.trainable_weights))
# Monitor loss.
self.gen_loss_tracker.update_state(g_loss)
self.disc_loss_tracker.update_state(d_loss)
return {
"g_loss": self.gen_loss_tracker.result(),
"d_loss": self.disc_loss_tracker.result(),
}
# Training the ConditionalGAN.
cond_gan = ConditionalGAN(discriminator=discriminator,
generator=generator,
latent_dim=latent_dim)
cond_gan.compile(
d_optimizer=keras.optimizers.Adam(learning_rate=0.0003),
g_optimizer=keras.optimizers.Adam(learning_rate=0.0003),
loss_fn=keras.losses.BinaryCrossentropy(from_logits=True),
)
cond_gan.fit(dataset, epochs=20)
# Interpolating between classes with the trained generator.
# First extract the trained generator from the Conditional GAN.
trained_gen = cond_gan.generator
# Choose the number of intermediate images that would be generated
# in between the interpolation + 2 (start and last images).
num_interpolations = 9 # @param {type: "integer"}
# Sample noise for the interpolation.
interpolation_noise = tf.random.normal(shape=(1, latent_dim))
interpolation_noise = tf.repeat(interpolation_noise,
repeats=num_interpolations)
interpolation_noise = tf.reshape(interpolation_noise,
(num_interpolations, latent_dim))
def interpolate_class(first_number, second_number):
# Convert the start and end labels to one-hot encoded vectors.
first_label = keras.utils.to_categorical([first_number], num_classes)
second_label = keras.utils.to_categorical([second_number], num_classes)
first_label = tf.cast(first_label, tf.float32)
second_label = tf.cast(second_label, tf.float32)
# Calculate the interpolation vector between the two labels.
percent_second_label = tf.linspace(0, 1, num_interpolations)[:, None]
percent_second_label = tf.cast(percent_second_label, tf.float32)
interpolation_labels = (first_label * (1 - percent_second_label) +
second_label * percent_second_label)
# Combine the noise and the labels and run inference with the
# generator.
noise_and_labels = tf.concat(
[interpolation_noise, interpolation_labels], 1)
fake = trained_gen.predict(noise_and_labels)
return fake
start_class = 1 # @param {type: "slider", min:0, max:9, step:1}
end_class = 5 # @param {type: "slider", min:0, max:9, step:1}
fake_images = interpolate_class(start_class, end_class)
# Here, first sample noise from a normal distribution and then
# repeat that for num_interpolation times and reshape the result
# accordingly. Then distribute it uniformly for num_interpolation
# with the label identities being present in some proportion.
fake_images *= 255
converted_images = fake_images.astype(np.uint8)
converted_images = tf.image.resize(converted_images,
(96, 96)).numpy().astype(np.uint8)
imageio.mimsave("animation.gif", converted_images, fps=1)
embed.embed_file("animation.gif")
# Can further improve the performance of this model with recipes
# like WGAN-GP. Conditional generation is also widely used in many
# modern image generation architectures like VQ-GANs, DALL-E, etc.
# Exit the program.
exit(0)
def to_gif(images):
converted_images = images.astype(np.uint8)
imageio.mimsave("animation.gif", converted_images, fps=10)
return embed.embed_file("animation.gif")
state = tf.expand_dims(state, 0)
action_probs, _ = model(state)
action = np.argmax(np.squeeze(action_probs))
state, _, done, _ = env.step(action)
state = tf.constant(state, dtype=tf.float32)
# Render screen every 10 steps
if i % 10 == 0:
screen = env.render(mode='rgb_array')
images.append(Image.fromarray(screen))
if done:
break
return images
# Save GIF image
images = render_episode(env, model, max_steps_per_episode)
image_file = 'cartpole-v0.gif'
# loop=0: loop forever, duration=1: play each frame for 1ms
images[0].save(image_file,
save_all=True,
append_images=images[1:],
loop=0,
duration=1)
import tensorflow_docs.vis.embed as embed
embed.embed_file(image_file)
noise_and_labels = tf.concat([interpolation_noise, interpolation_labels],
1)
fake = trained_gen.predict(noise_and_labels)
return fake
start_class = 1 # @param {type:"slider", min:0, max:9, step:1}
end_class = 5 # @param {type:"slider", min:0, max:9, step:1}
fake_images = interpolate_class(start_class, end_class)
"""
Here, we first sample noise from a normal distribution and then we repeat that for
`num_interpolation` times and reshape the result accordingly.
We then distribute it uniformly for `num_interpolation`
with the label indentities being present in some proportion.
"""
fake_images *= 255.0
converted_images = fake_images.astype(np.uint8)
converted_images = tf.image.resize(converted_images,
(96, 96)).numpy().astype(np.uint8)
imageio.mimsave("animation.gif", converted_images, fps=1)
embed.embed_file("animation.gif")
"""
We can further improve the performance of this model with recipes like
[WGAN-GP](https://keras.io/examples/generative/wgan_gp).
Conditional generation is also widely used in many modern image generation architectures like
[VQ-GANs](https://arxiv.org/abs/2012.09841), [DALL-E](https://openai.com/blog/dall-e/),
etc.
"""
def to_gif(images):
converted_images = np.clip(images * 255, 0, 255).astype(np.uint8)
imageio.mimsave('./animation.gif', converted_images, fps=25)
return embed.embed_file('./animation.gif')
def animate(images):
converted_images=np.clip(images*255,0,255).astype(np.uint8)
imageio.mimsave('animation.gif',converted_images)
return embed.embed_file('animation.gif')
def animate(images):
images = np.array(images)
print(np.amax(images[0]))
imageio.mimsave(path2, images)
return embed.embed_file(path2)
