hair_removal_GAN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "hair_removal_GAN.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyOIOEZJ3/nH+s9vx6rWJFIa",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/dlguswn3659/2020_Image_Lab/blob/master/hair_removal_GAN.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6Gchiw5u5n01",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\"\"\"\n",
        "Codes for training a hair-removal networking using WGAN-GP.\n",
        "For the WGAN-GP loss and algorithm parts,\n",
        "https://github.com/eriklindernoren/PyTorch-GAN/tree/master/implementations/wgan_gp is used\n",
        "\"\"\"\n",
        "\n",
        "\n",
        "from __future__ import print_function\n",
        "#%matplotlib inline\n",
        "import argparse\n",
        "import os\n",
        "import random\n",
        "import datetime\n",
        "import math\n",
        "import sys\n",
        "import csv\n",
        "from time import time\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.nn.parallel\n",
        "import torch.backends.cudnn as cudnn\n",
        "import torch.optim as optim\n",
        "import torch.autograd as autograd\n",
        "import torch.utils.data\n",
        "\n",
        "from torch.utils.data import DataLoader\n",
        "from torch.autograd import Variable\n",
        "\n",
        "from torchvision import datasets\n",
        "import torchvision.transforms as transforms\n",
        "import torchvision.utils as vutils\n",
        "\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fgvb3bDTBk6-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# ===================================\n",
        "#         discriminator\n",
        "# ===================================\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "class CNNDiscriminator(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(CNNDiscriminator, self).__init__()\n",
        "        self.model = nn.Sequential(\n",
        "            nn.Conv2d(3, 64, 4, 2, 1, bias=False),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "\n",
        "            nn.Conv2d(64, 128, 4, 2, 1, bias=False),\n",
        "            nn.BatchNorm2d(128),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "\n",
        "            nn.Conv2d(128, 256, 4, 2, 1, bias=False),\n",
        "            nn.BatchNorm2d(256),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "\n",
        "            nn.Conv2d(256, 512, 4, 2, 1, bias=False),\n",
        "            nn.BatchNorm2d(512),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "\n",
        "            nn.Conv2d(512, 1024, 4, 2, 1, bias=False),\n",
        "            # nn.BatchNorm2d(512),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "\n",
        "        )\n",
        "\n",
        "        self.linear = nn.Sequential(\n",
        "            nn.Linear(1024*4*4, 1),\n",
        "        )\n",
        "\n",
        "    def forward(self, img):\n",
        "        img = self.model(img)\n",
        "        img = img.reshape(img.size(0), -1)\n",
        "        validity = self.linear(img)\n",
        "        return validity"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "an6Mvi-PBrp1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# ===================================\n",
        "#         generator\n",
        "# ===================================\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "\n",
        "class ResizeCNNGenerator(nn.Module):\n",
        "    def __init__(self, in_channel, out_channel):\n",
        "        super(ResizeCNNGenerator, self).__init__()\n",
        "\n",
        "        self.in_channel = in_channel\n",
        "        self.out_channel = out_channel\n",
        "\n",
        "        # encoder layers\n",
        "        self.encoder1 = nn.Sequential(\n",
        "            nn.Conv2d(self.in_channel, 64, kernel_size=3, stride=1, padding=1, bias=False),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "        )\n",
        "        self.encoder2 = nn.Sequential(    \n",
        "            nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, bias=False),\n",
        "            nn.BatchNorm2d(64, 0.8),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "        )\n",
        "        self.encoder3 = nn.Sequential(    \n",
        "            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1, bias=False),\n",
        "            nn.BatchNorm2d(128, 0.8),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "        )\n",
        "        self.encoder4 = nn.Sequential(    \n",
        "            nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=1, bias=False),\n",
        "            nn.BatchNorm2d(128, 0.8),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "        )\n",
        "        self.encoder5 = nn.Sequential(    \n",
        "            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1, bias=False),\n",
        "            nn.BatchNorm2d(256, 0.8),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "        )\n",
        "\n",
        "        # decoder layers\n",
        "        self.decoder1 = nn.Sequential(\n",
        "            nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1, bias=False),\n",
        "            nn.BatchNorm2d(128, 0.8),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "        )\n",
        "        self.decoder2 = nn.Sequential(\n",
        "            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1, bias=False),\n",
        "            nn.BatchNorm2d(128, 0.8),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "        )\n",
        "        self.decoder3 = nn.Sequential(\n",
        "            nn.Upsample(scale_factor = 2, mode = 'nearest'), # increase size twice\n",
        "            nn.ReflectionPad2d(1),\n",
        "            nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=0, bias=False),\n",
        "            nn.BatchNorm2d(64, 0.8),\n",
        "            nn.LeakyReLU(0.2, inplace=True)\n",
        "        )\n",
        "        self.decoder4 = nn.Sequential(\n",
        "            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False),\n",
        "            nn.BatchNorm2d(64, 0.8),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "        )\n",
        "        self.decoder5 = nn.Sequential(\n",
        "            nn.Upsample(scale_factor = 2, mode = 'nearest'), # increase size twice\n",
        "            nn.ReflectionPad2d(1),\n",
        "            nn.Conv2d(64, self.out_channel, kernel_size=3, stride=1, padding=0, bias=False),\n",
        "            nn.Sigmoid()\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        # forward pass for encoder\n",
        "        out = self.encoder1(x)\n",
        "        out = self.encoder2(out)\n",
        "        out = self.encoder3(out)\n",
        "        out = self.encoder4(out)\n",
        "        out = self.encoder5(out)\n",
        "\n",
        "        # forward pass for decoder\n",
        "        out = self.decoder1(out)\n",
        "        out = self.decoder2(out)\n",
        "        out = self.decoder3(out)\n",
        "        out = self.decoder4(out)\n",
        "        out = self.decoder5(out)\n",
        "        return out"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Fzbtc-JsCJU3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# ===================================\n",
        "#         MelanomaDataset\n",
        "# ===================================\n",
        "\n",
        "import torch\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "from torch.utils.data import Dataset\n",
        "from PIL import Image\n",
        "\n",
        "\n",
        "class MelanomaDataset(Dataset):\n",
        "    \"\"\"Melanoma dataset\"\"\"\n",
        "\n",
        "    def __init__(self, csv_file, root_dir, label_type='target', img_format='dcm', transform=None):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "            csv_file (string): Path to the csv file with annotations.\n",
        "            root_dir (string): Directory with all the images.\n",
        "            label_type (string): Label type for each task.\n",
        "                                  * For the hair removal task -> 'hair'\n",
        "                                  * For the classification task -> 'target'\n",
        "            img_format (string): Image data type to load.\n",
        "                                  * dcm -> 'dcm'\n",
        "                                  * jpg -> 'jpg'\n",
        "            transform (callable, optional): Optional transform to be applied\n",
        "                on a sample.\n",
        "        \"\"\"\n",
        "        self.df = pd.read_csv(csv_file)\n",
        "        self.root_dir = root_dir\n",
        "        self.label_type = label_type\n",
        "        self.img_format = img_format\n",
        "        self.transform = transform\n",
        "    \n",
        "    def __len__(self):\n",
        "        return len(self.df)\n",
        "\n",
        "\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        if torch.is_tensor(idx):\n",
        "            idx = idx.tolist()\n",
        "\n",
        "        img_path = '{}/{}.{}'.format(self.root_dir, self.df.iloc[idx]['image_name'], self.img_format)\n",
        "\n",
        "        if self.img_format == 'jpg':\n",
        "            img = Image.open(img_path)\n",
        "            img = np.array(img) / 255\n",
        "            img = np.float32(img)\n",
        "        else:\n",
        "            ds = pydicom.read_file(img_path)\n",
        "            arr = ds.pixel_array\n",
        "            arr_scaled = arr / 255\n",
        "            img = arr_scaled\n",
        "            img = np.float32(img)\n",
        "        label = self.df.iloc[idx][self.label_type]\n",
        "\n",
        "        if self.transform:\n",
        "            img = self.transform(img)\n",
        "\n",
        "        data_dict = {'image': img, 'label': label}\n",
        "\n",
        "        return data_dict\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "# the following code snippet uses tensorflow.\n",
        "\n",
        "# import os\n",
        "# os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'\n",
        "# import tensorflow as tf\n",
        "# import tensorflow_io as tfio\n",
        "\n",
        "    # def __getitem__(self, idx):\n",
        "    #     if torch.is_tensor(idx):\n",
        "    #         idx = idx.tolist()\n",
        "\n",
        "    #     img_path = '{}/{}.{}'.format(self.root_dir, self.df.iloc[idx]['image_name'], self.img_format)\n",
        "\n",
        "    #     image_bytes = tf.io.read_file(img_path)\n",
        "    #     image = tfio.image.decode_dicom_image(image_bytes, dtype=tf.uint16)\n",
        "    #     print(image[0].shape)\n",
        "    #     print(image[0][0,0,:])\n",
        "    #     # print(image[0][0][0][0], image.shape)\n",
        "    #     img = image[0] / 255\n",
        "    #     img = np.float32(img)\n",
        "    #     label = self.df.iloc[idx][self.label_type]\n",
        "\n",
        "    #     if self.transform:\n",
        "    #         img = self.transform(img)\n",
        "\n",
        "    #     data_dict = {'image': img, 'label': label}\n",
        "\n",
        "    #     return data_dict"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "D9QpNv795r2G",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# ===================================\n",
        "#         Import custom codes\n",
        "# ===================================\n",
        "import config\n",
        "import utils\n",
        "from models import discriminator\n",
        "from models import generator\n",
        "from datasets import melanoma_dataset"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "W8EY4QIT6POf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# ===================================\n",
        "#         config.py\n",
        "# ===================================\n",
        "import argparse\n",
        "import os\n",
        "import numpy as np\n",
        "\n",
        "import torch\n",
        "\n",
        "def str2bool(v):\n",
        "    return v.lower() in ('true', '1')\n",
        "\n",
        "d = os.path.dirname\n",
        "parser = argparse.ArgumentParser(description='Unsupervised Segmentation incorporating Shape Prior via WGAN')\n",
        "path_arg = parser.add_argument_group('Experiment Config')\n",
        "path_arg = parser.add_argument_group('Data Config')\n",
        "path_arg = parser.add_argument_group('Networks Config')\n",
        "path_arg = parser.add_argument_group('Training Environment Config')\n",
        "path_arg = parser.add_argument_group('Coefficient Config')\n",
        "path_arg = parser.add_argument_group('Optimization Config')\n",
        "\n",
        "\n",
        "path_arg = parser.add_argument_group('Experiment Config')\n",
        "path_arg.add_argument('--file_prefix', type=str, default='experiment_name',\n",
        "    help='Path of model checkpoint to be save')\n",
        "path_arg.add_argument('--experiment_count', type=int, default=1,\n",
        "    help='Experiment count number to be used in filename')\n",
        "path_arg.add_argument('--dir_output', type=str, default='./csv',\n",
        "    help='Directory where output.csv will be stored')\n",
        "path_arg.add_argument('--sample_interval', type=int, default=200,\n",
        "    help='An interval to check training results from a sample')\n",
        "path_arg.add_argument('--num_plot_img', type=int, default=64,\n",
        "    help='The number of images to plot per batch')\n",
        "\n",
        "\n",
        "path_arg = parser.add_argument_group('Data Config')\n",
        "path_arg.add_argument('--dir_train_data_image', type=str, default='/home01/kaggle/train',\n",
        "    help='Directory of image data to be used')\n",
        "path_arg.add_argument('--dir_data_csv_hair', type=str, default='/home01/kaggle/train_hair.csv',\n",
        "    help='Directory of hair csv file')\n",
        "path_arg.add_argument('--dir_data_csv_non_hair', type=str, default='/home01/kaggle/train_non_hair.csv',\n",
        "    help='Directory of non hair csv file')\n",
        "path_arg.add_argument('--image_format', type=str, default='jpg',\n",
        "    help='Image format to be used')\n",
        "path_arg.add_argument('--height', type=int, default=64,\n",
        "    help='training image height to be resized by')\n",
        "path_arg.add_argument('--width', type=int, default=64,\n",
        "    help='training image width to be resized by')\n",
        "\n",
        "\n",
        "\n",
        "path_arg = parser.add_argument_group('Networks Config')\n",
        "path_arg.add_argument('--trained_ckpt_path', type=str, default=None,\n",
        "    help='Path of trained model checkpoint to be loaded')\n",
        "path_arg.add_argument('--num_in_channel', type=int, default=1,\n",
        "    help='Number of channel of input')\n",
        "path_arg.add_argument('--num_out_channel', type=int, default=1,\n",
        "    help='Number of channel of output')\n",
        "path_arg.add_argument('--network_d', type=str, default='vgg',\n",
        "    help='Network architecture to be used as a discriminator')\n",
        "path_arg.add_argument('--network_g', type=str, default='vgg',\n",
        "    help='Network architecture to be used as a generator')\n",
        "\n",
        "\n",
        "\n",
        "path_arg = parser.add_argument_group('Training Environment Config')\n",
        "path_arg.add_argument('--num_workers', type=int, default=16,\n",
        "    help='# of subprocesses to use for data loading for training')\n",
        "path_arg.add_argument('--multi_gpu', type=str2bool, default=True,\n",
        "    help='Decide whether to use multiple numbers of gpus or not')\n",
        "path_arg.add_argument('--num_gpu', type=int, default=4,\n",
        "    help='# of GPU to be used')\n",
        "path_arg.add_argument('--cuda_id', type=str, default='cuda:0',\n",
        "    help='GPU to be used')\n",
        "\n",
        "\n",
        "path_arg = parser.add_argument_group('Coefficient Config')\n",
        "path_arg.add_argument('--lambda_gp', type=int, default=10,\n",
        "    help='The size of gradient penalty')\n",
        "path_arg.add_argument('--lambda_distance', type=float, default=10,\n",
        "    help='The weight of a distance term in the generator loss')\n",
        "\n",
        "path_arg = parser.add_argument_group('Optimization Config')\n",
        "path_arg.add_argument('--num_epoch', type=int, default=5,\n",
        "    help='# of epochs to train for')\n",
        "path_arg.add_argument('--train_batch_size', type=int, default=64,\n",
        "    help='Batch size for training')\n",
        "path_arg.add_argument('--test_batch_size', type=int, default=64,\n",
        "    help='Batch size for testing')\n",
        "path_arg.add_argument('--lr_d', type=float, default=1e-3,\n",
        "    help='Fixed learning rate value for discriminator')\n",
        "path_arg.add_argument('--lr_g', type=float, default=1e-3,\n",
        "    help='Fixed learning rate value for generator')\n",
        "path_arg.add_argument('--beta1_d', type=float, default=0.5,\n",
        "    help='Beta1 hyperparam for Adam optimizers for discriminator')\n",
        "path_arg.add_argument('--beta1_g', type=float, default=0.5,\n",
        "    help='Beta1 hyperparam for Adam optimizers for generator')\n",
        "path_arg.add_argument('--beta2_d', type=float, default=0.999,\n",
        "    help='Beta2 hyperparam for Adam optimizers for discriminator')\n",
        "path_arg.add_argument('--beta2_g', type=float, default=0.999,\n",
        "    help='Beta2 hyperparam for Adam optimizers for generator')\n",
        "path_arg.add_argument('--num_discriminator', type=int, default=5,\n",
        "    help='The number of discriminator steps before one generator step')\n",
        "\n",
        "\n",
        "def get_config():\n",
        "    config = parser.parse_args()\n",
        "    print('[*] Configuration')\n",
        "    print(config)\n",
        "    \n",
        "    return config"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zss1VA_F6hQZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# ===================================\n",
        "#         utils.py\n",
        "# ===================================\n",
        "\n",
        "import numpy as np\n",
        "import torch\n",
        "\n",
        "\n",
        "def count_parameters(model):\n",
        "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
        "\n",
        "# to load dataset written in npy\n",
        "def npy_loader(path):\n",
        "    sample = torch.from_numpy(np.load(path))\n",
        "    return sample\n",
        "\n",
        "def requires_grad(model, flag=True):\n",
        "    for p in model.parameters():\n",
        "        p.requires_grad = flag"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Wby-fXSH5eFp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                    \n",
        "# ===================================\n",
        "#             Get config\n",
        "# ===================================\n",
        "train_config = config.get_config()\n",
        "config = train_config\n",
        "\n",
        "# config - experiment\n",
        "file_prefix = config.file_prefix\n",
        "experiment_count = config.experiment_count\n",
        "dir_output = config.dir_output\n",
        "sample_interval = config.sample_interval\n",
        "num_plot_img = config.num_plot_img\n",
        "\n",
        "# config - data\n",
        "dir_train_data_image = config.dir_train_data_image\n",
        "dir_data_csv_hair = config.dir_data_csv_hair\n",
        "dir_data_csv_non_hair = config.dir_data_csv_non_hair\n",
        "image_format = config.image_format\n",
        "height = config.height\n",
        "width = config.width\n",
        "\n",
        "# config - networks\n",
        "trained_ckpt_path = config.trained_ckpt_path\n",
        "num_in_channel = config.num_in_channel\n",
        "num_out_channel = config.num_out_channel\n",
        "network_d = config.network_d\n",
        "network_g = config.network_g\n",
        "\n",
        "# config - training env\n",
        "num_workers = config.num_workers\n",
        "multi_gpu = config.multi_gpu\n",
        "num_gpu = config.num_gpu\n",
        "cuda_id = config.cuda_id\n",
        "\n",
        "# config - coefficient\n",
        "lambda_gp = config.lambda_gp\n",
        "lambda_distance = config.lambda_distance \n",
        "\n",
        "# config - optimization\n",
        "num_epoch = config.num_epoch\n",
        "train_batch_size = config.train_batch_size\n",
        "test_batch_size = config.test_batch_size\n",
        "lr_d = config.lr_d\n",
        "lr_g = config.lr_g\n",
        "num_epoch = config.num_epoch\n",
        "beta1_d = config.beta1_d\n",
        "beta1_g = config.beta1_g\n",
        "beta2_d = config.beta2_d\n",
        "beta2_g = config.beta2_g\n",
        "num_discriminator = config.num_discriminator\n",
        "\n",
        "\n",
        "# ================================================\n",
        "#     Set Path & Files to Save Training Result\n",
        "# ================================================\n",
        "\n",
        "\n",
        "\n",
        "# Set file names or dirs to debug or save training result\n",
        "experiment_name = '{}_{}_lamdis{}_lamgp{}_lrd{}_lrg{}_bs{}_ndisc{}_nep{}_ex{}'.format(\n",
        "                                datetime.datetime.now().strftime('%Y%m%d'),\n",
        "                                file_prefix,\n",
        "                                lambda_distance,\n",
        "                                lambda_gp,\n",
        "                                lr_d,\n",
        "                                lr_g,\n",
        "                                train_batch_size,\n",
        "                                num_discriminator,\n",
        "                                num_epoch,\n",
        "                                experiment_count\n",
        "                                )\n",
        "experiment_name = experiment_name.replace(\".\",\"\")\n",
        "\n",
        "# dirs to save result files\n",
        "dir_output = '{}/output'.format(dir_output)\n",
        "dir_train_info = '{}/train_info'.format(dir_output)\n",
        "\n",
        "# dirs to save result images and learning curves\n",
        "dir_save_output_monitor_train = '{}/{}_monitor_train'.format(dir_output, experiment_name) # image to monitor train\n",
        "\n",
        "directories = [dir_output, dir_train_info, dir_save_output_monitor_train]\n",
        "\n",
        "# create dirs to save learning result if they don't exist.\n",
        "for directory in directories:\n",
        "    try:\n",
        "        os.mkdir(directory)\n",
        "        print(\"Directory \" , directory,  \" Created \") \n",
        "    except FileExistsError:\n",
        "        print(\"Directory \" , directory,  \" already exists\")\n",
        "\n",
        "\n",
        "# file path to save .txt file which contains training configuration\n",
        "txt_train_info = '{}/{}_train_info.txt'.format(dir_train_info, experiment_name)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uaahvhiR52x2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "# ===================================\n",
        "#  define functions to help training\n",
        "# ===================================\n",
        "\n",
        "# weights initialization function\n",
        "def weights_init(m):\n",
        "    classname = m.__class__.__name__\n",
        "    if classname.find('Conv') != -1:\n",
        "        nn.init.normal_(m.weight.data, 0.0, 0.02)\n",
        "    elif classname.find('BatchNorm') != -1:\n",
        "        nn.init.normal_(m.weight.data, 1.0, 0.02)\n",
        "        nn.init.constant_(m.bias.data, 0)\n",
        "\n",
        "\n",
        "# a function to compute gradient penalties for wgan\n",
        "torch.cuda.set_device(cuda_id)\n",
        "cuda = True if torch.cuda.is_available() else False\n",
        "Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n",
        "def compute_gradient_penalty(D, real_samples, fake_samples):\n",
        "    \"\"\"Calculates the gradient penalty loss for WGAN GP\"\"\"\n",
        "    # Random weight term for interpolation between real and fake samples\n",
        "    alpha = Tensor(np.random.random((real_samples.size(0), 1, 1, 1)))\n",
        "    # Get random interpolation between real and fake samples\n",
        "    interpolates = (alpha * real_samples + ((1 - alpha) * fake_samples)).requires_grad_(True)\n",
        "    d_interpolates = D(interpolates)\n",
        "    # fake = Variable(Tensor(real_samples.shape[0], 1).fill_(1.0), requires_grad=False)\n",
        "    fake = Variable(Tensor(real_samples.shape[0], 1).fill_(1.0), requires_grad=False).to(device)\n",
        "\n",
        "    # Get gradient w.r.t. interpolates\n",
        "    gradients = autograd.grad(\n",
        "        outputs=d_interpolates,\n",
        "        inputs=interpolates,\n",
        "        grad_outputs=fake,\n",
        "        create_graph=True,\n",
        "        retain_graph=True,\n",
        "        only_inputs=True,\n",
        "    )[0]\n",
        "    gradients = gradients.view(gradients.size(0), -1)\n",
        "    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()\n",
        "    return gradient_penalty\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F-aN1_Rc57av",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "# ===================================\n",
        "#             Load Data\n",
        "# ===================================\n",
        "\n",
        "\n",
        "\n",
        "transform_non_hair = transforms.Compose([\n",
        "                            transforms.ToTensor(),\n",
        "                            transforms.ToPILImage(),\n",
        "                            transforms.Resize((height, width)),\n",
        "                            transforms.ToTensor()\n",
        "])\n",
        "\n",
        "\n",
        "transform_train_hair = transforms.Compose([\n",
        "                            transforms.ToTensor(),\n",
        "                            transforms.ToPILImage(),\n",
        "                            transforms.Resize((height, width)),\n",
        "                            transforms.ToTensor()\n",
        "])\n",
        "\n",
        "dataset_non_hair = melanoma_dataset.MelanomaDataset(dir_data_csv_non_hair, dir_train_data_image, 'hair', image_format, transform=transform_non_hair)\n",
        "dataset_train_hair = melanoma_dataset.MelanomaDataset(dir_data_csv_hair, dir_train_data_image, 'hair', image_format, transform=transform_train_hair)\n",
        "\n",
        "\n",
        "loader_non_hair = torch.utils.data.DataLoader(\n",
        "    dataset = dataset_non_hair,\n",
        "    batch_size=train_batch_size,\n",
        "    shuffle=True,\n",
        "    num_workers=num_workers,\n",
        "    drop_last=True)\n",
        "\n",
        "loader_train_hair = torch.utils.data.DataLoader(\n",
        "    dataset = dataset_train_hair,\n",
        "    batch_size=train_batch_size,\n",
        "    shuffle=True,\n",
        "    num_workers=num_workers,\n",
        "    drop_last=True)\n",
        "\n",
        "\n",
        "# ===================================\n",
        "#           Set Train Env\n",
        "# ===================================\n",
        "if multi_gpu == True:\n",
        "    device = torch.device(\"cuda:0\" if (torch.cuda.is_available()) else \"cpu\")\n",
        "    ngpu = num_gpu # should be modified to specify multiple gpu ids to be used\n",
        "elif cuda_id:\n",
        "    device = torch.device(cuda_id)\n",
        "    ngpu = 1\n",
        "else:\n",
        "    device = torch.device('cpu')\n",
        "    ngpu = 0\n",
        "\n",
        "\n",
        "\n",
        "# ===================================\n",
        "#              Set Model\n",
        "# ===================================\n",
        "# init discriminator model instance\n",
        "if network_d == 'resnet':\n",
        "    pass # to be implemented\n",
        "elif network_d =='vgg':\n",
        "    model_d = discriminator.CNNDiscriminator().to(device)\n",
        "else:\n",
        "    raise ValueError('There is no such discriminator model')\n",
        "\n",
        "# init segmentation model instance\n",
        "if network_g == 'stylegan':\n",
        "    pass # to be implemented\n",
        "elif network_g =='resnet':\n",
        "    pass # to be implemented\n",
        "elif network_g =='vgg':\n",
        "    model_g = generator.ResizeCNNGenerator(num_in_channel, num_out_channel).to(device)\n",
        "elif network_g =='unet':\n",
        "    pass # to be implemented\n",
        "else:\n",
        "    raise ValueError('There is no such generator model')\n",
        "\n",
        "# decide whether to use multiple gpus\n",
        "if multi_gpu == True:\n",
        "    model_d = nn.DataParallel(model_d, list(range(num_gpu)))\n",
        "    model_g = nn.DataParallel(model_g, list(range(num_gpu)))\n",
        "\n",
        "# init model parameters\n",
        "model_d.apply(weights_init)\n",
        "if network_g =='vgg':\n",
        "    model_g.apply(weights_init)\n",
        "\n",
        "# optimizers\n",
        "optimizer_d = torch.optim.Adam(model_d.parameters(), lr=lr_d, betas=(beta1_d, beta2_d))\n",
        "optimizer_g = torch.optim.Adam(model_g.parameters(), lr=lr_g, betas=(beta1_g, beta2_g))\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UYJE3Z7o5_Xl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "# ===================================\n",
        "#           Save Train Info\n",
        "# ===================================\n",
        "with open(txt_train_info, 'a') as t:\n",
        "    t.write('[Data]     Hair: {}'.format(dir_data_csv_hair) + os.linesep)\n",
        "    t.write('[Data]     Non Hair: {}'.format(dir_data_csv_non_hair) + os.linesep)\n",
        "\n",
        "    t.write('[Optim]    Total Epoch: {}'.format(num_epoch) + os.linesep)\n",
        "    t.write('[Optim]    Train Batch Size: {}'.format(train_batch_size) + os.linesep)\n",
        "    t.write('[Optim]    Test Batch Size: {}'.format(test_batch_size) + os.linesep)\n",
        "    t.write('[Optim]    Learning Rate (Discriminator): {}'.format(lr_d) + os.linesep)\n",
        "    t.write('[Optim]    Learning Rate (Generator): {}'.format(lr_g) + os.linesep)\n",
        "    t.write('[Optim]    Beta1 (Discriminator): {}'.format(beta1_d) + os.linesep)\n",
        "    t.write('[Optim]    Beta1 (Generator): {}'.format(beta1_g) + os.linesep)\n",
        "    t.write('[Optim]    Beta2 (Discriminator): {}'.format(beta2_d) + os.linesep)\n",
        "    t.write('[Optim]    Beta2 (Generator): {}'.format(beta2_g) + os.linesep)\n",
        "    t.write('[Optim]    Num Discriminator: {}'.format(num_discriminator) + os.linesep)\n",
        "\n",
        "    t.write('[Coeff]    [Discriminator] Lambda GP: {}'.format(lambda_gp) + os.linesep)\n",
        "    t.write('[Coeff]    [Generator] Lambda Distance: {}'.format(lambda_distance) + os.linesep)\n",
        "\n",
        "    t.write('[Network]  In Channel: {}'.format(num_in_channel) + os.linesep)\n",
        "    t.write('[Network]  Out Channel: {}'.format(num_out_channel) + os.linesep)\n",
        "    t.write('[Network]  Discriminator Model: {}'.format(network_d) + os.linesep)\n",
        "    t.write('[Network]  Generator Model: {}'.format(network_g) + os.linesep)\n",
        "\n",
        "    t.write('**************** Discriminator Structure ****************' + os.linesep)\n",
        "    t.write('[*] num of params: {}'.format(utils.count_parameters(model_d)) + os.linesep)\n",
        "    t.write('[*] model structure: {}'.format(model_d) + os.linesep)\n",
        "\n",
        "    t.write('**************** Generator Model Structure ****************' + os.linesep)\n",
        "    t.write('[*] num of params: {}'.format(utils.count_parameters(model_g)) + os.linesep)\n",
        "    t.write('[*] model structure: {}'.format(model_g) + os.linesep)\n",
        "\n",
        "\n",
        "\n",
        "# ==================================================\n",
        "#      Init variables to save training results\n",
        "# ==================================================\n",
        "losses_d = []\n",
        "losses_g = []\n",
        "losses_distance = []\n",
        "\n",
        "# ==============================\n",
        "#         Start Training\n",
        "# ==============================\n",
        "time_total_start = time() # set a start time to monitor training time\n",
        "for epoch in range(num_epoch):\n",
        "    time_train_epoch_start = time() # set an epoch start time to monitor training time per epoch\n",
        "    if epoch == 0:\n",
        "        print('device:', device)\n",
        "    print(\"start train epoch{}:\".format(epoch))\n",
        "    dataloader_iterator = iter(loader_non_hair)\n",
        "\n",
        "    num_iters_d = 0\n",
        "    num_iters_g = 0\n",
        "\n",
        "    # ================================================\n",
        "    #                     Train\n",
        "    # ================================================\n",
        "    model_d.train()\n",
        "    model_g.train()\n",
        "\n",
        "    for i, data_hair in enumerate(loader_train_hair, 0):\n",
        "\n",
        "        # iterate dataloader_input at the same time\n",
        "        try:\n",
        "            data_non_hair = next(dataloader_iterator)\n",
        "        except StopIteration:\n",
        "            dataloader_iterator = iter(loader_non_hair)\n",
        "            data_non_hair = next(dataloader_iterator)\n",
        "\n",
        "        imgs_non_hair = Variable(data_non_hair['image'].type(Tensor)).to(device)\n",
        "        imgs_hair = Variable(data_hair['image'].type(Tensor)).to(device)\n",
        "\n",
        "\n",
        "        # -------------------------------\n",
        "        #       Train Discriminator\n",
        "        # ------------------------------- \n",
        "\n",
        "        # get hair-removed output from the generator model\n",
        "        output_hair_removed = model_g(imgs_hair)\n",
        "\n",
        "        # remove gradients on the computational graph of the discriminator\n",
        "        optimizer_d.zero_grad()\n",
        "        # get scores of non-hair images from the discriminator\n",
        "        validity_non_hair = model_d(imgs_non_hair)\n",
        "        # get scores of hair-removed images from the discriminator\n",
        "        validity_hair_removed = model_d(output_hair_removed)\n",
        "\n",
        "        # compute a loss for the discriminator\n",
        "        # compute Gradient_penalty for WGAN-GP loss\n",
        "        gradient_penalty = compute_gradient_penalty(model_d, imgs_non_hair.data, output_hair_removed.data)\n",
        "        # Adversarial loss using WGAN-GP\n",
        "        loss_d = -torch.mean(validity_non_hair) + torch.mean(validity_hair_removed) + (lambda_gp * gradient_penalty)\n",
        "        losses_d.append(loss_d.item())\n",
        "        loss_d.backward()\n",
        "\n",
        "        # optimize discriminator\n",
        "        optimizer_d.step()\n",
        "        num_iters_d += 1\n",
        "\n",
        "        if i % num_discriminator == 0:\n",
        "            model_g.train()\n",
        "            # ---------------------\n",
        "            #    Train Generator\n",
        "            # ---------------------\n",
        "            # remove gradients on the computational graph of the generator\n",
        "            optimizer_g.zero_grad()\n",
        "            # get hair-removed output from the generator model\n",
        "            output_hair_removed = model_g(imgs_hair)\n",
        "            validity_fake = model_d(output_hair_removed)\n",
        "\n",
        "            # compute a loss for the generator\n",
        "            # compute distance between the original hair images and their hair-removed outputs of generator\n",
        "            distance = F.l1_loss(imgs_hair, output_hair_removed)\n",
        "            # Final loss = Adversarial loss using WGAN-GP + distance\n",
        "            loss_g = -torch.mean(validity_fake) + (lambda_distance * distance)\n",
        "            losses_g.append(loss_g.item())\n",
        "            losses_distance.append(distance.item())\n",
        "            loss_g.backward()\n",
        "\n",
        "            # optimize generator\n",
        "            optimizer_g.step()\n",
        "            num_iters_g += 1\n",
        "\n",
        "            # ---------------------\n",
        "            #     Get Results\n",
        "            # ---------------------\n",
        "            print(\n",
        "                \"#-TRAIN-# [Epoch %d/%d] [Batch %d/%d] [G loss: %f] [D loss: %f] [Distance: %f]\"\n",
        "                % (epoch+1, num_epoch, i+1, len(loader_train_hair), loss_g.item(), loss_d.item(), distance.item())\n",
        "            )\n",
        "\n",
        "            total_time = time() - time_total_start\n",
        "            total_hours, _total_rest = divmod(total_time, 3600)\n",
        "            total_mins, total_secs = divmod(_total_rest, 60)\n",
        "            print('#-TRAIN-# Total Running Time: {}h {}m {}s'.format(total_hours, total_mins, total_secs))\n",
        "            train_epoch_time = time() - time_train_epoch_start\n",
        "            train_epoch_hours, _train_epoch_rest = divmod(train_epoch_time, 3600)\n",
        "            train_epoch_mins, train_epoch_secs = divmod(_train_epoch_rest, 60)\n",
        "            print('#-TRAIN-# Epoch Time: {}h {}m {}s'.format(train_epoch_hours, train_epoch_mins, train_epoch_secs))\n",
        "\n",
        "\n",
        "            # -----------------------------------------\n",
        "            #      Plot images to monitor training\n",
        "            # -----------------------------------------\n",
        "            if num_iters_g in [0, 1] or (num_iters_g * num_discriminator) % sample_interval == 0:\n",
        "\n",
        "                plot_hair_removed = output_hair_removed.cpu().detach().numpy()\n",
        "                plot_hair = imgs_hair.cpu().detach().numpy()\n",
        "                plot_non_hair = imgs_non_hair.cpu().detach().numpy()\n",
        "\n",
        "                plot_hair_removed = vutils.make_grid(torch.from_numpy(plot_hair_removed[:num_plot_img]), padding=2, pad_value=1)\n",
        "                plot_hair = vutils.make_grid(torch.from_numpy(plot_hair[:num_plot_img]), padding=2, pad_value=1)\n",
        "                plot_non_hair = vutils.make_grid(torch.from_numpy(plot_non_hair[:num_plot_img]), padding=2, pad_value=1)\n",
        "\n",
        "                imgs = [[plot_hair_removed, plot_hair, plot_non_hair]]\n",
        "                imgs_list = [plot_hair_removed, plot_hair, plot_non_hair]\n",
        "                imgs_names = ['hair removed', 'hair input', 'non-hair']\n",
        "                fig, axes = plt.subplots(len(imgs), len(imgs[0]), figsize=(18,18))\n",
        "                for plot_i, ax in enumerate(axes.flat):\n",
        "                    ax.axis(\"off\")\n",
        "                    ax.set_title(imgs_names[plot_i])\n",
        "                    ax.imshow(np.transpose(imgs_list[plot_i],(1,2,0)), vmin=0.0, vmax=1.0)\n",
        "                    if plot_i + 1 == len(imgs_list):\n",
        "                        break\n",
        "                plt.show()\n",
        "                file_name = '{}/results_{}_{}'.format(dir_save_output_monitor_train, epoch+1, i+1)\n",
        "                fig.savefig(file_name, bbox_inches='tight', pad_inches=0.1)\n",
        "                plt.clf()\n",
        "                plt.close()\n",
        "\n",
        "                # Plot loss & metric score curves\n",
        "                curve_titles = [\n",
        "                    \"Discriminator Loss\",\n",
        "                    \"Generator Loss\",\n",
        "                    \"Distance\"\n",
        "                                ]\n",
        "                curve_data = [[losses_d], [losses_g], [losses_distance]]\n",
        "                curve_labels = [[\"loss_d\"], [\"loss_g\"], [\"distance\"]]\n",
        "                curve_xlabels = [\"iterations\", \"iterations\", \"iterations\"]\n",
        "                curve_ylabels = [\"loss\", \"loss\", \"distance\"]\n",
        "                curve_filenames = [\"learn-curve-loss-d\", \"learn-curve-loss-g\", \"learn-curve-distance\"]\n",
        "\n",
        "                \n",
        "                for i_curve, curve_data in enumerate(curve_data):\n",
        "                    plt.figure(figsize=(10,5))\n",
        "                    plt.title(curve_titles[i_curve])\n",
        "                    for i_curve_data, curve_data_item in enumerate(curve_data):\n",
        "                        plt.plot(curve_data_item,label=curve_labels[i_curve][i_curve_data])\n",
        "                    plt.xlabel(curve_xlabels[i_curve])\n",
        "                    plt.ylabel(curve_ylabels[i_curve])\n",
        "                    plt.legend()\n",
        "                    file_name = '{}/{}'.format(dir_save_output_monitor_train, curve_filenames[i_curve])\n",
        "                    plt.show()\n",
        "                    plt.savefig(file_name, bbox_inches='tight', pad_inches=0.1)\n",
        "                    plt.clf()\n",
        "                    plt.close()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}