Comput Stats.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Computing Stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from model.losses import Stats \n",
    "import torch\n",
    "import torchvision \n",
    "import os \n",
    "from PIL import Image\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from functorch import vmap \n",
    "from torchmetrics import StructuralSimilarityIndexMeasure\n",
    "import lpips"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class Stats(nn.Module):\n",
    "    def __init__(self, lpips=False):\n",
    "        self.lpips = lpips\n",
    "        super().__init__()\n",
    "\n",
    "    def error_mse(self, im_pred, im_gt, mask = None):\n",
    "        \"\"\"\n",
    "        Computes MSE metric. Optionally applies mask.\n",
    "        \"\"\"\n",
    "        # Linearize.\n",
    "        im_pred = im_pred[..., :3].reshape(-1, 3)\n",
    "        im_gt = im_gt[..., :3].reshape(-1, 3)\n",
    "\n",
    "        # Mask?\n",
    "        if mask is not None:\n",
    "            # mask = mask.flatten()\n",
    "            # im_pred = im_pred[mask, :]\n",
    "            # im_gt = im_gt[mask, :]\n",
    "\n",
    "            # Use multiplication method as described in paper\n",
    "            mask = mask.reshape(-1,1)\n",
    "            im_pred = im_pred * mask[..., None]\n",
    "            im_gt = im_gt * mask[..., None]\n",
    "\n",
    "        mse = (im_pred - im_gt) ** 2\n",
    "        return mse.mean()\n",
    "\n",
    "    def error_psnr(self, im_pred, im_gt, mask=None):\n",
    "        \"\"\"\n",
    "        Computes PSNR metric. Optionally applies mask.\n",
    "        Assumes floats [0,1].\n",
    "        \"\"\"\n",
    "        mse = self.error_mse(im_pred, im_gt, mask)\n",
    "        # https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio\n",
    "        return 20 * torch.log10(torch.tensor(1.0)) - 10 * torch.log10(mse), mse\n",
    "\n",
    "\n",
    "    def error_ssim(self, im_pred, im_gt, pad=0, pad_mode='linear_ramp'):\n",
    "        \"\"\"Compute the LPIPS metric.\"\"\"\n",
    "        ssim = StructuralSimilarityIndexMeasure(data_range=1.0)\n",
    "        return ssim(im_pred, im_gt)\n",
    "\n",
    "    def error_lpips(self, im_pred, im_gt, mask=None, metric=None):\n",
    "        \"\"\"\n",
    "        Computes LPIPS metric. Optionally applies mask.\n",
    "        \"\"\"\n",
    "        # Mask?\n",
    "        if mask is not None:\n",
    "            print(mask.shape)\n",
    "            mask = mask.reshape(im_pred.shape[0], im_pred.shape[1], 1).repeat(3, axis=2)\n",
    "            print(mask.shape)\n",
    "            im_pred = im_pred * mask\n",
    "            im_gt = im_gt * mask\n",
    "\n",
    "        # To torch.\n",
    "        device = 'cuda'\n",
    "        print(im_gt.shape, im_pred.shape)\n",
    "        im_pred = (im_pred).to(device)\n",
    "        im_gt = (im_gt).to(device)\n",
    "        print(im_gt.shape, im_pred.shape)\n",
    "        # Make metric.\n",
    "        if metric is None:\n",
    "            # best forward scores\n",
    "            metric_a = lpips.LPIPS(net='alex').to(device)\n",
    "            # # closer to \"traditional\" perceptual loss, when used for optimization\n",
    "            metric_v = lpips.LPIPS(net='vgg').to(device)\n",
    "\n",
    "        # Compute metric.\n",
    "        loss_a = metric_a(im_pred, im_gt)\n",
    "        loss_v = metric_v(im_pred, im_gt)\n",
    "\n",
    "        return loss_a.item(), loss_v.item()\n",
    "\n",
    "    def forward(self, im_pred, im_gt, mask=None):\n",
    "        psnr, mse = self.error_psnr(im_pred, im_gt, mask=mask)\n",
    "        ssim = self.error_ssim(im_pred, im_gt)\n",
    "        if self.lpips:\n",
    "            lpips_alex, lpips_vgg = self.error_lpips(im_pred, im_gt, mask=mask)\n",
    "            stats = {\n",
    "                'psnr': psnr,\n",
    "                'lpips_alex': lpips_alex,\n",
    "                'lpips_vgg': lpips_vgg,\n",
    "                'ssim': ssim\n",
    "            }\n",
    "        else: \n",
    "            stats = {\n",
    "                'psnr': psnr,\n",
    "                'ssim': ssim\n",
    "            }\n",
    "        return stats"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##curls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'psnr': tensor(35.5179), 'ssim': tensor(0.8130)}\n"
     ]
    }
   ],
   "source": [
    "\n",
    "GT = torchvision.transforms.functional.to_tensor(Image.open('./data/DIY/curls/scan/image/000038.png'))\n",
    "folder = 'se3'\n",
    "num = '448'\n",
    "pred = torchvision.transforms.functional.to_tensor(Image.open('./out/DIY/curls_' + folder + '/rendering/'+num+'000_vis/0038_unisurf.png'))\n",
    "\n",
    "get_stats = Stats(lpips=False)\n",
    "\n",
    "print(get_stats(pred.unsqueeze(0), GT.unsqueeze(0)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##toby"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'psnr': tensor(24.6651), 'ssim': tensor(0.4442)}\n"
     ]
    }
   ],
   "source": [
    "\n",
    "GT = torchvision.transforms.functional.to_tensor(Image.open('/ubc/cs/research/kmyi/svsamban/research/unisurf/data/DIY/toby/scan/image/000040.png'))\n",
    "folder = 'ctrl'\n",
    "num = '447'\n",
    "pred_file = '/ubc/cs/research/kmyi/svsamban/research/unisurf/out/DIY/toby_' + folder + '/rendering/'+num+'000_vis/0040_unisurf.png'\n",
    "pred = torchvision.transforms.functional.to_tensor(Image.open(pred_file))\n",
    "\n",
    "get_stats = Stats(lpips=False)\n",
    "\n",
    "print(get_stats(pred.unsqueeze(0), GT.unsqueeze(0)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'psnr': tensor(11.2125), 'ssim': tensor(0.3147)}\n"
     ]
    }
   ],
   "source": [
    "#METANLR \n",
    "\n",
    "folder = 'se3'\n",
    "\n",
    "pic ='/ubc/cs/research/kmyi/svsamban/research/metanlrpp/logs/MYcurls_'+folder+'/meshes/model_final/test_benchmark/f000000_v038_color_pred.png'\n",
    "pred = torchvision.transforms.functional.to_tensor(Image.open(pic))\n",
    "\n",
    "GT = torchvision.transforms.functional.to_tensor(Image.open('./data/DIY/curls/scan/image/000038.png'))\n",
    "GT = torchvision.transforms.Resize(pred.shape[1:])(GT)\n",
    "mask = torchvision.transforms.functional.to_tensor(Image.open('./data/DIY/curls/scan/mask/000038.png'))\n",
    "mask = torchvision.transforms.Resize(pred.shape[1:])(mask)#.squeeze()\n",
    "get_stats = Stats(lpips=False)\n",
    "print(get_stats(pred.unsqueeze(0), GT.unsqueeze(0), mask=mask))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'psnr': tensor(19.9665), 'ssim': tensor(0.3065)}\n"
     ]
    }
   ],
   "source": [
    "#METANLR \n",
    "\n",
    "folder = 'control'\n",
    "\n",
    "pic ='/ubc/cs/research/kmyi/svsamban/research/metanlrpp/logs/MYtoby_'+folder+'/meshes/model_final/test_benchmark/f000000_v040_color_pred.png'\n",
    "pred = torchvision.transforms.functional.to_tensor(Image.open(pic))\n",
    "\n",
    "GT = torchvision.transforms.functional.to_tensor(Image.open('./data/DIY/toby/scan/image/000040.png'))\n",
    "GT = torchvision.transforms.Resize(pred.shape[1:])(GT)\n",
    "mask = torchvision.transforms.functional.to_tensor(Image.open('./data/DIY/toby/scan/mask/000040.png'))\n",
    "mask = torchvision.transforms.Resize(pred.shape[1:])(mask)#.squeeze()\n",
    "get_stats = Stats(lpips=False)\n",
    "print(get_stats(pred.unsqueeze(0), GT.unsqueeze(0), mask=mask))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'psnr': tensor(25.2052), 'ssim': tensor(0.4740)}\n"
     ]
    }
   ],
   "source": [
    "#NERF\n",
    "gt = '/ubc/cs/research/kmyi/svsamban/research/data-nerfies/curls/rgb/4x/right1_000041.png'\n",
    "gt = '/ubc/cs/research/kmyi/svsamban/research/unisurf/data/DIY/curlstest/scan/image/000038.png'\n",
    "GT = torchvision.transforms.functional.to_tensor(Image.open(gt))\n",
    "# pred = '/ubc/cs/research/kmyi/svsamban/research/hypernerf/out/curls50/renders/00250000/val/rgb_right1_000041.png'\n",
    "# pred = '/ubc/cs/research/kmyi/svsamban/research/nerfies/out/curls50/renders/00250000/val/rgb_right1_000041.png'\n",
    "pred = '/ubc/cs/research/kmyi/svsamban/research/unisurf/out/DIY/curls_ebt/rendering/1000000_vis/0038_unisurf.png'\n",
    "pred = torchvision.transforms.functional.to_tensor(Image.open(pred))\n",
    "\n",
    "get_stats = Stats(lpips=False)\n",
    "\n",
    "print(get_stats(pred.unsqueeze(0), GT.unsqueeze(0)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(600, 800)\n"
     ]
    }
   ],
   "source": [
    "from PIL import Image\n",
    "image = Image.open('/ubc/cs/research/kmyi/svsamban/research/unisurf/data/DIY/curlstest/scan/image/000038.png')\n",
    "print(image.size)\n",
    "# image.thumbnail((600, 800))\n",
    "# image.save('/ubc/cs/research/kmyi/svsamban/research/unisurf/data/DIY/curlstest/scan/image/000038.png')\n",
    "# print(image.size) # Output: (400, 350) -->\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.8 ('unisurf2')",
   "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.10.8 | packaged by conda-forge | (main, Nov 22 2022, 08:23:14) [GCC 10.4.0]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "a798088355447effe12dbac206397bf2d5a84e0d94f3ff66d84cb0806de32ad9"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}