ColorizeMountainAdversarial.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "ColorizeMountainAdversarial.ipynb",
      "provenance": [],
      "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/CanKeles5/ColorizeMountainAdversarial/blob/master/ColorizeMountainAdversarial.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ci2A-4mfcsr3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%reload_ext autoreload\n",
        "%autoreload 2\n",
        "%matplotlib inline"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "h2f8qdaTdbMU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torchvision\n",
        "import torchvision.datasets as datasets\n",
        "import torchvision.transforms as transforms\n",
        "import matplotlib.pyplot as plt\n",
        "import torch.utils.data.sampler\n",
        "import torchvision.transforms.functional as TF\n",
        "from torch.utils.data import DataLoader as DL\n",
        "from torch.utils.data import *\n",
        "from PIL import Image, ImageFilter\n",
        "import os\n",
        "import cv2\n",
        "import numpy\n",
        "import random\n",
        "import fnmatch"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sGLRU8cBdbJj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "print(device)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ppj81AKddbFV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "! mkdir -p ~/.kaggle/\n",
        "! mv kaggle.json ~/.kaggle/"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aCUm1ywadbDI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "! kaggle datasets download -d puneet6060/intel-image-classification"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dEFJiFj0dbAp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "path = \"/content/intel-image-classification.zip\"\n",
        "to = \"/content/dataset\"\n",
        "! unzip -q -n {path} -d {to}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pYBZD_wuda8U",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "image_folder = \"/content/dataset/seg_train/seg_train/mountain\"\n",
        "image_paths = []\n",
        "\n",
        "for dirname, _, filenames in os.walk(image_folder):\n",
        "    for filename in filenames:\n",
        "        if(fnmatch.fnmatch(dirname, '*mountain*')):\n",
        "            image_paths.append(os.path.join(dirname, filename))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RXOcu6hBda41",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class MyDataset(Dataset):\n",
        "  def __init__(self, image_paths, train=True):\n",
        "    self.image_paths = image_paths\n",
        "  \n",
        "  def transforms(self, image):\n",
        "    tfms = transforms.Resize(size=(256, 256))\n",
        "    \n",
        "    tfms = transforms.Compose([\n",
        "                              transforms.Resize(size=(256, 256)),\n",
        "                              transforms.RandomHorizontalFlip(),\n",
        "                              transforms.RandomRotation(degrees=5),\n",
        "                              transforms.RandomPerspective(distortion_scale=0.05)\n",
        "                              ])\n",
        "\n",
        "    image = tfms(image)\n",
        "    image = image.filter(ImageFilter.MedianFilter())\n",
        "    image = TF.to_tensor(image)\n",
        "    \n",
        "    return image\n",
        "\n",
        "  def __getitem__(self, index):\n",
        "    image = Image.open(self.image_paths[index])\n",
        "    x = self.transforms(image)\n",
        "\n",
        "    return x\n",
        "  \n",
        "  def __len__(self):\n",
        "    return len(self.image_paths)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oHB4u2gJdij0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class ValidateSet(Dataset):\n",
        "  def __init__(self, image_paths, train=True):\n",
        "    self.image_paths = image_paths\n",
        "  \n",
        "  def transforms(self, image):\n",
        "    resize = transforms.Resize(size=(256, 256))\n",
        "    image = resize(image)\n",
        "    image = image.filter(ImageFilter.MedianFilter())\n",
        "    image = TF.to_tensor(image)\n",
        "    \n",
        "    return image\n",
        "\n",
        "  def __getitem__(self, index):\n",
        "    image = Image.open(self.image_paths[index])\n",
        "    x = self.transforms(image)\n",
        "\n",
        "    return x\n",
        "  \n",
        "  def __len__(self):\n",
        "    return len(self.image_paths)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hltUu5IMdig5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "dataset = MyDataset(image_paths[0:2500])\n",
        "len(dataset)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iaxZln7ndieu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "validate_set = ValidateSet(image_paths[2500: 2512])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yGLjhHZUdicd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "train_indices = range(0,2300)\n",
        "test_indices = range(2300, 2500)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0sW9t16rdiaU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "batch_size = 4\n",
        "train_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, sampler=SubsetRandomSampler(train_indices))\n",
        "test_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, sampler=SubsetRandomSampler(test_indices))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jeOBfyU4diXz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def en_double_conv(in_channels, out_channels):\n",
        "    return nn.Sequential(\n",
        "        nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1),\n",
        "        nn.LeakyReLU(0.2, inplace=True),\n",
        "        nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),\n",
        "        nn.BatchNorm2d(out_channels),\n",
        "        nn.LeakyReLU(0.2, inplace=True)\n",
        "    )\n",
        "\n",
        "def dec_double_conv(in_channels, out_channels):\n",
        "  return nn.Sequential(\n",
        "      nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1),\n",
        "      nn.ReLU(inplace=True),\n",
        "      nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),\n",
        "      nn.BatchNorm2d(out_channels),\n",
        "      nn.ReLU(inplace=True)\n",
        "  )\n",
        "\n",
        "n=8\n",
        "\n",
        "class G(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "\n",
        "        self.dconv_1 = en_double_conv(1, n)\n",
        "        self.dconv_2 = en_double_conv(n, n*2)\n",
        "        self.dconv_3 = en_double_conv(n*2, n*4)\n",
        "        self.dconv_4 = en_double_conv(n*4, n*8)\n",
        "        self.dconv_5 = en_double_conv(n*8, n*8)\n",
        "        self.dconv_6 = en_double_conv(n*8, n*8)\n",
        "        self.dconv_7 = en_double_conv(n*8, n*8)\n",
        "        self.dconv_8 = en_double_conv(n*8, n*8)\n",
        "\n",
        "        self.dropout = nn.Dropout(0.5)\n",
        "        self.maxpool = nn.MaxPool2d(2)\n",
        "\n",
        "        self.TConv8 = nn.ConvTranspose2d(n*8, n*8, 4, 2, 1)\n",
        "        self.TConv7 = nn.ConvTranspose2d(n*8*2, n*8*2, 4, 2, 1)\n",
        "        self.TConv6 = nn.ConvTranspose2d(n*8*3, n*8*3, 4, 2, 1)\n",
        "        self.TConv5 = nn.ConvTranspose2d(n*8*4, n*8*4, 4, 2, 1)\n",
        "        self.TConv4 = nn.ConvTranspose2d(n*8*5, n*8*5, 4, 2, 1)\n",
        "        self.TConv3 = nn.ConvTranspose2d(n*44, n*44, 4, 2, 1)\n",
        "        self.TConv2 = nn.ConvTranspose2d(n*46, n*46, 4, 2, 1)\n",
        "        self.TConv1 = nn.ConvTranspose2d(n*47, 3, 4, 2, 1)\n",
        "        \n",
        "    def forward(self, x):\n",
        "        conv1 = self.dconv_1(x)\n",
        "        conv1 = self.maxpool(conv1)\n",
        "\n",
        "        conv2 = self.dconv_2(conv1)\n",
        "        conv2 = self.maxpool(conv2)\n",
        "\n",
        "        conv3 = self.dconv_3(conv2)\n",
        "        conv3 = self.maxpool(conv3)\n",
        "\n",
        "        conv4 = self.dconv_4(conv3)\n",
        "        conv4 = self.maxpool(conv4)\n",
        "\n",
        "        conv5 = self.dconv_5(conv4)\n",
        "        conv5 = self.maxpool(conv5)\n",
        "\n",
        "        conv6 = self.dconv_6(conv5)\n",
        "        conv6 = self.maxpool(conv6)\n",
        "\n",
        "        conv7 = self.dconv_7(conv6)\n",
        "        conv7 = self.maxpool(conv7)\n",
        "\n",
        "        conv8 = self.dconv_8(conv7)\n",
        "        conv8 = self.maxpool(conv8)\n",
        "\n",
        "        x = self.TConv8(conv8)\n",
        "\n",
        "        x = torch.cat([x, conv7], dim=1)\n",
        "        x = self.TConv7(x)\n",
        "        x = self.dropout(x)\n",
        "\n",
        "        x = torch.cat([x, conv6], dim=1)\n",
        "        x = self.TConv6(x)\n",
        "        x = self.dropout(x)\n",
        "\n",
        "        x = torch.cat([x, conv5], dim=1)\n",
        "        x = self.TConv5(x)\n",
        "        x = self.dropout(x)\n",
        "\n",
        "        x = torch.cat([x, conv4], dim=1)\n",
        "        x = self.TConv4(x)\n",
        "        x = self.dropout(x)\n",
        "\n",
        "        x = torch.cat([x, conv3], dim=1)\n",
        "        x = self.TConv3(x)\n",
        "        x = self.dropout(x)\n",
        "\n",
        "        x = torch.cat([x, conv2], dim=1)\n",
        "        x = self.TConv2(x)\n",
        "        x = self.dropout(x)\n",
        "\n",
        "        x = torch.cat([x, conv1], dim=1)\n",
        "        x = self.TConv1(x)\n",
        "        x = nn.Tanh()(x)\n",
        "\n",
        "        return x"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4XzavRU4diSa",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class D(nn.Module):\n",
        "  def __init__(self):\n",
        "    super(D, self).__init__()\n",
        "    self.main = nn.Sequential(        \n",
        "        nn.Conv2d(3, 16, 2, 2, 0),\n",
        "        nn.LeakyReLU(0.2, inplace=True),\n",
        "        \n",
        "        nn.Conv2d(16, 32, 2, 2, 0),\n",
        "        nn.BatchNorm2d(32),\n",
        "        nn.LeakyReLU(0.2, inplace=True),\n",
        "        \n",
        "        nn.Conv2d(32, 64, 2, 2, 0),\n",
        "        nn.BatchNorm2d(64),\n",
        "        nn.LeakyReLU(0.2, inplace=True),\n",
        "        \n",
        "        nn.Conv2d(64, 128, 2, 2, 0),\n",
        "        nn.BatchNorm2d(128),\n",
        "        nn.LeakyReLU(0.2, inplace=True),\n",
        "        \n",
        "        nn.Conv2d(128, 256, 2, 2, 0),\n",
        "        nn.BatchNorm2d(256),\n",
        "        nn.LeakyReLU(0.2, inplace=True),\n",
        "        \n",
        "        nn.Conv2d(256, 512, 2, 2, 0),\n",
        "        nn.BatchNorm2d(512),\n",
        "        nn.LeakyReLU(0.2, inplace=True),\n",
        "        \n",
        "        nn.Conv2d(512, 1024, 2, 2, 0),\n",
        "        nn.BatchNorm2d(1024),\n",
        "        nn.LeakyReLU(0.2, inplace=True),\n",
        "\n",
        "        nn.Conv2d(1024, 1, 2, 2, 0),\n",
        "        nn.Sigmoid()\n",
        "    )\n",
        "    \n",
        "  def forward(self, im):\n",
        "    return self.main(im)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "M-q1IF_RdiPq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "Generator = G().to(device)\n",
        "Discriminator = D().to(device)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fu62sNafd0cl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def weights_init(m):\n",
        "    if isinstance(m, nn.Conv2d):\n",
        "        m.weight.data.normal_(0, 0.02)\n",
        "        m.bias.data.normal_(0, 0.001)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9XJ1Ig68d0aR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "weights_init(Generator)\n",
        "weights_init(Discriminator)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wDx0YIhmd0Xm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(\"Number of parameters in Generator: \", sum([p.numel() for p in Generator.parameters()]))\n",
        "print(\"Number of parameters in Discriminator: \", sum([p.numel() for p in Discriminator.parameters()]))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iWdHjOStd0UT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "criterion = nn.BCELoss()\n",
        "adv_criterion = nn.BCELoss()\n",
        "l1_criterion = nn.L1Loss()\n",
        "G_optim = torch.optim.Adam(Generator.parameters(), lr=6e-6)\n",
        "D_optim = torch.optim.Adam(Discriminator.parameters(), lr=6e-6)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fmpUtOPDd0Rt",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "Discriminator.train()\n",
        "Generator.train()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t-ihucFGd0O5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def save_pic(epoch_no):\n",
        "  Generator.eval()\n",
        "  im = dataset[0]\n",
        "  im = (0.2989*im[0,:,:] + 0.5870*im[1,:,:] + 0.1140*im[2,:,:])\n",
        "  im = im.unsqueeze(0).unsqueeze(0).cuda()\n",
        "  output = Generator(im)\n",
        "\n",
        "  p = output[0].detach().cpu()\n",
        "  p = p.clamp(0.0, 1.0)\n",
        "  \n",
        "  PIL_img = transforms.ToPILImage()(p)\n",
        "  PIL_img = PIL_img.save(str(epoch_no) + \".jpg\")\n",
        "  Generator.train()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NjEGaF1ed0Lf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "D_losses_train = []\n",
        "G_losses_train = []\n",
        "\n",
        "D_losses_test = []\n",
        "G_losses_test = []"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_D_n6-2ieMyg",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def shuffle_data(fake_im, real_im):\n",
        "  batch_size=fake_im.shape[0]\n",
        "  data=torch.cat((fake_im, real_im),dim=0)\n",
        "  labels=torch.cat((torch.zeros(batch_size), torch.ones(batch_size)))\n",
        "  \n",
        "  return data, labels"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EV8phKgUeMv7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(len(train_indices))\n",
        "print(len(test_indices))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GDLRl8I8eMtM",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "n_epochs = 300\n",
        "\n",
        "for epoch in range(n_epochs):\n",
        "  D_train_loss = 0.0\n",
        "  G_train_loss = 0.0\n",
        "\n",
        "  D_test_loss = 0.0\n",
        "  G_test_loss = 0.0\n",
        "  for i, real_im in enumerate(train_loader):\n",
        "    gray_scale_im=(0.2989*real_im[:,0,:,:] + 0.5870*real_im[:,1,:,:] + 0.1140*real_im[:,2,:,:])\n",
        "    gray_scale_im=gray_scale_im.unsqueeze(1).to(device)\n",
        "    ##########Train the discriminator##########\n",
        "    D_optim.zero_grad()\n",
        "    real_im=real_im.to(device)\n",
        "    fake_img = Generator(gray_scale_im)\n",
        "\n",
        "    data, labels = shuffle_data(fake_img, real_im)\n",
        "    guess = Discriminator(data)\n",
        "\n",
        "    D_loss = criterion(guess, labels.to(device))\n",
        "    D_train_loss += D_loss.item()\n",
        "    D_loss.backward()\n",
        "    D_optim.step()\n",
        "    ###########################################\n",
        "    \n",
        "    ############Train the generator############\n",
        "    G_optim.zero_grad()\n",
        "    fake_img = Generator(gray_scale_im)\n",
        "    guess = Discriminator(fake_img).view(-1)\n",
        "    G_loss_adv = adv_criterion(guess, torch.ones(4).to(device))\n",
        "    G_loss_l1 = l1_criterion(fake_img, real_im.to(device))\n",
        "    G_loss = G_loss_adv + G_loss_l1*100\n",
        "    G_train_loss += G_loss.item()\n",
        "    G_loss.backward()\n",
        "    G_optim.step()\n",
        "    ###########################################\n",
        "  \n",
        "  for i, real_im in enumerate(test_loader):\n",
        "    real_im=real_im.to(device)\n",
        "    gray_scale_im=(0.2989*real_im[:,0,:,:] + 0.5870*real_im[:,1,:,:] + 0.1140*real_im[:,2,:,:])\n",
        "    gray_scale_im=gray_scale_im.unsqueeze(1).to(device)\n",
        "\n",
        "    fake_img = Generator(gray_scale_im)\n",
        "\n",
        "    guess = Discriminator(fake_img)\n",
        "    adv_loss = adv_criterion(guess, torch.ones(4).to(device))\n",
        "    l1_loss = l1_criterion(fake_img, real_im)\n",
        "    G_loss = adv_loss + l1_loss*100\n",
        "    G_test_loss += G_loss.item()\n",
        "  \n",
        "  G_train_loss = G_train_loss/len(train_indices)\n",
        "  G_test_loss = G_test_loss/len(test_indices)\n",
        "  print(\"Epoch \" + str(epoch) + \", Train: \" + str(G_train_loss) + \" , Test: \" + str(G_test_loss))\n",
        "  G_losses_train.append(G_train_loss)\n",
        "  G_losses_test.append(G_test_loss)\n",
        "  save_pic(epoch)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SV_A6IiFeMqi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "plt.plot(G_losses_train)\n",
        "plt.plot(G_losses_test)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lIoCLSVeeMm4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def d_save_pic(i):\n",
        "  Generator.eval()\n",
        "  im = real = validate_set[i]\n",
        "  gray_scale = im = (0.2989*im[0,:,:] + 0.5870*im[1,:,:] + 0.1140*im[2,:,:])\n",
        "  im = im.unsqueeze(0).unsqueeze(0).cuda()\n",
        "  \n",
        "  output = Generator(im)\n",
        "  p = output[0].detach().cpu()\n",
        "  p = p.clamp(0.0, 1.0)\n",
        "  \n",
        "  real_img = transforms.ToPILImage()(real)\n",
        "  real_img = real_img.save(\"real\"+str(i)+\".jpg\")\n",
        "\n",
        "  PIL_img = transforms.ToPILImage()(p)\n",
        "  PIL_img = PIL_img.save(\"res\" + str(i) + \".jpg\")\n",
        "\n",
        "  bw_img = transforms.ToPILImage()(gray_scale)\n",
        "  bw_img = bw_img.save(\"bw\" + str(i) + \".jpg\")\n",
        "\n",
        "  Generator.train()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6S42-L9dei-s",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "for i in range(12):\n",
        "    d_save_pic(i)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}