deconv_ml_experiments/eval_superres.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "30ee9e6a-b33f-4bb2-aee2-94129862e8d2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "633849ebce294857b837be718a6c6917",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "269b409238f844e192a993b75d974445",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9755382a94a94d15b9bfa3b8274e0b68",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a5cdc4e95680435a898b45d895711c82",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f0100682100>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib widget\n",
    "if 'NOTEBOOK_INITIATED_FLAG' not in globals():\n",
    "    NOTEBOOK_INITIATED_FLAG = True\n",
    "    %cd ..\n",
    "\n",
    "import os \n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import image\n",
    "import matplotlib.colors as colors\n",
    "\n",
    "import examplesetup\n",
    "import lib.deconvolution as deconvolution\n",
    "import lib.utils as utils\n",
    "import lib.xtrace as xtrace\n",
    "import lib.datagen as datagen\n",
    "import cv2\n",
    "from numpy.random import default_rng\n",
    "rng = default_rng()\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "import numpy as np\n",
    "\n",
    "model_name = \"superres_nodeconv_abs\"\n",
    "model = keras.models.load_model(f'trained_networks/{model_name}.h5')\n",
    "\n",
    "#img = image.imread(\"example_image.png\").sum(axis=2)\n",
    "img_9 = np.load('example_files/test_data_9.npy')/1e6\n",
    "img_0 = np.load('example_files/test_data_0.npy')/1e6\n",
    "\n",
    "img, true_img = datagen.SuperRes(1, 1, (400, 400), 4)[0]\n",
    "img = img[0,:,:,0]\n",
    "true_img = true_img[0,:,:,0]\n",
    "recovered_img = model(img[np.newaxis,:,:, np.newaxis], training=False).numpy()[0,:,:,0]\n",
    "recovered_img_0 = model(img_0[np.newaxis,:,:, np.newaxis], training=False).numpy()[0,:,:,0]\n",
    "recovered_img_9 = model(img_9[np.newaxis,:,:, np.newaxis], training=False).numpy()[0,:,:,0]\n",
    "percent = 99\n",
    "perc = np.percentile(img_0, percent)\n",
    "#G = xtrace.depth_spill_psf(datagen.global_config, *utils.ray_grid(datagen.global_config[\"dimensions\"]))\n",
    "#recovered_img_conv_0 = deconvolution.richard_lucy(distorted_img_0, G, 6)\n",
    "\n",
    "fig, axs = plt.subplots(1, 3, figsize=(10, 5), sharex=True, sharey=True)\n",
    "axs[0].imshow(true_img, vmin=0, vmax=perc)\n",
    "axs[0].set_title(\"True Image\")\n",
    "axs[1].imshow(img, vmin=0, vmax=perc,extent=[0.5, 400, 400, 0.5])\n",
    "axs[1].set_title(\"Small Img\")\n",
    "axs[2].imshow(recovered_img, vmin=0, vmax=perc)\n",
    "axs[2].set_title(\"Recovered\")\n",
    "\n",
    "fig, axs = plt.subplots(2, 2, figsize=(10, 10), sharex=True, sharey=True)\n",
    "axs[0,0].imshow(img_0, vmin=0, vmax=perc, extent=[0.5, 400, 400, 0.5])\n",
    "axs[0,0].set_title(\"Detector\")\n",
    "axs[0,1].imshow(recovered_img_0, vmin=0, vmax=perc)\n",
    "axs[0,1].set_title(\"Recovered\")\n",
    "axs[1,0].imshow(img_9, vmin=0, vmax=perc, extent=[0.5, 400, 400, 0.5])\n",
    "axs[1,0].set_title(\"Detector\")\n",
    "axs[1,1].imshow(recovered_img_9, vmin=0, vmax=perc)\n",
    "axs[1,1].set_title(\"Recovered\")\n",
    "\n",
    "G = xtrace.depth_spill_psf(datagen.global_config, *utils.ray_grid(datagen.global_config[\"dimensions\"]))\n",
    "recovered_img_conv_0 = deconvolution.richard_lucy(img_0, G, 6)\n",
    "distorted_conv = (G.get()@recovered_img_conv_0.flatten()).reshape(recovered_img_conv_0.shape)\n",
    "\n",
    "config, _ = utils.upsample_transform(datagen.global_config, np.zeros(1), 4)\n",
    "G_up = xtrace.depth_spill_psf(config, *utils.ray_grid(config[\"dimensions\"]))\n",
    "distorted_ml = utils.downsample_img((G_up.get()@recovered_img_0.flatten()).reshape(recovered_img_0.shape), 4)\n",
    "\n",
    "mm = np.max(img_0)\n",
    "diff_rl = np.abs(img_0 - distorted_conv)/mm\n",
    "#diff_rl = np.sign(diff_rl)*np.sqrt(np.abs(diff_rl))\n",
    "diff_ml = np.abs(img_0 - distorted_ml)/mm\n",
    "#diff_ml = np.sign(diff_ml)*np.sqrt(np.abs(diff_ml))\n",
    "\n",
    "\n",
    "vmax = np.max(np.maximum(diff_rl, diff_ml))\n",
    "norm = colors.PowerNorm(gamma=0.5, vmin=0, vmax=vmax)\n",
    "fig, axs = plt.subplots(1, 2, figsize=(20, 10), sharex=True, sharey=True, constrained_layout=True)\n",
    "#pcm_1 = axs[0].imshow(distorted_img_0, vmax=perc, vmin=0)\n",
    "pcm = axs[0].imshow(diff_rl, norm=norm)\n",
    "axs[0].set_title(\"Difference (original - RL)\")\n",
    "axs[1].imshow(diff_ml, norm=norm)\n",
    "axs[1].set_title(\"Difference (original - ML)\")\n",
    "fig.colorbar(pcm, ax=axs, shrink=1, location='bottom')\n",
    "plt.show()\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(img_0)\n",
    "plt.title(\"Original Image\")\n",
    "plt.colorbar()\n",
    "\n",
    "\n",
    "#Add local energy loss constraint? (ex: 10x10 area should contain about the same ammount of energy)\n",
    "#Add positivity loss constraint?\n",
    "#Try the Pecent contraint with lower randomness value -> trains on more realistic schenarios"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "716ad78a-4cce-4756-a82a-9fb030ee5fdf",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8762369-0f36-4bcf-973c-4ce5f42ce138",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "HDF5 / Simple Analysis / GPU",
   "language": "python",
   "name": "maxiv-jhub-docker-kernel-hdf5"
  },
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}