train_image.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/0010.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=9e-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='Learning rate scheduler')\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"
   ]
  },
  {
   "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)/3)}\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": "36c1145d",
   "metadata": {},
   "source": [
    "### Model Configureations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22d602b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Harmonizer Configurations\n",
    "MLP_configs={'task': 'image',\n",
    "             'model': 'resnet34',\n",
    "             'truncated_layer':5,\n",
    "             'in_channels': 64,             \n",
    "             'hidden_channels': [64, 32, 4],\n",
    "             'mlp_bias':0.3120,\n",
    "             'activation_layer': nn.SiLU,\n",
    "             'GT': torch.tensor(im).to(device)[None,...].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=30.0,\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": "db268c3d",
   "metadata": {},
   "source": [
    "## Training Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f58b062",
   "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",
    "\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).reshape(H * W, 3)[None, ...].to(device)\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, ...].to(device)\n",
    "        b_indices = b_indices.to(device)\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[:, 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 PSNR\n",
    "    with torch.no_grad():\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",
    "    # Adjust learning rate using a scheduler if applicable\n",
    "    if args.using_schedular:\n",
    "        if args.inr_model == 'incode' and 30 < step:\n",
    "            scheduler.step()\n",
    "        else:\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 loss 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[step].item(),\n",
    "                                                                     psnr.item()))\n",
    "        \n",
    "        # Plot\n",
    "        fig, axes = plt.subplots(1, 2, figsize=(6, 6))\n",
    "        axes[0].set_title('Ground Truth')\n",
    "        axes[0].imshow(im)\n",
    "        axes[0].axis('off')\n",
    "        axes[1].set_title('Reconstructed')\n",
    "        axes[1].imshow(best_img)\n",
    "        axes[1].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": "c86e4b9e",
   "metadata": {},
   "source": [
    "# Convergance Rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1dc2a55",
   "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 Representation', 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
}