{"cells":[{"metadata":{},"cell_type":"markdown","source":"## CycleGAN for Image translation with Residual Blocks and TensorFlow 2"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"from keras.models import Model, Sequential\nfrom keras.optimizers import Adam\nfrom keras.layers import Input, Conv2D, Conv2DTranspose, Activation, LeakyReLU, Concatenate\nfrom keras.initializers import RandomNormal","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# install keras.contrib with the InstanceNormalization layer\n!pip install git+https://www.github.com/keras-team/keras-contrib.git","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# import the new instance normalization layer\nfrom keras_contrib.layers.normalization.instancenormalization import InstanceNormalization","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# create the discriminator for the GAN\ndef discriminator(input_shape):\n    # kernel initializer for our layers\n    init = RandomNormal(stddev = 0.02)\n    # image input\n    inputs = Input(shape = input_shape)\n    # first conv layer with 32 filters\n    x = Conv2D(16, 5, 2, padding='same', kernel_initializer=init, name='conv_0')(inputs)\n    x = LeakyReLU(alpha=0.2)(x)\n    # conv layers with instance normalization\n    x = Conv2D(32, 5, 2, padding='same', kernel_initializer=init, name='conv_1')(x)\n    x = InstanceNormalization(axis=-1)(x)\n    x = LeakyReLU(alpha=0.2)(x)\n    \n    x = Conv2D(64, 5, 2, padding='same', kernel_initializer=init, name='conv_2')(x)\n    x = InstanceNormalization(axis=-1)(x)\n    x = LeakyReLU(alpha=0.2)(x)\n    \n    x = Conv2D(128, 5, 2, padding='same', kernel_initializer=init, name='conv_3')(x)\n    x = InstanceNormalization(axis=-1)(x)\n    x = LeakyReLU(alpha=0.2)(x)\n    \n    x = Conv2D(256, 5, padding='same', kernel_initializer=init, name='conv_5')(x)\n    x = InstanceNormalization(axis=-1)(x)\n    x = LeakyReLU(alpha=0.2)(x)\n    \n    # output layer\n    outputs = Conv2D(1, 5, padding='same', kernel_initializer=init, name='output')(x)\n    # define the discriminator\n    model = Model(inputs, outputs)\n    # compile the model\n    model.compile(loss='mse', optimizer=Adam(lr=0.0001, beta_1=0.5))\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# define image shape constant\nIMG_SHAPE = (256, 256, 3)\n# declare the model\nmodel = discriminator(IMG_SHAPE)\n# model summary\nmodel.summary()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# create residual blocks for the generator\ndef res_block(filters, inputs):\n    # kernel weights initializer\n    init = RandomNormal(stddev=0.02)\n    x = Conv2D(filters, 3, padding='same', kernel_initializer=init)(inputs)\n    x = InstanceNormalization(axis=-1)(x)\n    x = Activation('selu')(x)\n    x = Conv2D(filters, 3, padding='same', kernel_initializer=init)(x)\n    x = InstanceNormalization(axis=-1)(x)\n    # concatenate second conv layer with the inputs\n    x = Concatenate()([x, inputs])\n    return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# generator function\ndef generator(img_shape = (256, 256, 3), n_blocks = 8):\n    # weight initialization\n    init = RandomNormal(stddev=0.02)\n    inputs = Input(shape = (256, 256, 3))\n    x = Conv2D(16, 5, padding='same', kernel_initializer=init)(inputs)\n    x = InstanceNormalization(axis=-1)(x)\n    x = Activation('selu')(x)\n    \n    x = Conv2D(32, 3, 2, padding='same', kernel_initializer=init)(x)\n    x = InstanceNormalization(axis=-1)(x)\n    x = Activation('selu')(x)\n    \n    x = Conv2D(64, 3, 2, padding='same', kernel_initializer=init)(x)\n    x = InstanceNormalization(axis=-1)(x)\n    x = Activation('selu')(x)\n    \n    # add residual blocks to our generator\n    for _ in range(n_blocks):\n        x = res_block(128, x)\n    \n    # transpose convolutions\n    x = Conv2DTranspose(32, 3, strides = 2, padding='same', kernel_initializer=init)(x)\n    x = InstanceNormalization(axis=-1)(x)\n    x = Activation('selu')(x)\n    \n    x = Conv2DTranspose(64, 3, 2, padding='same', kernel_initializer=init)(x)\n    x = InstanceNormalization(axis=-1)(x)\n    x = Activation('selu')(x)\n    \n    # output layer\n    x = Conv2D(3, 7, padding='same', kernel_initializer=init)(x)\n    x = InstanceNormalization(axis=-1)(x)\n    outputs = Activation('tanh')(x)\n    \n    # create the model\n    model = Model(inputs, outputs)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = generator()\nmodel.summary()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%cd ..\n!ls ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# define the two generators and two discriminators for CycleGAN\ngenerator_BtoA = generator(IMG_SHAPE)\ngenerator_AtoB = generator(IMG_SHAPE)\ndiscriminator_A = discriminator(IMG_SHAPE)\ndiscriminator_B = discriminator(IMG_SHAPE)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# define composite model\ndef composite_model(g1, d1, g2, image_shape):\n    # mark the first generator as trainable\n\tg1.trainable = True\n\t# freeze discriminator\n\td1.trainable = False\n\t# freeze second generator\n\tg2.trainable = False\n\tinput_gen = Input(shape=image_shape)\n\tgen1_out = g1(input_gen)\n\toutput_d = d1(gen1_out)\n\tinput_id = Input(shape=image_shape)\n\toutput_id = g1(input_id)\n\toutput_f = g2(gen1_out)\n\tgen2_out = g2(input_id)\n\toutput_b = g1(gen2_out)\n\tmodel = Model([input_gen, input_id], [output_d, output_id, output_f, output_b])\n\toptimizer = Adam(lr=0.0001, beta_1=0.5)\n\t# compile model with L1 loss\n\tmodel.compile(loss=['mse', 'mae', 'mae', 'mae'], loss_weights=[1, 5, 10, 10], optimizer=optimizer)\n\treturn model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from PIL import Image\nimport os\nimport numpy as np\n\nmonet_path = 'input/gan-getting-started/monet_jpg/'\nmonet_array = []\n\n# monet images to npz\nfor _, file in enumerate(os.listdir(monet_path)): \n    image = Image.open(monet_path + file)\n    single_array = np.array(image)\n    monet_array.append(single_array)\n# save the images\nnp.savez('monet_compressed.npz', monet_array)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nx1 = np.load('monet_compressed.npz')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# photos to npz\nphoto_path = 'input/gan-getting-started/photo_jpg/'\nphoto_array_1 = []\nc = 0\n# photo images to npz (limit the length to prevent memory error)\nfor _, file in enumerate(os.listdir(photo_path)): \n    c+=1\n    image = Image.open(photo_path + file)\n    single_array_1 = np.array(image)\n    photo_array_1.append(single_array_1)\n    if c == 3000:\n        break\n# save the images\nnp.savez('photos1_compressed.npz', photo_array_1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# load the photos\nx2 = np.load('photos1_compressed.npz')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# scale the data to -1 and 1\n\nX1 = (x1['arr_0'] - 127.5) / 127.5","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"X2 = (x2['arr_0'] - 127.5) / 127.5","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from matplotlib import pyplot as plt","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# translate an image without training\nplt.axis(False)\ntranslation = generator_BtoA(np.reshape(X2[1], (1, 256, 256, 3)))\nprint('TRANSLATED IMAGE WITHOUT TRAINING')\nplt.imshow(np.reshape(translation, (256, 256, 3)))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print('ORIGINAL IMAGE')\nplt.axis(False)\nplt.imshow(X2[1])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from numpy.random import randint\ndef generate_real_samples(dataset, n_samples, patch_shape):\n\t# choose random instances\n\tix = randint(0, dataset.shape[0], n_samples)\n\t# retrieve selected images\n\tX = dataset[ix]\n\t# generate 'real' class labels (1)\n\ty = np.ones((n_samples, patch_shape, patch_shape, 1))\n\treturn X, y","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# generate fake images (batch size)\ndef generate_fake_samples(g_model, dataset, patch_shape):\n\t# generate the image\n\tX = g_model(dataset)\n\t# create 'fake' class labels (0)\n\ty = np.zeros((len(X), patch_shape, patch_shape, 1))\n\treturn X, y","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def summarize_performance(step, g_model, trainX, name, n_samples=5):\n\tX_in, _ = generate_real_samples(trainX, n_samples, 0)\n\t# generate target images\n\tX_out, _ = generate_fake_samples(g_model, X_in, 0)\n\t# rescale to 0 - 1\n\tX_in = (X_in + 1) / 2.0\n\tX_out = (X_out + 1) / 2.0\n\tfor i in range(n_samples):\n\t\tplt.subplot(2, n_samples, 1 + i)\n\t\tplt.axis('off')\n\t\tplt.imshow(X_in[i])\n\t# plot target image\n\tfor i in range(n_samples):\n\t\tplt.subplot(2, n_samples, 1 + n_samples + i)\n\t\tplt.axis('off')\n\t\tplt.imshow(X_out[i])\n\t# save plot to file\n\tfilename1 = '%s_generated_plot_%06d.png' % (name, (step+1))\n\tplt.savefig(filename1)\n\tplt.close()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# create a pool of fake images to trace the training\nfrom random import random\ndef update_image_pool(pool, images, max_size=50):\n\tselected = list()\n\tfor image in images:\n\t\tif len(pool) < max_size:\n\t\t\t# stock the pool\n\t\t\tpool.append(image)\n\t\t\tselected.append(image)\n\t\telif random() < 0.5:\n\t\t\t# use image, but don't add it to the pool\n\t\t\tselected.append(image)\n\t\telse:\n\t\t\t# replace an existing image and use replaced image\n\t\t\tix = randint(0, len(pool))\n\t\t\tselected.append(pool[ix])\n\t\t\tpool[ix] = image\n\treturn np.asarray(selected)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# train the cycleGAN models for 50 epochs with a batch size of 1\ndef train(d_model_A, d_model_B, g_model_AtoB, g_model_BtoA, c_model_AtoB, c_model_BtoA, dataset):\n    # train for 10 epochs\n\tn_epochs, n_batch, = 10, 1\n\t# determine the output square shape of the discriminator\n\tn_patch = d_model_A.output_shape[1]\n\t# unpack dataset\n\ttrainA, trainB = dataset\n\t# prepare image pool for fakes\n\tpoolA, poolB = list(), list()\n\t# calculate the number of batches per training epoch\n\tbat_per_epo = int(len(trainA) / n_batch)\n\t# calculate the number of training iterations\n\tn_steps = bat_per_epo * n_epochs\n\t# manually enumerate epochs\n\tfor i in range(n_steps):\n\t\tX_realA, y_realA = generate_real_samples(trainA, n_batch, n_patch)\n\t\tX_realB, y_realB = generate_real_samples(trainB, n_batch, n_patch)\n\t\t# generate fake images\n\t\tX_fakeA, y_fakeA = generate_fake_samples(g_model_BtoA, X_realB, n_patch)\n\t\tX_fakeB, y_fakeB = generate_fake_samples(g_model_AtoB, X_realA, n_patch)\n\t\t# update images\n\t\tX_fakeA = update_image_pool(poolA, X_fakeA)\n\t\tX_fakeB = update_image_pool(poolB, X_fakeB)\n\t\tg_loss2, _, _, _, _  = c_model_BtoA.train_on_batch([X_realB, X_realA], [y_realA, X_realA, X_realB, X_realA])\n\t\tdA_loss1 = d_model_A.train_on_batch(X_realA, y_realA)\n\t\tdA_loss2 = d_model_A.train_on_batch(X_fakeA, y_fakeA)\n\t\tg_loss1, _, _, _, _ = c_model_AtoB.train_on_batch([X_realA, X_realB], [y_realB, X_realB, X_realA, X_realB])\n\t\tdB_loss1 = d_model_B.train_on_batch(X_realB, y_realB)\n\t\tdB_loss2 = d_model_B.train_on_batch(X_fakeB, y_fakeB)\n\t\t# summarize performance\n\t\tprint('>%d, dA[%.3f,%.3f] dB[%.3f,%.3f] g[%.3f,%.3f]' % (i+1, dA_loss1,dA_loss2, dB_loss1,dB_loss2, g_loss1,g_loss2))\n\t\t# evaluate the model performance every so often\n\t\tif (i+1) % (bat_per_epo * 1) == 0:\n\t\t\t# plot A->B translation\n\t\t\tsummarize_performance(i, g_model_AtoB, trainA, 'AtoB')\n\t\t\t# plot B->A translation\n\t\t\tsummarize_performance(i, g_model_BtoA, trainB, 'BtoA')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# create the composite model for A to B\ncomposite_AtoB = composite_model(generator_AtoB, discriminator_B,generator_BtoA, IMG_SHAPE)\ncomposite_BtoA = composite_model(generator_BtoA, discriminator_A, generator_AtoB, IMG_SHAPE)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# finally, train the model\ntrain(discriminator_A, discriminator_B, generator_AtoB,generator_BtoA,  composite_AtoB,\n     composite_BtoA, [X1, X2])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# save the model\n%cd working\ngenerator_BtoA.save('monet_translator.h5')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"translation = generator_BtoA(np.reshape(X2[600], (1, 256, 256, 3)))\nprint('TRANSLATED IMAGE AFTER TRAINING')\nplt.imshow(np.reshape(translation, (256, 256, 3)))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print('ORIGINAL BASE IMAGE')\nplt.imshow(X2[600])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}