RCAN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RCAN\n",
    "\n",
    "This notebook implements RCAN 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "RCAN model\n",
    "\"\"\"\n",
    "class ChannelAttention(nn.Module):\n",
    "    def __init__(self, num_features, reduction):\n",
    "        super(ChannelAttention, self).__init__()\n",
    "        self.module = nn.Sequential(\n",
    "            nn.AdaptiveAvgPool2d(1),\n",
    "            nn.Conv2d(num_features, num_features // reduction, kernel_size=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(num_features // reduction, num_features, kernel_size=1),\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x * self.module(x)\n",
    "\n",
    "class RCAB(nn.Module):\n",
    "    def __init__(self, num_features, reduction):\n",
    "        super(RCAB, self).__init__()\n",
    "        self.module = nn.Sequential(\n",
    "            nn.Conv2d(num_features, num_features, kernel_size=3, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(num_features, num_features, kernel_size=3, padding=1),\n",
    "            ChannelAttention(num_features, reduction)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x + self.module(x)\n",
    "\n",
    "class RG(nn.Module):\n",
    "    def __init__(self, num_features, num_rcab, reduction):\n",
    "        super(RG, self).__init__()\n",
    "        self.module = [RCAB(num_features, reduction) for _ in range(num_rcab)]\n",
    "        self.module.append(nn.Conv2d(num_features, num_features, kernel_size=3, padding=1))\n",
    "        self.module = nn.Sequential(*self.module)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x + self.module(x)\n",
    "\n",
    "class RCAN(nn.Module):\n",
    "    def __init__(self, args):\n",
    "        super(RCAN, self).__init__()\n",
    "        scale = args.scale\n",
    "        num_features = args.num_features\n",
    "        num_rg = args.num_rg\n",
    "        num_rcab = args.num_rcab\n",
    "        reduction = args.reduction\n",
    "\n",
    "        self.sf = nn.Conv2d(3, num_features, kernel_size=3, padding=1)\n",
    "        # RIR layers\n",
    "        self.rgs = nn.Sequential(*[RG(num_features, num_rcab, reduction) for _ in range(num_rg)])\n",
    "        self.conv1 = nn.Conv2d(num_features, num_features, kernel_size=3, padding=1)\n",
    "        self.upscale = nn.Sequential(\n",
    "            nn.Conv2d(num_features, num_features * (scale ** 2), kernel_size=3, padding=1),\n",
    "            nn.PixelShuffle(scale)\n",
    "        )\n",
    "        self.conv2 = nn.Conv2d(num_features, 3, kernel_size=3, padding=1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.sf(x)\n",
    "        residual = x\n",
    "        x = self.rgs(x)\n",
    "        x = self.conv1(x)\n",
    "        x += residual\n",
    "        x = self.upscale(x)\n",
    "        x = self.conv2(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Setup the dataset\n",
    "\"\"\"\n",
    "def create_data(path, output):\n",
    "\n",
    "    h5_file = h5py.File(os.path.join(path, output), 'w')\n",
    "\n",
    "    hr_image_path = os.path.join(path, 'images_stage3/*.png')\n",
    "    lr_image_path = os.path.join(path, 'images_stage5/*.png')\n",
    "\n",
    "    hr_image_list = glob.glob(hr_image_path)\n",
    "    lr_image_list = glob.glob(lr_image_path)\n",
    "\n",
    "    hr_imgs = []\n",
    "    lr_imgs = []\n",
    "\n",
    "    for i in range(len(hr_image_list)):\n",
    "\n",
    "        # open image\n",
    "        hr = Image.open(hr_image_list[i]).convert('RGB')\n",
    "        lr = Image.open(lr_image_list[i]).convert('RGB')\n",
    "\n",
    "        # convert data type\n",
    "        hr = np.array(hr).astype(np.float32)\n",
    "        lr = np.array(lr).astype(np.float32)\n",
    "\n",
    "        # transpose and normalize\n",
    "        hr = np.transpose(hr, axes=[2, 0, 1])\n",
    "        lr = np.transpose(lr, axes=[2, 0, 1])\n",
    "\n",
    "        hr /= 255.0\n",
    "        lr /= 255.0\n",
    "\n",
    "        hr_imgs.append(hr)\n",
    "        lr_imgs.append(lr)\n",
    "\n",
    "    hr_imgs = np.array(hr_imgs)\n",
    "    lr_imgs = np.array(lr_imgs)\n",
    "\n",
    "    h5_file.create_dataset('lr', np.shape(lr_imgs), h5py.h5t.IEEE_F32LE, data=lr_imgs)\n",
    "    h5_file.create_dataset('hr', np.shape(hr_imgs), h5py.h5t.IEEE_F32LE, data=hr_imgs)\n",
    "\n",
    "    h5_file.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "create_data('dataset/train', 'train_full_s.h5')\n",
    "create_data('dataset/val', 'val_full_s.h5')"
   ]
  },
  {
   "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'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Model Setup\n",
    "\"\"\"\n",
    "torch.manual_seed(123)\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "ParamStrct = namedtuple('ParamStrct', 'scale num_features num_rg num_rcab reduction')\n",
    "param = ParamStrct(4, 64, 6, 16, 16)\n",
    "batch_size = 1\n",
    "\n",
    "model = RCAN(param).to(device)\n",
    "criterion = nn.L1Loss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Setup data loader\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)"
   ]
  },
  {
   "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": [
    "\"\"\"\n",
    "Train and val the model\n",
    "\"\"\"\n",
    "best_weights = copy.deepcopy(model.state_dict())\n",
    "best_epoch = 0\n",
    "best_psnr = 0.0\n",
    "\n",
    "num_epoch = 20\n",
    "\n",
    "for epoch in range(num_epoch):\n",
    "    model.train()\n",
    "    epoch_losses = AverageMeter()\n",
    "\n",
    "    with tqdm(total=(len(train_dataset) - len(train_dataset) % batch_size)) as t:\n",
    "        t.set_description('epoch: {}/{}'.format(epoch, num_epoch - 1))\n",
    "        \n",
    "        # training\n",
    "        for data in train_dataloader:\n",
    "            inputs, labels = data\n",
    "\n",
    "            inputs = inputs.to(device, dtype=torch.float)\n",
    "            labels = labels.to(device, dtype=torch.float)\n",
    "\n",
    "            preds = model(inputs)\n",
    "\n",
    "            loss = criterion(preds, labels)\n",
    "\n",
    "            epoch_losses.update(loss.item(), len(inputs))\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            t.set_postfix(loss='{:.6f}'.format(epoch_losses.avg))\n",
    "            t.update(len(inputs))\n",
    "\n",
    "        torch.save(model.state_dict(), os.path.join('weight_rcan', 'epoch_{}.pth'.format(epoch)))\n",
    "\n",
    "        # validation\n",
    "        model.eval()\n",
    "        epoch_psnr = AverageMeter()\n",
    "\n",
    "        for data in eval_dataloader:\n",
    "            inputs, labels = data\n",
    "\n",
    "            inputs = inputs.to(device, dtype=torch.float)\n",
    "            labels = labels.to(device, dtype=torch.float)\n",
    "\n",
    "            with torch.no_grad():\n",
    "                preds = model(inputs).clamp(0.0, 1.0)\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(model.state_dict())\n",
    "\n",
    "    print('best epoch: {}, psnr: {:.2f}'.format(best_epoch, best_psnr))\n",
    "    torch.save(best_weights, os.path.join('weight_rcan', 'best.pth'))"
   ]
  },
  {
   "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 = RCAN(param).to(device)\n",
    "state_dict = model.state_dict()\n",
    "for n, p in torch.load('weight_rcan/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()"
   ]
  },
  {
   "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_rcan_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
}