train_denoising.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df9452a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import skimage\n",
    "import os\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "import argparse\n",
    "import cv2\n",
    "from scipy import io\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision.models as models\n",
    "import torch.optim.lr_scheduler as lr_scheduler\n",
    "from pytorch_msssim import ssim\n",
    "\n",
    "from modules import utils\n",
    "from modules.models import INR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7de91a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser(description='INCODE')\n",
    "\n",
    "# Shared Parameters\n",
    "parser.add_argument('--input',type=str, default='./incode_data/Image/0055.png', help='Input image path')\n",
    "parser.add_argument('--inr_model',type=str, default='incode', help='[gauss, mfn, relu, siren, wire, wire2d, ffn, incode]')\n",
    "parser.add_argument('--lr',type=float, default=1.5e-4, help='Learning rate')\n",
    "parser.add_argument('--using_schedular', type=bool, default=True, help='Whether to use schedular')\n",
    "parser.add_argument('--scheduler_b', type=float, default=0.1, help='LR scheduler [set it from 0.05 to 0.1]')\n",
    "parser.add_argument('--maxpoints', type=int, default=256*256, help='Batch size')\n",
    "parser.add_argument('--niters', type=int, default=501, help='Number if iterations')\n",
    "parser.add_argument('--steps_til_summary', type=int, default=100, help='Number of steps till summary visualization')\n",
    "\n",
    "# INCODE Parameters\n",
    "parser.add_argument('--a_coef',type=float, default=0.1993, help='a coeficient')\n",
    "parser.add_argument('--b_coef',type=float, default=0.0196, help='b coeficient')\n",
    "parser.add_argument('--c_coef',type=float, default=0.0588, help='c coeficient')\n",
    "parser.add_argument('--d_coef',type=float, default=0.0269, help='d coeficient')\n",
    "\n",
    "\n",
    "args = parser.parse_args(args=[])\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c268067e",
   "metadata": {},
   "source": [
    "## Loading Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5d16c2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "im = utils.normalize(plt.imread(args.input).astype(np.float32), True)\n",
    "im = cv2.resize(im, None, fx=1/4, fy=1/4, interpolation=cv2.INTER_AREA)\n",
    "H, W, _ = im.shape\n",
    "\n",
    "# Create a noisy image using realistic sensor measurement\n",
    "tau = 4e1 # Photon noise\n",
    "noise_snr = 2 # Readout noise\n",
    "\n",
    "im_noisy = utils.measure(im, noise_snr, tau)\n",
    "im_noisy_gt = utils.normalize(im_noisy, True)[None, ...].astype(np.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53ab2563",
   "metadata": {},
   "source": [
    "## Defining Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4c717b7",
   "metadata": {},
   "source": [
    "### Defining desired Positional Encoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64cdf32b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Frequency Encoding\n",
    "pos_encode_freq = {'type':'frequency', 'use_nyquist': True, 'mapping_input': int(max(H, W))}\n",
    "\n",
    "# Gaussian Encoding\n",
    "pos_encode_gaus = {'type':'gaussian', 'scale_B': 10, 'mapping_input': 256}\n",
    "\n",
    "# No Encoding\n",
    "pos_encode_no = {'type': None}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65abb4f4",
   "metadata": {},
   "source": [
    "### Model Configureations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22d602b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Harmonizer Configurations\n",
    "MLP_configs={'task': 'denoising',\n",
    "             'model': 'resnet34',\n",
    "             'truncated_layer':5,\n",
    "             'in_channels': 64,             \n",
    "             'hidden_channels': [32, 16, 8, 4],\n",
    "             'mlp_bias':0.0005,\n",
    "             'activation_layer': nn.SiLU,\n",
    "             'GT': torch.tensor(im_noisy_gt).to(device).permute(0, 3, 1, 2)\n",
    "            }\n",
    "\n",
    "### Model Configurations\n",
    "model = INR(args.inr_model).run(in_features=2,\n",
    "                                out_features=3, \n",
    "                                hidden_features=256,\n",
    "                                hidden_layers=3,\n",
    "                                first_omega_0=10.0, # Change it from 5 to 20 according to the noise level; higher noise -> lower omega values)\n",
    "                                hidden_omega_0=30.0,\n",
    "                                pos_encode_configs=pos_encode_no, \n",
    "                                MLP_configs = MLP_configs\n",
    "                               ).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1e99429",
   "metadata": {},
   "source": [
    "## Training Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc475e0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optimizer setup\n",
    "if args.inr_model == 'wire':\n",
    "    args.lr = args.lr * min(1, args.maxpoints / (H * W))\n",
    "optim = torch.optim.Adam(lr=args.lr, params=model.parameters())\n",
    "scheduler = lr_scheduler.LambdaLR(optim, lambda x: args.scheduler_b ** min(x / args.niters, 1))\n",
    "\n",
    "# Initialize lists for PSNR and MSE values\n",
    "psnr_values = []\n",
    "mse_array = torch.zeros(args.niters, device=device)\n",
    "mse_array_noisy = torch.zeros(args.niters, device=device)\n",
    "\n",
    "# Initialize best loss value as positive infinity\n",
    "best_loss = torch.tensor(float('inf'))\n",
    "\n",
    "# Generate coordinate grid\n",
    "coords = utils.get_coords(H, W, dim=2)[None, ...]\n",
    "\n",
    "# Convert input image to a tensor and reshape\n",
    "gt = torch.tensor(im).cuda().reshape(H * W, 3)[None, ...]\n",
    "gt_noisy = torch.tensor(im_noisy).cuda().reshape(H*W, 3)[None, ...]\n",
    "\n",
    "\n",
    "# Initialize a tensor for reconstructed data\n",
    "rec = torch.zeros_like(gt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a452e7cd",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for step in tqdm(range(args.niters)):\n",
    "    # Randomize the order of data points for each iteration\n",
    "    indices = torch.randperm(H*W)\n",
    "\n",
    "    # Process data points in batches\n",
    "    for b_idx in range(0, H*W, args.maxpoints):\n",
    "        b_indices = indices[b_idx:min(H*W, b_idx+args.maxpoints)]\n",
    "        b_coords = coords[:, b_indices, ...].cuda()\n",
    "        b_indices = b_indices.cuda()\n",
    "        \n",
    "        # Calculate model output\n",
    "        if args.inr_model == 'incode':\n",
    "            model_output, coef = model(b_coords)  \n",
    "        else:\n",
    "            model_output = model(b_coords) \n",
    "\n",
    "        # Update the reconstructed data\n",
    "        with torch.no_grad():\n",
    "            rec[:, b_indices, :] = model_output\n",
    "\n",
    "        # Calculate the output loss\n",
    "        output_loss = ((model_output - gt_noisy[:, b_indices, :])**2).mean()\n",
    "        \n",
    "        if args.inr_model == 'incode':\n",
    "            # Calculate regularization loss for 'incode' model\n",
    "            a_coef, b_coef, c_coef, d_coef = coef[0]  \n",
    "            reg_loss = args.a_coef * torch.relu(-a_coef) + \\\n",
    "                       args.b_coef * torch.relu(-b_coef) + \\\n",
    "                       args.c_coef * torch.relu(-c_coef) + \\\n",
    "                       args.d_coef * torch.relu(-d_coef)\n",
    "\n",
    "            # Total loss for 'incode' model\n",
    "            loss = output_loss + reg_loss \n",
    "        else: \n",
    "            # Total loss for other models\n",
    "            loss = output_loss\n",
    "\n",
    "        # Perform backpropagation and update model parameters\n",
    "        optim.zero_grad()\n",
    "        loss.backward()\n",
    "        optim.step()\n",
    "\n",
    "    # Calculate and log mean squared error (MSE) and PSNR\n",
    "    with torch.no_grad():\n",
    "        mse_array_noisy[step] = ((gt_noisy - rec)**2).mean().item()\n",
    "        mse_array[step] = ((gt - rec)**2).mean().item()\n",
    "        psnr = -10*torch.log10(mse_array[step])\n",
    "        psnr_values.append(psnr.item())\n",
    "        \n",
    "        \n",
    "    # Adjust learning rate using a scheduler if applicable\n",
    "    if args.using_schedular:\n",
    "        scheduler.step()\n",
    "\n",
    "    # Prepare reconstructed image for visualization\n",
    "    imrec = rec[0, ...].reshape(H, W, 3).detach().cpu().numpy()\n",
    "\n",
    "    # Check if the current iteration's MSE is the best so far\n",
    "    if (mse_array[step] < best_loss) or (step == 0):\n",
    "        best_loss = mse_array[step]\n",
    "        best_img = imrec\n",
    "        best_img = (best_img - best_img.min()) / (best_img.max() - best_img.min())\n",
    "\n",
    "    # Display intermediate results at specified intervals\n",
    "    if step % args.steps_til_summary == 0:\n",
    "        print(\"Epoch: {} | Total Loss: {:.5f} | PSNR: {:.4f}\".format(step, \n",
    "                                                                     mse_array_noisy[step].item(),\n",
    "                                                                     psnr.item())) \n",
    "        \n",
    "        # Plot\n",
    "        im_noisy = (im_noisy - im_noisy.min()) / (im_noisy.max() - im_noisy.min())\n",
    "        \n",
    "        fig, axes = plt.subplots(1, 3, figsize=(10, 10))\n",
    "        axes[0].set_title('Ground Truth')\n",
    "        axes[0].imshow(im)\n",
    "        axes[0].axis('off')\n",
    "        axes[1].set_title('Noisy Image')\n",
    "        axes[1].imshow(im_noisy)\n",
    "        axes[1].axis('off')\n",
    "        axes[2].set_title('Denoised')\n",
    "        axes[2].imshow(best_img)\n",
    "        axes[2].axis('off')\n",
    "        plt.show()\n",
    "\n",
    "        \n",
    "# Print maximum PSNR achieved during training\n",
    "print('--------------------')\n",
    "print('Max PSNR:', max(psnr_values))\n",
    "print('--------------------')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6bf3625",
   "metadata": {},
   "source": [
    "# Convergance Rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b7d115b",
   "metadata": {},
   "outputs": [],
   "source": [
    "font = {'font': 'Times New Roman', 'size': 12}\n",
    "\n",
    "plt.figure()\n",
    "axfont = {'family' : 'Times New Roman', 'weight' : 'regular', 'size'   : 10}\n",
    "plt.rc('font', **axfont)\n",
    "\n",
    "plt.plot(np.arange(len(psnr_values[:-1])), psnr_values[:-1], label = f\"{(args.inr_model).upper()}\")\n",
    "plt.xlabel('# Epochs', fontdict=font)\n",
    "plt.ylabel('PSNR (dB)', fontdict=font)\n",
    "plt.title('Image Denoising', fontdict={'family': 'Times New Roman', 'size': 12, 'weight': 'bold'})\n",
    "plt.legend()\n",
    "plt.grid(True, color='lightgray')\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "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.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}