notebooks/example_deblur_na_dmsp.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd ..\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tensorflow as tf\n",
    "\n",
    "tf.get_logger().setLevel(\"ERROR\")\n",
    "import tensorflow_datasets as tfds\n",
    "\n",
    "import tensorflow_datasets_bw as datasets\n",
    "from tensorflow_datasets_bw import visualize\n",
    "import dppp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load an example dataset and example kernels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = (\n",
    "    tfds.load(name=\"set14\", split=\"test\")\n",
    "    .map(datasets.get_value(\"hr\"))\n",
    "    .map(datasets.to_float32)\n",
    "    .map(datasets.from_255_to_1_range)\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.random.set_seed(99)  # Random seed for reproducability\n",
    "mode = \"wrap\"  # Border mode ('constant' or 'wrap')\n",
    "noise_stddev = 0.04\n",
    "image = datasets.get_one_example(images, index=0)[None, ...]\n",
    "kernel = dppp.conv2D_filter_rgb(dppp.NB_DEBLURRING_LEVIN_KERNELS[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Degrade the image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "degraded = dppp.blur(image, kernel, noise_stddev, clip_final=False, mode=mode)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reconstruct the image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the denoiser\n",
    "denoiser, (denoiser_min_stddev, denoiser_max_stddev) = dppp.load_denoiser(\n",
    "    \"models/drugan+_0.0-0.2.h5\"\n",
    ")\n",
    "\n",
    "# Define the noise stddev used for the stochastic evaluation of the prior\n",
    "if denoiser_min_stddev <= 0.1 <= denoiser_max_stddev:\n",
    "    denoiser_stddev = 0.1\n",
    "else:\n",
    "    denoiser_stddev = denoiser_max_stddev\n",
    "\n",
    "# Callbacks\n",
    "callbacks = [\n",
    "    # Print the PSNR every 20th step\n",
    "    dppp.callback_print_psnr(\"psnr\", 20, image),\n",
    "]\n",
    "\n",
    "# Reconstruct the image\n",
    "# NOTE: The noise_stddev is not given to the algorithm\n",
    "reconstructed = dppp.dmsp_deblur_nb(\n",
    "    degraded=degraded,\n",
    "    kernel=kernel,\n",
    "    denoiser=denoiser,\n",
    "    denoiser_stddev=denoiser_stddev,\n",
    "    num_steps=300,\n",
    "    conv_mode=mode,\n",
    "    callbacks=callbacks,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print PSNR and LPIPS\n",
    "psnr = dppp.psnr(image, reconstructed).numpy()[0]\n",
    "lpips = dppp.lpips_alex(image, reconstructed).numpy()[0]\n",
    "print(f\"Reconstructed PSNR: {psnr:0.2f}, LPIPS: {lpips:0.4f}\")\n",
    "\n",
    "# Visualize\n",
    "visualize.draw_images([image[0], degraded[0], reconstructed[0]])\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}