SRGAN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SRGAN\n",
    "\n",
    "This notebook implements SRGAN model along with training and test data creation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Import Library\n",
    "\"\"\"\n",
    "from torch import nn\n",
    "import h5py\n",
    "import numpy as np\n",
    "import glob\n",
    "import os\n",
    "from PIL import Image\n",
    "from torch.utils.data import Dataset\n",
    "import torch.optim as optim\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "from collections import namedtuple\n",
    "import copy\n",
    "import math\n",
    "from torch.autograd import Variable\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "SRGAN model\n",
    "\"\"\"\n",
    "class Generator(nn.Module):\n",
    "    def __init__(self, scale_factor):\n",
    "        upsample_block_num = int(math.log(scale_factor, 2))\n",
    "\n",
    "        super(Generator, self).__init__()\n",
    "        self.block1 = nn.Sequential(\n",
    "            nn.Conv2d(3, 32, kernel_size=9, padding=4),\n",
    "            nn.PReLU()\n",
    "        )\n",
    "        self.block2 = ResidualBlock(32)\n",
    "        self.block3 = ResidualBlock(32)\n",
    "        self.block4 = ResidualBlock(32)\n",
    "        self.block5 = ResidualBlock(32)\n",
    "        self.block6 = ResidualBlock(32)\n",
    "        self.block7 = nn.Sequential(\n",
    "            nn.Conv2d(32, 32, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(32)\n",
    "        )\n",
    "        block8 = [UpsampleBLock(32, 2) for _ in range(upsample_block_num)]\n",
    "        block8.append(nn.Conv2d(32, 3, kernel_size=9, padding=4))\n",
    "        self.block8 = nn.Sequential(*block8)\n",
    "\n",
    "    def forward(self, x):\n",
    "        block1 = self.block1(x)\n",
    "        block2 = self.block2(block1)\n",
    "        block3 = self.block3(block2)\n",
    "        block4 = self.block4(block3)\n",
    "        block5 = self.block5(block4)\n",
    "        block6 = self.block6(block5)\n",
    "        block7 = self.block7(block6)\n",
    "        block8 = self.block8(block1 + block7)\n",
    "\n",
    "        return (torch.tanh(block8) + 1) / 2\n",
    "\n",
    "\n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Discriminator, self).__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Conv2d(3, 32, kernel_size=3, padding=1),\n",
    "            nn.LeakyReLU(0.2),\n",
    "\n",
    "            nn.Conv2d(32, 32, kernel_size=3, stride=2, padding=1),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.LeakyReLU(0.2),\n",
    "\n",
    "            nn.Conv2d(32, 64, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.LeakyReLU(0.2),\n",
    "\n",
    "            nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.LeakyReLU(0.2),\n",
    "\n",
    "            nn.Conv2d(64, 128, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.LeakyReLU(0.2),\n",
    "\n",
    "            nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=1),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.LeakyReLU(0.2),\n",
    "\n",
    "            nn.Conv2d(128, 256, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.LeakyReLU(0.2),\n",
    "\n",
    "            nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.LeakyReLU(0.2),\n",
    "\n",
    "            nn.AdaptiveAvgPool2d(1),\n",
    "            nn.Conv2d(512, 512, kernel_size=1),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            nn.Conv2d(512, 1, kernel_size=1)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.size(0)\n",
    "        return torch.sigmoid(self.net(x).view(batch_size))\n",
    "\n",
    "\n",
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, channels):\n",
    "        super(ResidualBlock, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)\n",
    "        self.bn1 = nn.BatchNorm2d(channels)\n",
    "        self.prelu = nn.PReLU()\n",
    "        self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)\n",
    "        self.bn2 = nn.BatchNorm2d(channels)\n",
    "\n",
    "    def forward(self, x):\n",
    "        residual = self.conv1(x)\n",
    "        residual = self.bn1(residual)\n",
    "        residual = self.prelu(residual)\n",
    "        residual = self.conv2(residual)\n",
    "        residual = self.bn2(residual)\n",
    "\n",
    "        return x + residual\n",
    "\n",
    "\n",
    "class UpsampleBLock(nn.Module):\n",
    "    def __init__(self, in_channels, up_scale):\n",
    "        super(UpsampleBLock, self).__init__()\n",
    "        self.conv = nn.Conv2d(in_channels, in_channels * up_scale ** 2, kernel_size=3, padding=1)\n",
    "        self.pixel_shuffle = nn.PixelShuffle(up_scale)\n",
    "        self.prelu = nn.PReLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv(x)\n",
    "        x = self.pixel_shuffle(x)\n",
    "        x = self.prelu(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Dataset feeding\n",
    "\"\"\"\n",
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, h5_file):\n",
    "        super(CustomDataset, self).__init__()\n",
    "        self.h5_file = h5_file\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        with h5py.File(self.h5_file, 'r') as f:\n",
    "            return f['lr'][idx], f['hr'][idx]\n",
    "\n",
    "    def __len__(self):\n",
    "        with h5py.File(self.h5_file, 'r') as f:\n",
    "            return len(f['lr'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Loss Functions\n",
    "\"\"\"\n",
    "from torchvision.models.vgg import vgg16\n",
    "\n",
    "# TV loss is optional but implemented in paper\n",
    "class TVLoss(nn.Module):\n",
    "    def __init__(self, tv_loss_weight=1):\n",
    "        super(TVLoss, self).__init__()\n",
    "        self.tv_loss_weight = tv_loss_weight\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.size()[0]\n",
    "        h_x = x.size()[2]\n",
    "        w_x = x.size()[3]\n",
    "        count_h = self.tensor_size(x[:, :, 1:, :])\n",
    "        count_w = self.tensor_size(x[:, :, :, 1:])\n",
    "        h_tv = torch.pow((x[:, :, 1:, :] - x[:, :, :h_x - 1, :]), 2).sum()\n",
    "        w_tv = torch.pow((x[:, :, :, 1:] - x[:, :, :, :w_x - 1]), 2).sum()\n",
    "        return self.tv_loss_weight * 2 * (h_tv / count_h + w_tv / count_w) / batch_size\n",
    "\n",
    "    @staticmethod\n",
    "    def tensor_size(t):\n",
    "        return t.size()[1] * t.size()[2] * t.size()[3]\n",
    "\n",
    "class GeneratorLoss(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(GeneratorLoss, self).__init__()\n",
    "        # use VGG16 for loss calculation\n",
    "        vgg = vgg16(pretrained=True, progress=False)\n",
    "        loss_network = nn.Sequential(*list(vgg.features)[:31]).eval()\n",
    "        for param in loss_network.parameters():\n",
    "            param.requires_grad = False\n",
    "        self.loss_network = loss_network\n",
    "        self.mse_loss = nn.MSELoss()\n",
    "        self.tv_loss = TVLoss()\n",
    "\n",
    "    def forward(self, out_labels, out_images, target_images):\n",
    "        # Adversarial Loss\n",
    "        adversarial_loss = torch.mean(1 - out_labels)\n",
    "        # Perception Loss\n",
    "        perception_loss = self.mse_loss(self.loss_network(out_images), self.loss_network(target_images))\n",
    "        # Image Loss\n",
    "        image_loss = self.mse_loss(out_images, target_images)\n",
    "        # TV Loss\n",
    "        tv_loss = self.tv_loss(out_images)\n",
    "        return image_loss + 0.001 * adversarial_loss + 0.006 * perception_loss + 2e-8 * tv_loss\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Setup data loader\n",
    "\"\"\"\n",
    "batch_size = 1\n",
    "\n",
    "train_dataset = CustomDataset('dataset/train/train_full_s.h5')\n",
    "train_dataloader = DataLoader(dataset=train_dataset,\n",
    "                                  batch_size=batch_size,\n",
    "                                  shuffle=True,\n",
    "                                  num_workers=0,\n",
    "                                  pin_memory=True,\n",
    "                                  drop_last=True)\n",
    "eval_dataset = CustomDataset('dataset/val/val_full_s.h5')\n",
    "eval_dataloader = DataLoader(dataset=eval_dataset, batch_size=batch_size)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Setup network parameter\n",
    "\"\"\"\n",
    "upscale_factor = 4\n",
    "num_epoch = 20\n",
    "\n",
    "torch.manual_seed(123)\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Setup network\n",
    "\"\"\"\n",
    "netG = Generator(upscale_factor)\n",
    "netD = Discriminator()\n",
    "generator_criterion = GeneratorLoss()\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    netG.to(device)\n",
    "    netD.to(device)\n",
    "    generator_criterion.to(device)\n",
    "\n",
    "optimizerG = optim.Adam(netG.parameters())\n",
    "optimizerD = optim.Adam(netD.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Util function to measure error\n",
    "\"\"\"\n",
    "class AverageMeter(object):\n",
    "    def __init__(self):\n",
    "        self.reset()\n",
    "\n",
    "    def reset(self):\n",
    "        self.val = 0\n",
    "        self.avg = 0\n",
    "        self.sum = 0\n",
    "        self.count = 0\n",
    "\n",
    "    def update(self, val, n=1):\n",
    "        self.val = val\n",
    "        self.sum += val * n\n",
    "        self.count += n\n",
    "        self.avg = self.sum / self.count\n",
    "\n",
    "\"\"\"\n",
    "Calculate PSNR\n",
    "\"\"\"\n",
    "def calc_psnr(img1, img2):\n",
    "    return 10. * torch.log10(1. / torch.mean((img1 - img2) ** 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "results = {'d_loss': [], 'g_loss': [], 'd_score': [], 'g_score': [], 'psnr': []}\n",
    "best_weights = copy.deepcopy(netG.state_dict())\n",
    "best_epoch = 0\n",
    "best_psnr = 0.0\n",
    "\n",
    "for epoch in range(1, num_epoch + 1):\n",
    "\n",
    "    epoch_losses = AverageMeter()\n",
    "    netG.train()\n",
    "    netD.train()\n",
    "\n",
    "    with tqdm(total=(len(train_dataset) - len(train_dataset) % 1)) as t:\n",
    "        t.set_description('epoch: {}/{}'.format(epoch, 30 - 1))\n",
    "\n",
    "        running_results = {'batch_sizes': 0, 'd_loss': 0, 'g_loss': 0, 'd_score': 0, 'g_score': 0}\n",
    "\n",
    "        # training\n",
    "        netG.train()\n",
    "        netD.train()\n",
    "\n",
    "        for data in train_dataloader:\n",
    "            inputs, labels = data\n",
    "\n",
    "            g_update_first = True\n",
    "            batch_size = inputs.size(0)\n",
    "            running_results['batch_sizes'] += batch_size\n",
    "\n",
    "            # Update D network\n",
    "            real_img = Variable(labels).to(device, dtype=torch.float)\n",
    "            z = Variable(inputs).to(device, dtype=torch.float)\n",
    "\n",
    "            fake_img = netG(z)\n",
    "\n",
    "            netD.zero_grad()\n",
    "            real_out = netD(real_img).mean()\n",
    "            fake_out = netD(fake_img).mean()\n",
    "            d_loss = 1 - real_out + fake_out\n",
    "            d_loss.backward(retain_graph=True)\n",
    "\n",
    "            # Update G network\n",
    "            netG.zero_grad()\n",
    "            g_loss = generator_criterion(fake_out, fake_img, real_img)\n",
    "            g_loss.backward()\n",
    "\n",
    "            epoch_losses.update(g_loss.item(), len(inputs))\n",
    "\n",
    "            optimizerD.step()\n",
    "            optimizerG.step()\n",
    "\n",
    "            # Loss for current batch\n",
    "            running_results['g_loss'] += g_loss.item() * batch_size\n",
    "            running_results['d_loss'] += d_loss.item() * batch_size\n",
    "            running_results['d_score'] += real_out.item() * batch_size\n",
    "            running_results['g_score'] += fake_out.item() * batch_size\n",
    "\n",
    "            t.set_description(desc='[%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f' % (\n",
    "                epoch, num_epoch, running_results['d_loss'] / running_results['batch_sizes'],\n",
    "                running_results['g_loss'] / running_results['batch_sizes'],\n",
    "                running_results['d_score'] / running_results['batch_sizes'],\n",
    "                running_results['g_score'] / running_results['batch_sizes']))\n",
    "            t.update(len(inputs))\n",
    "\n",
    "        torch.save(netG.state_dict(), 'weight_srgan/netG_epoch_%d.pth' % epoch)\n",
    "        torch.save(netD.state_dict(), 'weight_srgan/netD_epoch_%d.pth' % epoch)\n",
    "\n",
    "        # validation\n",
    "        netG.eval()\n",
    "        epoch_psnr = AverageMeter()\n",
    "\n",
    "        with torch.no_grad():\n",
    "            val_images = []\n",
    "            for data in eval_dataloader:\n",
    "                inputs, labels = data\n",
    "                inputs = inputs.to(device, dtype=torch.float)\n",
    "                labels = labels.to(device, dtype=torch.float)\n",
    "\n",
    "                preds = netG(inputs)\n",
    "\n",
    "                epoch_psnr.update(calc_psnr(preds, labels), len(inputs))\n",
    "\n",
    "            print('eval psnr: {:.2f}'.format(epoch_psnr.avg))\n",
    "\n",
    "            if epoch_psnr.avg > best_psnr:\n",
    "                best_epoch = epoch\n",
    "                best_psnr = epoch_psnr.avg\n",
    "                best_weights = copy.deepcopy(netG.state_dict())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Evaluate the model with test set\n",
    "\"\"\"\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "model = Generator(upscale_factor).to(device)\n",
    "state_dict = model.state_dict()\n",
    "for n, p in torch.load('weight_srgan/best.pth', map_location=lambda storage, loc: storage).items():\n",
    "    if n in state_dict.keys():\n",
    "        state_dict[n].copy_(p)\n",
    "    else:\n",
    "        raise KeyError(n)\n",
    "\n",
    "model.eval()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from pytorch_ssim import pytorch_ssim\n",
    "\n",
    "lr_image_path = 'dataset/test/images_stage5/*.png'\n",
    "lr_image_list = glob.glob(lr_image_path)\n",
    "hr_image_path = 'dataset/test/images_stage3/*.png'\n",
    "hr_image_list = glob.glob(hr_image_path)\n",
    "\n",
    "psnr_total = 0\n",
    "ssim_total = 0\n",
    "\n",
    "for i, img in enumerate(lr_image_list):\n",
    "    image = Image.open(img).convert('RGB')\n",
    "    image = np.array(image).astype(np.float32)\n",
    "    image = np.transpose(image, axes=[2, 0, 1])\n",
    "    image /= 255.0\n",
    "\n",
    "    image = torch.from_numpy(image).to(device)\n",
    "    image = image.unsqueeze(0)\n",
    "\n",
    "    label = Image.open(hr_image_list[i]).convert('RGB')\n",
    "    label = np.array(label).astype(np.float32)\n",
    "    label = np.transpose(label, axes=[2, 0, 1])\n",
    "    label /= 255.0\n",
    "    label = torch.from_numpy(label).to(device)\n",
    "    label = label.unsqueeze(0)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        preds = model(image).clamp(0.0, 1.0)\n",
    "\n",
    "    psnr = calc_psnr(label, preds)\n",
    "    psnr_total += psnr\n",
    "    ssim = pytorch_ssim.ssim(label, preds)\n",
    "    ssim_total += ssim\n",
    "    print('PSNR: {:.2f}'.format(psnr))\n",
    "    print('SSIM: {:.2f}'.format(ssim))\n",
    "\n",
    "    output = preds.mul_(255.0).clamp_(0.0, 255.0).squeeze(0).permute(1, 2, 0).byte().cpu().numpy()\n",
    "    output = Image.fromarray(output)\n",
    "    output.save(f'result_srgan_new/img_{i}.png')\n",
    "\n",
    "psnr_total /= len(lr_image_list)\n",
    "ssim_total /= len(lr_image_list)\n",
    "print('PSNR_T: {:.4f}'.format(psnr_total))\n",
    "print('SSIM_T: {:.4f}'.format(ssim_total))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}