clean_DCGANS.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "clean_DCGANS.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyPeH0fErCCep14jVGRaacVM",
      "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/tejasbana/DCGANS/blob/main/clean_DCGANS.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uBOoC_RWYRFS"
      },
      "source": [
        "import torch \n",
        "import torchvision\n",
        "import torchvision.transforms as T\n",
        "import matplotlib.pyplot as plt\n",
        "import os\n",
        "from torchvision.datasets import ImageFolder\n",
        "from torch.utils.data import DataLoader\n",
        "import torch.nn as nn\n",
        "from skimage.color import rgb2lab, lab2rgb, rgb2gray, gray2rgb\n",
        "from skimage.transform import resize\n",
        "from skimage.io import imsave\n",
        "import numpy as np\n",
        "import random\n",
        "import cv2\n",
        "%matplotlib inline"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hC-hztKsZQxw"
      },
      "source": [
        "!gdown --id 19N6B89I5ATZaP9s8zkS_IbP4LIm_V5Hu\n",
        "!unzip ./dataset.zip\n",
        "\n",
        "!mkdir images\n",
        "!mv /content/dataset/Test /content/images"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TA9vYFDfZYK_"
      },
      "source": [
        "image_size = 128\n",
        "batch_size = 16\n",
        "\n",
        "\n",
        "data_dir = \"/content/dataset\"\n",
        "print(len(os.listdir(data_dir +\"/Train\")))\n",
        "print(data_dir +\"/Train\")\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t2ZOnCFpZc_Q"
      },
      "source": [
        "train_ds = ImageFolder(data_dir , transform=T.Compose([ T.Resize(image_size),\n",
        "                                                        T.ToTensor()          ]))\n",
        "\n",
        "train_loader = DataLoader(train_ds , batch_size , num_workers=3 , pin_memory=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gsN5XWcZbvW3"
      },
      "source": [
        "**RGB_TENSOR_BATCH   to    LAB_TENSOR_BACTH**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QVq439BiZgo1"
      },
      "source": [
        "def convert_RGB_bacth_to_lab_batch(real_images):\n",
        "  xb = to_cpu(real_images , 'cpu')\n",
        "  numpy_rgb_batch = xb.permute(0,2,3,1).numpy()                   #convert rgb image tensor of range[0,1) to numpy array\n",
        "  numpy_lab_batch = rgb2lab(numpy_rgb_batch)                      # numpy LAB images of range[-127,127] shape : (batch_size,128,128,3)\n",
        "  numpy_lab_batch[:,:, :, 0] *= 255 / 100                         #Transform the numpy lab images to images of range [0, 1] \n",
        "  numpy_lab_batch[:,:, :, 1] += 128\n",
        "  numpy_lab_batch[:,:, :, 2] += 128\n",
        "  numpy_lab_batch /= 255\n",
        "  torch_lab_batch  = torch.from_numpy(np.transpose(numpy_lab_batch, (0,3,1,2))).type(torch.float32)       # torch LAB image shape:(batch_size,3,128,128)\n",
        "  return torch_lab_batch.detach()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z3o4nk2dcHjM"
      },
      "source": [
        "# MOVE TO GPU OR CPU"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OwBVTUZncGNx"
      },
      "source": [
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "\n",
        "def to_cpu(data, device='cpu'):\n",
        "    \"\"\"Move tensor(s) to chosen device\"\"\"\n",
        "    if isinstance(data, (list,tuple)):\n",
        "        return [to_device(x, device) for x in data]\n",
        "    return data.to(device, non_blocking=True)\n",
        "\n",
        "def to_device(data, device):\n",
        "    \"\"\"Move tensor(s) to chosen device\"\"\"\n",
        "    if isinstance(data, (list,tuple)):\n",
        "        return [to_device(x, device) for x in data]\n",
        "    return data.to(device, non_blocking=True)\n",
        "\n",
        "class DeviceDataLoader():\n",
        "    \"\"\"Wrap a dataloader to move data to a device\"\"\"\n",
        "    def __init__(self, dl, device):\n",
        "        self.dl = dl\n",
        "        self.device = device\n",
        "        \n",
        "    def __iter__(self):\n",
        "        \"\"\"Yield a batch of data after moving it to device\"\"\"\n",
        "        for b in self.dl: \n",
        "          img , l = b\n",
        "          img = convert_RGB_bacth_to_lab_batch(img)\n",
        "          b = img , l\n",
        "          yield to_device(b, self.device)\n",
        "\n",
        "    def __len__(self):\n",
        "        \"\"\"Number of batches\"\"\"\n",
        "        return len(self.dl)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PJahlUAuAazu"
      },
      "source": [
        "# Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SrxVwCfZAdUB"
      },
      "source": [
        "class UNet(torch.nn.Module):\n",
        "\n",
        "  def unet_conv(self , ch_in , ch_out , is_leaky):\n",
        "    if is_leaky:\n",
        "      return nn.Sequential(\n",
        "          nn.Conv2d(ch_in , ch_out , 3 , padding=1),\n",
        "          nn.BatchNorm2d(ch_out),\n",
        "          nn.LeakyReLU(0.2),\n",
        "          nn.Conv2d(ch_out , ch_out , 3 , padding=1),\n",
        "          nn.BatchNorm2d(ch_out),\n",
        "          nn.LeakyReLU(0.2)\n",
        "      )\n",
        "    else:\n",
        "      return nn.Sequential(\n",
        "          nn.Conv2d(ch_in , ch_out , 3 , padding=1),\n",
        "          nn.BatchNorm2d(ch_out),\n",
        "          nn.ReLU(),\n",
        "          nn.Conv2d(ch_out , ch_out , 3 , padding=1),\n",
        "          nn.BatchNorm2d(ch_out),\n",
        "          nn.ReLU()\n",
        "      )\n",
        "    \n",
        "  def up(self,ch_in,ch_out):\n",
        "    return nn.Sequential(\n",
        "        nn.ConvTranspose2d(ch_in , ch_out , 3, 2 , 1 ,1),\n",
        "        nn.ReLU()\n",
        "        )\n",
        "  \n",
        "  def __init__(self, is_leaky):\n",
        "    super(UNet,self).__init__()\n",
        "\n",
        "    # First encoding layer\n",
        "    self.conv1 = self.unet_conv(1,64, is_leaky)                     # IN : 128 x 128  , OUT : 128 x 128\n",
        "    # Second encoding layer\n",
        "    self.conv2 = self.unet_conv(64,128 , is_leaky)                  # IN : 128 x 128  , OUT : 64 x 64\n",
        "    # Third encoding layer\n",
        "    self.conv3 = self.unet_conv(128,256 , is_leaky)                 # IN : 64 x 64  , OUT : 32 x 32\n",
        "    # Forth encoding layer\n",
        "    self.conv4 = self.unet_conv(256,512, is_leaky)                  # IN : 32 x 32  , OUT : 16 x 16\n",
        "    # Fifth encoding layer\n",
        "    self.conv5 = self.unet_conv(512,1024, is_leaky)                 # IN : 16 x 16  , OUT : 8 x 8\n",
        "    # sixth enconding layer\n",
        "    self.conv6 = self.unet_conv(1024,2048 , is_leaky)               # IN : 8 x 8  , OUT : 4 x 4\n",
        "    # Seventh encoding layer\n",
        "    self.conv7 = self.unet_conv(2048,1024 , is_leaky)               # IN : 4 x 4  , OUT : 2 x 2\n",
        "\n",
        "\n",
        "    #Pooling layer\n",
        "    self.pool = nn.MaxPool2d(2)\n",
        "\n",
        "\n",
        "    # First Upsampling layer\n",
        "    self.up1 = self.up(1024,2048)                                   # IN : 2 x 2  , OUT : 4 x 4\n",
        "    # Second Upsampling layer\n",
        "    self.up2 = self.up(2048,1024)                                   # IN : 4 x 4  , OUT : 8 x 8\n",
        "    # Third Upsampling layer\n",
        "    self.up3 = self.up(1024,512)                                    # IN : 8 x 8  , OUT : 16 x 16\n",
        "    # Fourth Upsampling layer\n",
        "    self.up4 = self.up(512,256)                                     # IN : 16 x 16  , OUT : 32 x 32\n",
        "    # Fifth Upsampling layer\n",
        "    self.up5 = self.up(256,128)                                     # IN : 32 x 32  , OUT : 64 x 64\n",
        "    # Sixth Upsampling layer\n",
        "    self.up6 = self.up(128,64)                                      # IN : 64 x 64  , OUT : 128 x 128\n",
        "\n",
        "\n",
        "    # First Decoding layer\n",
        "    self.conv8  = self.unet_conv(4096 ,2048, False)\n",
        "    # Second Decoding layer\n",
        "    self.conv9  = self.unet_conv(2048, 1024, False)\n",
        "    # Third Decoding layer\n",
        "    self.conv10 = self.unet_conv(1024, 512 , False)\n",
        "    # Fourth Decoding layer\n",
        "    self.conv11 = self.unet_conv(512,  256 , False)\n",
        "    # Fifth Decoding layer \n",
        "    self.conv12 = self.unet_conv(256,  128 , False)\n",
        "    # Sixth Decoding layer\n",
        "    self.conv13 = self.unet_conv(128,  64  , False)\n",
        "\n",
        "\n",
        "    #Last layer\n",
        "    self.conv14 = nn.Conv2d(64,2,1)                   #IN_channel : 64 , OUT: 2 , Kernel_size = 1\n",
        "\n",
        "  \n",
        "  def forward(self, x):\n",
        "\n",
        "    #Encoding Path\n",
        "    x1 = self.conv1(x)\n",
        "    x2 = self.conv2(self.pool(x1))\n",
        "    x3 = self.conv3(self.pool(x2))\n",
        "    x4 = self.conv4(self.pool(x3))\n",
        "    x5 = self.conv5(self.pool(x4))\n",
        "    x6 = self.conv6(self.pool(x5))\n",
        "    x7 = self.conv7(self.pool(x6))\n",
        "\n",
        "    #Decoding Path\n",
        "    x  = self.conv8( torch.cat( ( x6 ,  self.up1(x7) ),1 ) )\n",
        "    x  = self.conv9( torch.cat( ( x5 ,  self.up2(x) ), 1 ) )\n",
        "    x  = self.conv10( torch.cat(( x4 , self.up3(x) ), 1 ) )\n",
        "    x  = self.conv11( torch.cat(( x3 , self.up4(x) ), 1 ) )\n",
        "    x  = self.conv12( torch.cat(( x2 , self.up5(x) ), 1 ) )\n",
        "    x  = self.conv13( torch.cat(( x1 , self.up6(x) ), 1 ) )\n",
        "\n",
        "    x = self.conv14(x)\n",
        "    m = nn.Tanh()\n",
        "    x = m(x)\n",
        "\n",
        "    return x\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UB07wJjEQLhY"
      },
      "source": [
        "# discriminator"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qdkL39afQxmx"
      },
      "source": [
        "class DNet(torch.nn.Module):\n",
        "  def unet_conv(self, ch_in , ch_out):\n",
        "    return nn.Sequential(\n",
        "        nn.Conv2d(ch_in , ch_out , 3, padding=1),\n",
        "        nn.BatchNorm2d(ch_out),\n",
        "        nn.LeakyReLU(0.2),\n",
        "        nn.Conv2d(ch_out , ch_out , 3, padding=1),\n",
        "        nn.BatchNorm2d(ch_out),\n",
        "        nn.LeakyReLU(0.2)\n",
        "    )\n",
        "\n",
        "  def __init__(self):\n",
        "    super(DNet,self).__init__()\n",
        "\n",
        "    # First layer\n",
        "    self.conv1 = self.unet_conv(3,64)\n",
        "    # Second layer\n",
        "    self.conv2 = self.unet_conv(64,64)\n",
        "    # Third layer \n",
        "    self.conv3 = self.unet_conv(64,128)\n",
        "    # Fourth layer\n",
        "    self.conv4 = self.unet_conv(128,128)\n",
        "    # Fifth layer\n",
        "    self.conv5 = self.unet_conv(128,256)\n",
        "    # Sixth layer\n",
        "    self.conv6 = self.unet_conv(256,512)\n",
        "    # Seventh layer\n",
        "    self.conv7 = self.unet_conv(512,1024)\n",
        "\n",
        "    #Pooling layer\n",
        "    self.pool = nn.MaxPool2d(2)\n",
        "\n",
        "    #Last layer\n",
        "    self.linear = nn.Linear(2*2*1024 , 1)\n",
        "\n",
        "  def forward(self,x):\n",
        "    x1 = self.conv1(x)\n",
        "    x2 = self.conv2(self.pool(x1))\n",
        "    x3 = self.conv3(self.pool(x2))\n",
        "    x4 = self.conv4(self.pool(x3))\n",
        "    x5 = self.conv5(self.pool(x4))\n",
        "    x6 = self.conv6(self.pool(x5))\n",
        "    x7 = self.conv7(self.pool(x6))\n",
        "\n",
        "    x8 = x5.reshape(-1,2*2*1024)\n",
        "    m = nn.Sigmoid()\n",
        "    x = m(self.linear(x8))\n",
        "\n",
        "    return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t_15a5OHc-FI"
      },
      "source": [
        "# Move to GPU"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sqUxXczbUFtg"
      },
      "source": [
        "discriminator = DNet()\n",
        "generator = UNet(True)\n",
        "generator.cuda()\n",
        "discriminator.cuda()\n",
        "\n",
        "train_loader = DeviceDataLoader(train_loader , device)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7w5Ltw1AdCzl"
      },
      "source": [
        "# Test"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PfiITdMEZGZw"
      },
      "source": [
        "for batch in train_loader:\n",
        "  img , l = batch\n",
        "  #lab_batch = convert_RGB_bacth_to_lab_batch(img) \n",
        "  #gray_batch = lab_batch[:,0,:,:]\n",
        "  #gray_batch = to_device(gray_batch , device)\n",
        "  gray_batch = img[:,0,:,:]\n",
        "  gray_batch = gray_batch.unsqueeze(1)\n",
        "  print(gray_batch.shape)\n",
        "  out = unet(gray_batch)\n",
        "  print(out.shape)\n",
        "  out = discriminator(torch.cat([gray_batch , out],1))\n",
        "  print(out.shape)\n",
        "  break;"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7Hc6cixMdSdB"
      },
      "source": [
        "# Fit"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xzHGNFuFdWJq"
      },
      "source": [
        "d_optimizer = torch.optim.Adam(discriminator.parameters() , betas=(0.5,0.999) , lr = 0.0002)\n",
        "g_optimizer = torch.optim.Adam(generator.parameters() , betas=(0.5,0.999) , lr = 0.0002)\n",
        "\n",
        "d_criterion = nn.BCELoss()\n",
        "g_criterion_1 = nn.BCELoss()\n",
        "g_criterion_2 = nn.L1Loss()\n",
        "\n",
        "def train(epochs):\n",
        "  losses_g = []\n",
        "  losses_d = []\n",
        "  g_lambda = 100\n",
        "  smooth = 0.1\n",
        "  \n",
        "\n",
        "  for epoch in range(epochs):\n",
        "    d_running_loss = 0.0\n",
        "    g_running_loss = 0.0\n",
        "\n",
        "    for lab_batch , _ in train_loader:\n",
        "      # split the lab color space images into luminescence and chrominance channels.\n",
        "      l_images = lab_batch[:,0,:,:]\n",
        "      c_images = lab_batch[:,1:,:,:]\n",
        "      # shift the source and target images into the range [-1, 1].\n",
        "      mean = torch.Tensor([0.5])\n",
        "      l_images = l_images - mean.expand_as(l_images).cuda()\n",
        "      l_images = 2*l_images\n",
        "      l_images = l_images.unsqueeze(1)\n",
        "\n",
        "      c_images = c_images - mean.expand_as(c_images).cuda()\n",
        "      c_images = 2*c_images\n",
        "\n",
        "      batch_size = l_images.shape[0]\n",
        "\n",
        "      # fake images are generated by passing them through the generator.\n",
        "      fake_images = generator(l_images)\n",
        "\n",
        "      # Train the discriminator. The loss would be the sum of the losses over\n",
        "\t    # the source and fake images, with greyscale images as the condition.\n",
        "      d_optimizer.zero_grad()\n",
        "      d_loss = 0\n",
        "      preds = discriminator(cat([l_images, c_images] , 1))\n",
        "      d_real_loss = d_criterion(preds.squeeze(1) , ((1 - smooth) * torch.ones(batch_size)).cuda() )\n",
        "\n",
        "      preds = discriminator(cat([l_images, fake_images] , 1))\n",
        "      d_fake_loss = d_criterion(preds.squeeze(1) , (torch.zeros(batch_size)).cuda())\n",
        "\n",
        "      d_loss = d_real_loss + d_fake_loss\n",
        "      d_loss.backward(retain_graph=True)\n",
        "      d_optimizer.step()\n",
        "\n",
        "      # Train the generator. The loss would be the sum of the adversarial loss\n",
        "\t    # due to the GAN and L1 distance loss between the fake and target images.\n",
        "\n",
        "      g_optimizer.zero_grad()\n",
        "      g_loss = 0\n",
        "      fake_preds = discriminator(cat([l_images , fake_images] , 1))\n",
        "      g_fake_loss = g_criterion_1(fake_preds.squeeze(1) , (torch.ones(batch_size)).cuda())         # generator loss 1\n",
        "\n",
        "      g_image_distance_loss = g_lambda * g_criterion_2(fake_images , c_images)         # generator loss 2\n",
        "\n",
        "      g_loss = g_fake_loss + g_image_distance_loss\n",
        "      g_loss.backward(retain_graph=True)\n",
        "      g_optimizer.step()\n",
        "\n",
        "      # print statistics on pre-defined intervals.\n",
        "      d_running_loss += d_loss.detach()\n",
        "      g_running_loss += g_fake_loss.detach()\n",
        "\n",
        "    print('Epoch : {} , g_epoch_loss : {:.4f} , d_epoch_loss : {:.4f}'.format(epoch,g_running_loss,d_running_loss))\n",
        "    losses_g.append(d_running_loss)\n",
        "    losses_d.append(g_running_loss)\n",
        "  return losses_g , losses_d\n",
        "\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S7ftT9grxc7p"
      },
      "source": [
        "discriminator.load_state_dict(torch.load('/content/drive/MyDrive/weights/discriminator.ckpt'))\n",
        "generator.load_state_dict(torch.load('/content/drive/MyDrive/weights/generator.ckpt'))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sTnEtUlsw4oD"
      },
      "source": [
        "history = train(5)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6f-T0ixHw66r"
      },
      "source": [
        "torch.save(generator.state_dict() , 'generator.ckpt')\n",
        "torch.save(discriminator.state_dict() , 'discriminator.ckpt')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hWB42gNdnGOn"
      },
      "source": [
        "def denorm(l_image ,fake_image):\n",
        "  l_image = l_image * 100\n",
        "  mean = torch.Tensor([0.5])  \n",
        "  fake_image = fake_image + mean.expand_as(fake_image)\n",
        "  fake_image *= 255\n",
        "  fake_image -= 128\n",
        "  torch_image = cat([l_image ,fake_image] , 0).detach()\n",
        "  numpy_image = torch_image.permute(1,2,0).numpy()\n",
        "  rgb_image = lab2rgb(numpy_image)\n",
        "  plt.imshow(rgb_image)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jdidggcrKede"
      },
      "source": [
        "rgb_image , l = train_ds[2222]                                      #take single image from dataset\n",
        "numpy_rgb_image = rgb_image.permute(1,2,0).numpy()                #convert rgb image tensor of range[0,1) to numpy array \n",
        "numpy_lab_image = rgb2lab(numpy_rgb_image)                        # numpy LAB images of range[-127,127] shape : (128,128,3)\n",
        "#numpy_lab_image = cv2.cvtColor(numpy_rgb_image, cv2.COLOR_RGB2LAB)\n",
        "\n",
        "numpy_lab_image[:, :, 0] *= 255 / 100\n",
        "numpy_lab_image[:, :, 1] += 128\n",
        "numpy_lab_image[:, :, 2] += 128\n",
        "numpy_lab_image /= 255                                            #Transform the numpy lab images to images of range [0, 1] \n",
        "\n",
        "torch_lab_image = torch.from_numpy(np.transpose(numpy_lab_image, (2, 0, 1)))       # torch LAB image\n",
        "gray_image = torch_lab_image[0,:,:].type(torch.float32)                             #shape [2 dim]\n",
        "mean = torch.Tensor([0.5])  \n",
        "gray_image = gray_image - mean.expand_as(gray_image)              \n",
        "\n",
        "l_image = gray_image.type(torch.float32)\n",
        "gray_image = gray_image.unsqueeze(0)\n",
        "gray_image = gray_image.unsqueeze(0).type(torch.float32)\n",
        "\n",
        "fake_image = generator(gray_image.cuda())\n",
        "fake_image = to_cpu(fake_image)\n",
        "fake_image = fake_image.squeeze(0)                      # remove batch channel\n",
        "print(fake_image.shape)\n",
        "l_image = l_image.unsqueeze(0)\n",
        "denorm(l_image ,fake_image)\n",
        "\n",
        "print(l_image.shape)\n",
        "pred = cat([l_image ,fake_image] , 0).detach()\n",
        "pred = pred.permute(1,2,0).numpy()\n",
        "#print(pred)\n",
        "#plt.imshow(pred)\n",
        "#plt.imshow(numpy_lab_image)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}