dcgan_fix_512_512.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import time\n",
    "import pathlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import PIL\n",
    "import cv2\n",
    "import imageio\n",
    "from sklearn.utils import shuffle\n",
    "from glob import glob\n",
    "from ast import literal_eval\n",
    "\n",
    "import tensorflow as tf\n",
    "import tensorflow.keras as keras\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.preprocessing.image import ImageDataGenerator\n",
    "from tensorflow.keras.models import model_from_json\n",
    "from keras.preprocessing import image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Only use 1 GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# train label for load/save path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "label = 'castle'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# path to read file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FILE_PATH = './resize_512_512/{}/*.jpg'.format(label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_imgs(file_path,counts):\n",
    "    imgs_list = glob(file_path)[:counts]\n",
    "    imgs = []\n",
    "    for i in tqdm_notebook(imgs_list):\n",
    "        img = cv2.imread(i)[:,:,::-1].astype(np.float32)/255.\n",
    "        imgs.append(img)\n",
    "    imgs = np.array(imgs)\n",
    "    \n",
    "    return imgs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_generator(latent_dim, output_size):\n",
    "    filter_num = [256, 128, 64 , 32]\n",
    "    generator_input = keras.Input(shape=(latent_dim,))\n",
    "    height, width = output_size\n",
    "    \n",
    "    x = layers.Dense(filter_num[0] * int(height//16) * int(width//16))(generator_input)\n",
    "    x = layers.LeakyReLU(0.2)(x)\n",
    "    x = layers.Reshape((int(height//16), int(width//16), filter_num[0]))(x) \n",
    "    \n",
    "    #### 32*32*256\n",
    "    x = layers.UpSampling2D(size=(2, 2))(x)\n",
    "    x = layers.Conv2DTranspose(filter_num[0],(3,3),strides=(1,1),padding='same', kernel_initializer = 'he_normal')(x)  \n",
    "    x = layers.LeakyReLU(0.2)(x)\n",
    "    x = layers.Conv2DTranspose(filter_num[0],(3,3),strides=(1,1),padding='same', kernel_initializer = 'he_normal')(x)\n",
    "    x = layers.LeakyReLU(0.2)(x)\n",
    "\n",
    "    #### 64*64*128\n",
    "           \n",
    "    x = layers.UpSampling2D(size=(2, 2))(x)\n",
    "    x = layers.Conv2DTranspose(filter_num[1],(3,3),strides=(1,1),padding='same', kernel_initializer = 'he_normal')(x)\n",
    "    x = layers.LeakyReLU(0.2)(x)   \n",
    "    x = layers.Conv2DTranspose(filter_num[1],(3,3),strides=(1,1),padding='same', kernel_initializer = 'he_normal')(x)\n",
    "    x = layers.LeakyReLU(0.2)(x) \n",
    "    \n",
    "    #### 128*128*64\n",
    "    \n",
    "    x = layers.UpSampling2D(size=(2, 2))(x)\n",
    "    x = layers.Conv2DTranspose(filter_num[2],(3,3),strides=(1,1),padding='same', kernel_initializer = 'he_normal')(x)\n",
    "    x = layers.LeakyReLU(0.2)(x)   \n",
    "    x = layers.Conv2DTranspose(filter_num[2],(3,3),strides=(1,1),padding='same', kernel_initializer = 'he_normal')(x)\n",
    "    x = layers.LeakyReLU(0.2)(x) \n",
    "    \n",
    "    #### 256*256*32\n",
    "    \n",
    "    x = layers.UpSampling2D(size=(2, 2))(x)\n",
    "    x = layers.Conv2DTranspose(filter_num[3],(3,3),strides=(1,1),padding='same', kernel_initializer = 'he_normal')(x)\n",
    "    x = layers.LeakyReLU(0.2)(x)   \n",
    "    x = layers.Conv2DTranspose(filter_num[3],(3,3),strides=(1,1),padding='same', kernel_initializer = 'he_normal')(x)\n",
    "    x = layers.LeakyReLU(0.2)(x) \n",
    "    x = layers.Conv2DTranspose(3,(1,1),strides=(1,1),padding='same',activation='linear', kernel_initializer = 'he_normal')(x)\n",
    "    \n",
    "    #### 512*512*3\n",
    "\n",
    "    return keras.models.Model(generator_input,x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_discriminator(input_size):\n",
    "    height, width, channels = input_size\n",
    "    filter_num = [32,64,128,256]\n",
    "    \n",
    "    discriminator_input = layers.Input(shape=(height, width, channels))\n",
    "    \n",
    "    x = layers.Conv2D(filter_num[0], 3, padding = 'same', kernel_initializer = 'he_normal')(discriminator_input)\n",
    "    x = layers.LeakyReLU(0.2)(x)\n",
    "    x = layers.Conv2D(filter_num[0], 3, padding = 'same', kernel_initializer = 'he_normal')(x)\n",
    "    x = layers.LeakyReLU(0.2)(x)\n",
    "    x = layers.Conv2D(filter_num[0], 3, padding = 'same', kernel_initializer = 'he_normal')(x)\n",
    "    x = layers.LeakyReLU(0.2)(x)\n",
    "    x = layers.AveragePooling2D()(x)\n",
    "    \n",
    "    #### 256*256*32\n",
    "\n",
    "    x = layers.Conv2D(filter_num[1], 3, strides = 1, padding = 'same', kernel_initializer = 'he_normal')(x)\n",
    "    x = layers.LeakyReLU(0.2)(x)\n",
    "    x = layers.Conv2D(filter_num[1], 3, strides = 1, padding = 'same', kernel_initializer = 'he_normal')(x)\n",
    "    x = layers.LeakyReLU(0.2)(x)\n",
    "    x = layers.AveragePooling2D()(x)\n",
    "    \n",
    "    #### 128*128*64\n",
    "    \n",
    "    x = layers.Conv2D(filter_num[2], 3, strides = 1, padding = 'same', kernel_initializer = 'he_normal')(x)\n",
    "    x = layers.LeakyReLU(0.2)(x)\n",
    "    x = layers.Conv2D(filter_num[2], 3, strides = 1, padding = 'same', kernel_initializer = 'he_normal')(x)\n",
    "    x = layers.LeakyReLU(0.2)(x)\n",
    "    x = layers.AveragePooling2D()(x)\n",
    "    \n",
    "    #### 64*64*128\n",
    "    \n",
    "    x = layers.Conv2D(filter_num[3], 3, strides = 1, padding = 'same', kernel_initializer = 'he_normal')(x)\n",
    "    x = layers.LeakyReLU(0.2)(x)\n",
    "    x = layers.Conv2D(filter_num[3], 3, strides = 1, padding = 'same', kernel_initializer = 'he_normal')(x)\n",
    "    x = layers.LeakyReLU(0.2)(x)\n",
    "    x = layers.AveragePooling2D()(x)\n",
    "    \n",
    "    #### 32*32*256\n",
    "\n",
    "    x = layers.Flatten()(x)\n",
    "    x = layers.Dense(1, activation='sigmoid')(x)\n",
    "\n",
    "    return keras.models.Model(discriminator_input, x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_GAN(G, D):\n",
    "    D.trainable = False\n",
    "    gan_input = G.input\n",
    "    gan_output = D(G(gan_input))\n",
    "    gan = keras.models.Model(gan_input, gan_output)\n",
    "    return gan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_model(root_folder_name, model_dict):\n",
    "    model_path = './{}/model/'.format(root_folder_name)\n",
    "    os.makedirs(model_path, exist_ok=True)\n",
    "    \n",
    "    for key,model in model_dict.items():\n",
    "        model_json = model.to_json()\n",
    "        \n",
    "        with open(model_path + '{}.json'.format(key), 'w') as json_file:\n",
    "            json_file.write(model_json)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(root_folder_name, imgs, iterations=10000, bs=8, is_load_weight=False, pre_low_step=0):\n",
    "    start_time_all = time.time()\n",
    "    \n",
    "    # build net work\n",
    "    iterations = iterations\n",
    "    latent_dim = 200\n",
    "    height_1, width_1 = 512, 512\n",
    "\n",
    "    G1 = build_generator(latent_dim, (height_1, width_1))\n",
    "    D1 = build_discriminator((height_1, width_1, 3))\n",
    "    GAN1 = build_GAN(G1, D1)\n",
    "    \n",
    "    if is_load_weight:\n",
    "        model_path = './{}/model/'.format(root_folder_name)\n",
    "        weight_path = './{}/weight/record/'.format(root_folder_name)\n",
    "        \n",
    "        with open(model_path+'G1.json', 'r') as json_file:\n",
    "            temp = json_file.read()\n",
    "            G1 = model_from_json(temp)\n",
    "            G1.load_weights(weight_path+'g1_{}.h5'.format(pre_low_step))\n",
    "    \n",
    "        with open(model_path+'D1.json', 'r') as json_file:\n",
    "            temp = json_file.read()\n",
    "            D1 = model_from_json(temp)\n",
    "            D1.load_weights(weight_path+'d1_{}.h5'.format(pre_low_step))\n",
    "\n",
    "    optimizer = keras.optimizers.Adam(lr=0.0001, beta_1=0.5)\n",
    "    D1.compile(loss='binary_crossentropy', optimizer=optimizer)\n",
    "    GAN1 = build_GAN(G1, D1)\n",
    "    optimizer = keras.optimizers.Adam(lr=0.0001, beta_1=0.5)\n",
    "    GAN1.compile(loss='binary_crossentropy', optimizer=optimizer)\n",
    "\n",
    "    model_dict = {'G1':G1,'D1':D1,'GAN1':GAN1}\n",
    "    save_model(root_folder_name,model_dict)\n",
    "    \n",
    "    # create folder\n",
    "    os.makedirs('./{}/result_image/'.format(root_folder_name),exist_ok=True)\n",
    "    os.makedirs('./{}/weight/latest/'.format(root_folder_name),exist_ok=True)\n",
    "    os.makedirs('./{}/weight/record/'.format(root_folder_name),exist_ok=True)\n",
    "    save_dir = './{}/result_image/'.format(root_folder_name)\n",
    "    weight_path = './{}/weight/latest/'.format(root_folder_name)\n",
    "    weight_record_path = './{}/weight/record/'.format(root_folder_name)\n",
    "\n",
    "    # start training loop\n",
    "    start = 0\n",
    "    start_time = time.time()\n",
    "\n",
    "    low_iteration = 600\n",
    "    high_iteration = 1000\n",
    "    pre_high_step = 0\n",
    "    batch_size = bs\n",
    "    batch_num = len(imgs) // batch_size\n",
    "\n",
    "    imgs_temp = imgs[:batch_size * batch_num]\n",
    "\n",
    "    for step in range(iterations):\n",
    "        imgs_temp = shuffle(imgs_temp)\n",
    "\n",
    "        for low_step in range(batch_num):\n",
    "            real_images = imgs_temp[low_step*batch_size:(low_step+1)*batch_size]\n",
    "            real_images = (real_images-0.5)*2\n",
    "            random_latent_vectors = np.random.normal(size=(batch_size, latent_dim))\n",
    "\n",
    "            generated_images = G1.predict(random_latent_vectors)\n",
    "\n",
    "            labels = np.concatenate([np.zeros((batch_size, 1)),\n",
    "                                 np.ones((batch_size, 1))])\n",
    "\n",
    "            labels_real = 0.9*np.ones((batch_size, 1)) \n",
    "            labels_fake = np.zeros((batch_size, 1)) \n",
    "\n",
    "            d_loss_real = D1.train_on_batch(real_images, labels_real)\n",
    "            d_loss_fake = D1.train_on_batch(generated_images, labels_fake)\n",
    "            d_loss = 0.5*np.add(d_loss_real, d_loss_fake)\n",
    "\n",
    "            random_latent_vectors = np.random.normal(size=(batch_size, latent_dim))\n",
    "\n",
    "            misleading_targets = np.ones((batch_size, 1))\n",
    "\n",
    "            g_loss = GAN1.train_on_batch(random_latent_vectors, misleading_targets)\n",
    "\n",
    "            if low_step % 100 == 0:\n",
    "                # save model weights\n",
    "                G1.save_weights(weight_path+'g1.h5')\n",
    "                D1.save_weights(weight_path+'d1.h5')\n",
    "                \n",
    "                step_indicator = step*low_iteration+low_step+pre_low_step\n",
    "                \n",
    "                if step_indicator % 1000 == 0:\n",
    "                    G1.save_weights(weight_record_path+'g1_{}.h5'.format(step_indicator))\n",
    "                    D1.save_weights(weight_record_path+'d1_{}.h5'.format(step_indicator))\n",
    "\n",
    "                # print metrics\n",
    "                print('low resolution, discriminator loss at step %s: %s' % (step_indicator, d_loss))\n",
    "                print('low resolution, adversarial loss at step %s: %s' % (step_indicator, g_loss))\n",
    "                display_grid = np.zeros((4*height_1,width_1,3))\n",
    "\n",
    "                for j in range(4):\n",
    "                    display_grid[j*height_1:(j+1)*height_1,0:width_1,:] = generated_images[j]\n",
    "\n",
    "                img = image.array_to_img((display_grid[:,:,::-1]*127.5)+127.5, scale=False)\n",
    "                img.save(os.path.join(save_dir, 'low_generated_' + str(step*low_iteration+low_step+pre_low_step) + '.png'))\n",
    "                print(\"--- %s seconds ---\" % (time.time() - start_time))\n",
    "                start_time = time.time()        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__=='__main__':\n",
    "    \n",
    "    # load image\n",
    "    # imgs_length = len(glob(FILE_PATH))\n",
    "    imgs_length=4000\n",
    "    imgs = read_imgs(FILE_PATH, imgs_length)\n",
    "    \n",
    "    # train\n",
    "    train(label,imgs)\n",
    "    \n",
    "    # train with pretrain weights\n",
    "    # train(label, imgs, is_load_weight=True, pre_low_step=60000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 286
    },
    "colab_type": "code",
    "id": "gl7jcC7TdPTG",
    "outputId": "e9d7161f-1f84-44d1-dae8-5563b75d283f"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "dcgan.ipynb",
   "provenance": [],
   "toc_visible": true,
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}