PlantSR.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2abe9809-4bd1-4de6-acff-a29cef5286e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn,optim\n",
    "from torch.backends import cudnn\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "from torch.utils.data import Dataset\n",
    "import os\n",
    "import numpy as np\n",
    "import cv2\n",
    "import torch\n",
    "from torchvision import transforms\n",
    "from torch.utils.data import Dataset\n",
    "import random\n",
    "from math import sqrt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5e74a55b-14b7-4c4d-aa48-783e1c3219d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SRDataset(Dataset):\n",
    "    def __init__(self, root, upscale_factor):\n",
    "        super(SRDataset, self).__init__()\n",
    "        self.hr_path = os.path.join(root, 'train_64')\n",
    "        self.upscale_factor = upscale_factor\n",
    "        self.hr_filenames = sorted(os.listdir(self.hr_path))\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        hr_image = cv2.imread(os.path.join(self.hr_path, self.hr_filenames[index]))\n",
    "        hr_image = cv2.cvtColor(hr_image, cv2.COLOR_BGR2RGB)\n",
    "        h, w, _ = hr_image.shape\n",
    "\n",
    "        ## make sure same demension\n",
    "        h -= h % self.upscale_factor\n",
    "        w -= w % self.upscale_factor\n",
    "        hr_image = hr_image[:h, :w]\n",
    "\n",
    "        lr_image = cv2.resize(hr_image, (int(w // self.upscale_factor),int(h // self.upscale_factor)), interpolation=cv2.INTER_LINEAR)\n",
    "\n",
    "\n",
    "        if random.random() > 0.5:  \n",
    "            lr_image = cv2.flip(lr_image, 1)\n",
    "            hr_image = cv2.flip(hr_image, 1)\n",
    "        \n",
    "        transform = transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])\n",
    "        ])\n",
    "        lr_image = transform(lr_image)\n",
    "        hr_image = transform(hr_image)\n",
    "\n",
    "\n",
    "\n",
    "        return lr_image, hr_image\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.hr_filenames)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03046720-0094-4703-8b65-0819b000f7e7",
   "metadata": {},
   "source": [
    "## train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e7ab9a65-b566-47b9-9251-ba2b9c57462d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "import torch.nn as nn\n",
    "import torch\n",
    "\n",
    "upscale= 2\n",
    "train_dataset = SRDataset(root='./data/PlantSR_dataset/', upscale_factor=upscale)\n",
    "train_loader = DataLoader(dataset=train_dataset, batch_size=32, shuffle=True)\n",
    "\n",
    "device = 'cuda:0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c592409e-cd55-44d2-8554-8dbbbf37d9ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "## upscale = 2/3/4\n",
    "from models.PlantSR import PlantSR\n",
    "\n",
    "outPath = \"outputs\"\n",
    "lr = 1e-4  \n",
    "\n",
    "device = torch.device('cuda:0')\n",
    "if upscale == 4:\n",
    "    model = PlantSR(scale=upscale,num_features=96,n_resgroups=16,n_resblocks=4,reduction=16)\n",
    "if upscale == 2:\n",
    "    model = PlantSR(scale=upscale,num_features=32,n_resgroups=16,n_resblocks=4,reduction=16)\n",
    "if upscale == 3:\n",
    "    model = PlantSR(scale=upscale,num_features=64,n_resgroups=16,n_resblocks=4,reduction=16)\n",
    "model.to(device)\n",
    "criterion = nn.L1Loss()\n",
    "optimizer = optim.Adam(model.parameters(),lr=lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d2020c06-acfd-41e0-9552-c87bf5da2e67",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## load the pretrained model (if have one)\n",
    "\n",
    "# model_path = 'ckpts/PlantSR_x2_best.pth'\n",
    "# model.load_state_dict(torch.load(model_path), strict=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dfd32df-a578-4c40-bb84-171355b77609",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/15], Batch [22934/38997], Loss: 0.0151"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "import sys\n",
    "\n",
    "\n",
    "start_epoch = 0\n",
    "num_epochs = 15\n",
    "\n",
    "for epoch in range(start_epoch,num_epochs):\n",
    "    model.train()\n",
    "    for batch_idx, (lr_images, hr_images) in enumerate(train_loader):\n",
    "        lr_images = lr_images.to(device)\n",
    "        hr_images = hr_images.to(device)\n",
    "\n",
    "        sr_images = model(lr_images.float())\n",
    "\n",
    "        loss = criterion(sr_images, hr_images)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if (batch_idx+1) % 1 == 0:\n",
    "            sys.stdout.write('\\rEpoch [{}/{}], Batch [{}/{}], Loss: {:.4f}'\n",
    "                             .format(epoch+1, num_epochs, batch_idx+1, len(train_loader), loss.item()))\n",
    "            sys.stdout.flush()\n",
    "        # if(batch_idx%10000) == 0:\n",
    "        #     torch.save(model.state_dict(), 'outputs/PlantSR_x2_{}_{}.pth'.format(batch_idx,epoch+1))\n",
    "\n",
    "    print(\"\\n\")\n",
    "    if (epoch+1) % 1 == 0:\n",
    "        torch.save(model.state_dict(), 'outputs/PlantSR_x2_{}.pth'.format(epoch+1))\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from models.PlantSR import PlantSR\n",
    "# from PLtest import PlantSR\n",
    "import torch\n",
    "\n",
    "upscale = 2\n",
    "device = 'cuda'\n",
    "# model_path = r'outputs/PlantSR_x4_20000_10.pth'\n",
    "model_path = r'outputs/PlantSR_x2_1.pth'\n",
    "device = torch.device('cuda:0')\n",
    "if upscale == 4:\n",
    "    model = PlantSR(scale=upscale,num_features=96,n_resgroups=16,n_resblocks=4,reduction=16)\n",
    "if upscale == 2:\n",
    "    model = PlantSR(scale=upscale,num_features=32,n_resgroups=16,n_resblocks=4,reduction=16)\n",
    "if upscale == 3:\n",
    "    model = PlantSR(scale=upscale,num_features=64,n_resgroups=16,n_resblocks=4,reduction=16)\n",
    "\n",
    "model.load_state_dict(torch.load(model_path), strict=True)\n",
    "model.eval()\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2 as cv2\n",
    "import numpy as np\n",
    "import torch.nn.functional as F\n",
    "from calulate_psnr_ssim import *\n",
    "import os\n",
    "\n",
    "test_psnr = 0\n",
    "test_ssim = 0\n",
    "image_count = 0\n",
    "\n",
    "test_path = \"./data/PlantSR_dataset/test\"\n",
    "\n",
    "for filename in os.listdir(test_path):\n",
    "    if filename.endswith((\".png\",\".jpg\")):\n",
    "        image_count+=1\n",
    "        print(image_count)\n",
    "        file_path = os.path.join(test_path, filename)\n",
    "        \n",
    "        hr_img = cv2.imread(file_path, cv2.IMREAD_COLOR).astype(np.float32) \n",
    "        h, w, _ = hr_img.shape\n",
    "\n",
    "        ## make sure same dimension\n",
    "        h -= h % upscale\n",
    "        w -= w % upscale\n",
    "        hr_img = hr_img[:h, :w]\n",
    "        \n",
    "        lr_image = cv2.resize(hr_img, (w // upscale, h // upscale), interpolation=cv2.INTER_LINEAR) \n",
    "        lr_image = lr_image/255.\n",
    "        lr_image = torch.from_numpy(np.transpose(lr_image[:, :, [2, 1, 0]],\n",
    "                                                (2, 0, 1))).float()\n",
    "        lr_image = lr_image.unsqueeze(0).to(device)\n",
    "        \n",
    "        \n",
    "        with torch.no_grad():\n",
    "            output = model(lr_image)\n",
    "\n",
    "        output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()\n",
    "        output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))\n",
    "        output = (output * 255.0)\n",
    "\n",
    "        # if (output.shape[2] != hr_img.shape[2]) or (output.shape[3] != hr_img.shape[3]):\n",
    "        #     output = F.interpolate(output, size=(h, w), mode='bilinear', align_corners=False)\n",
    "            \n",
    "        psnr = calc_psnr(hr_img, output)\n",
    "        ssim = calc_ssim(hr_img, output)\n",
    "        test_psnr += psnr\n",
    "        test_ssim += ssim\n",
    "\n",
    "test_psnr = test_psnr/image_count\n",
    "test_ssim = test_ssim/image_count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('test psnr: {:.2f}'.format(test_psnr))\n",
    "print('test ssim: {:.4f}'.format(test_ssim))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from models.PlantSR import PlantSR\n",
    "import torch\n",
    "\n",
    "upscale = 4\n",
    "device = 'cuda'\n",
    "model_path = r'outputs/PlantSR_x4_best.pth'\n",
    "\n",
    "device = torch.device('cuda:0')\n",
    "if upscale == 4:\n",
    "    model = PlantSR(scale=upscale,num_features=96,n_resgroups=16,n_resblocks=4,reduction=16)\n",
    "if upscale == 2:\n",
    "    model = PlantSR(scale=upscale,num_features=32,n_resgroups=16,n_resblocks=4,reduction=16)\n",
    "if upscale == 3:\n",
    "    model = PlantSR(scale=upscale,num_features=64,n_resgroups=16,n_resblocks=4,reduction=16)\n",
    "\n",
    "model.load_state_dict(torch.load(model_path), strict=True)\n",
    "model.eval()\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "input_folder = \"data/PlantSR_dataset/YourData\"\n",
    "output_folder = \"data/PlantSR_dataset/YourDatax2\"\n",
    "\n",
    "os.makedirs(output_folder, exist_ok=True)\n",
    "\n",
    "for filename in os.listdir(input_folder):\n",
    "    if filename.endswith((\".jpg\", \".jpeg\", \".png\")):\n",
    "        img_path = os.path.join(input_folder, filename)\n",
    "\n",
    "        img = cv2.imread(img_path, cv2.IMREAD_COLOR).astype(np.float32) /255.\n",
    "        h, w, _ = img.shape\n",
    "        img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()\n",
    "        img = img.unsqueeze(0).to(device)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            output = model(img)\n",
    "\n",
    "        output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()\n",
    "        output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))\n",
    "        output = (output * 255.0).round().astype(np.uint8)\n",
    "\n",
    "        save_path = os.path.join(output_folder, filename)\n",
    "\n",
    "        cv2.imwrite(save_path, output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}