import numpy as np
from matplotlib import pyplot as plt
from keras.datasets import mnist
from gan import GAN
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train.reshape([-1, 28, 28, 1]) / 255
x_test = x_test.reshape([-1, 28, 28, 1]) / 255
x = np.concatenate([x_train, x_test])
gan = GAN()
gan.gan.summary()
gan.train(x, 100, batch_size=100)
gan.save('number_gan/gan.h5')
noise = np.random.normal(size=[10, 8 * 8])
images = gan.generator.predict(noise)
images = images.reshape([-1, 28, 28])
plt.gray()
for i in range(10):
plt.subplot(2, 5, i + 1)
plt.imshow(images[i])
plt.show()
opts.generator_frequency = 1
opts.generator_dropout = 0.3
opts.label_softness = 0.2
opts.batch_size = 128
opts.epoch_number = 300
#Note: if you are using jupyter notebook you might need to disable workers:
opts.workers_nbr = 0
#List of tranformations applied to each input image
opts.transforms = [
transforms.Resize(int(opts.image_size), Image.BICUBIC),
transforms.ToTensor(), #do not forget to transform image into a tensor
transforms.Normalize((.5, .5, .5), (.5, .5, .5))
#RandomNoise(0.01) #adding noise might prevent the discriminator from over-fitting
]
#Build GAN
model = GAN(opts)
#Display GAN architecture (note: only work if cuda is enabled)
summary(model.generator.cuda(), input_size=(opts.latent_dim, 1, 1))
summary(model.discriminator.cuda(),
input_size=(opts.channels_nbr, opts.image_size, opts.image_size))
#Start training
model.train()
#Save model
model.save(opts.output_path)
model = GAN(gex_size=train.shape[1], num_cells_generate=test.shape[0])
model.compile()
model.build(input_shape=(model.hyperparams.batch_size, model.hyperparams.latent_dim)) # req. for subclassed models
# process data for training
train_tf = tf.data.Dataset.from_tensor_slices(train.X). \
cache(). \
shuffle(buffer_size=train.shape[0], seed=utils.RANDOM). \
batch(batch_size=model.hyperparams.batch_size * strategy.num_replicas_in_sync, num_parallel_calls=tf.data.AUTOTUNE). \
prefetch(buffer_size=tf.data.AUTOTUNE)
train_tf_distributed = strategy.experimental_distribute_dataset(train_tf)
test_tf = tf.data.Dataset.from_tensor_slices(test.X). \
cache(). \
shuffle(buffer_size=test.shape[0], seed=utils.RANDOM). \
prefetch(buffer_size=tf.data.AUTOTUNE)
tb_callback = tf.keras.callbacks.TensorBoard(log_dir=utils.LOG_DIR + datetime.datetime.now().strftime("%Y%m%d-%H%M%S"),
update_freq='epoch',
write_graph=False,
profile_batch=0)
model.fit(x=train_tf_distributed,
epochs=model.hyperparams.num_epochs,
steps_per_epoch=int (train.shape[0] / model.hyperparams.batch_size), # steps = # batches per epoch
callbacks=[tb_callback,
TrainingMetrics(tb_callback),
tf.keras.callbacks.ModelCheckpoint(period=100, filepath='training/model_{epoch}')])
# Save model
model.save("training/model_final")
def main(num_faces, num_epochs, cuda, verbose=False):
# set computation device (None/CPU if in development mode, CUDA otherwise)
device = torch.device("cuda:0") if cuda else None
# load faces
masked_dir = "../data/masked"
masked_suffix = "_Mask.jpg"
unmasked_dir = "../data/unmasked"
unmasked_suffix = ".png"
masked_faces, unmasked_faces, idx_to_face_id = utils.load_faces(
num_faces, masked_dir, masked_suffix, unmasked_dir, unmasked_suffix)
if verbose:
print("loaded {} faces...".format(num_faces))
# split data into training and testing sets
split = int(0.8 * num_faces)
train_input, train_output = (
masked_faces[:split],
torch.Tensor((range(0, split))).long(),
)
test_input, test_output = (
masked_faces[split:],
torch.Tensor(range(split, num_faces)).long(),
)
static_faces = unmasked_faces[:split]
# instantiate GAN
generator = Generator(learning_rate=2e-3)
projector = Projector(load_path="../models/projector.pt")
discriminator = Discriminator()
gan = GAN(generator, projector, discriminator, device=device)
if verbose:
print("instantiated GAN...")
# compute and store unmasked discriminator embeddings
gan.compute_unmasked_embeddings(unmasked_faces)
# train
if verbose:
print("training initiated...")
gan.fit(train_input,
static_faces,
train_output,
num_epochs=num_epochs,
verbose=verbose)
if verbose:
print("\ntraining complete...")
# save models
save_dir = "../models"
suffix = time.strftime("%Y%m%d_%H%M%S")
gan.save(save_dir, suffix)
if verbose:
print("models saved under '{}/<model>_{}'...".format(save_dir, suffix))
# display sample masks and faces
plt.figure()
fig, axes = plt.subplots(2, 5)
fig.set_figwidth(20)
fig.set_figheight(7)
for idx in range(5):
# original image
face_id = idx_to_face_id[idx]
original_img = Image.open("../data/masked/{}_Mask.jpg".format(face_id))
axes[0, idx].imshow(original_img)
axes[0, idx].get_xaxis().set_ticks([])
axes[0, idx].get_yaxis().set_ticks([])
# generated mask image
mask = (gan.generator(torch.rand(1, 100).to(device))
if device else gan.generator())
masked_tensor = masked_faces[idx].unsqueeze(0)
if device:
masked_tensor = masked_tensor.to(device)
masked_image = gan.project_mask(mask, masked_tensor, process=True)[0]
masked_image = torch.transpose(masked_image, 0, 1)
masked_image = torch.transpose(masked_image, 1, 2)
masked_image = masked_image.cpu().detach().numpy()
axes[1, idx].imshow(masked_image)
axes[1, idx].get_xaxis().set_ticks([])
axes[1, idx].get_yaxis().set_ticks([])
plt.savefig("../figures/sample_masks.png")
# evaluate accuracy
train_accuracy = gan.evaluate(train_input, train_output)
test_accuracy = gan.evaluate(test_input, test_output)
masked_accuracy = gan.discriminator_evaluate(masked_faces, unmasked_faces)
unmasked_accuracy = gan.discriminator_evaluate(unmasked_faces,
unmasked_faces)
print("\nfacial recognition accuracy for...")
print(" random choice:\t\t{:.1f}%".format(100 / num_faces))
print(" training images:\t\t{:.1f}%".format(100 * train_accuracy))
print(" testing images:\t\t{:.1f}%".format(100 * test_accuracy))
print(" original masked images:\t{:.1f}%".format(100 * masked_accuracy))
print(" original unmasked images:\t{:.1f}%".format(100 *
unmasked_accuracy))
# write results to file
file_path = "../data/accuracy.txt"
with open(file_path, "w") as file:
file.write("facial recognition accuracy for...")
file.write("\n random choice:\t\t{:.1f}%".format(100 / num_faces))
file.write("\n training images:\t\t{:.1f}%".format(100 *
train_accuracy))
file.write("\n testing images:\t\t{:.1f}%".format(100 *
test_accuracy))
file.write("\n original masked images:\t{:.1f}%".format(
100 * masked_accuracy))
file.write("\n original unmasked images:\t{:.1f}%".format(
100 * unmasked_accuracy))
if verbose:
print("\nsaved results...")
print("done:)")
