notebooks/example_inpaint_hqs.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%cd ..\n",
    "%config Completer.use_jedi = False\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tensorflow as tf\n",
    "import imageio\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 image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "examples = [\"kate\", \"library\", \"vase\", \"vase2\"]\n",
    "example_id = 2\n",
    "data_path = os.path.join(\"notebooks\", \"data\", \"inpainting\")\n",
    "image = imageio.imread(os.path.join(data_path, f\"{examples[example_id]}.png\"))\n",
    "mask = imageio.imread(os.path.join(data_path, f\"{examples[example_id]}_mask.png\"))\n",
    "\n",
    "image = datasets.from_255_to_1_range(datasets.to_float32(tf.constant(image)))[None, ...]\n",
    "mask = tf.broadcast_to(\n",
    "    datasets.from_255_to_1_range(datasets.to_float32(tf.constant(mask)))[\n",
    "        None, ..., None\n",
    "    ],\n",
    "    shape=tf.shape(image),\n",
    ")\n",
    "\n",
    "image_masked = mask * image\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple for comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import skimage.restoration\n",
    "\n",
    "inpainted_skimage = skimage.restoration.inpaint_biharmonic(\n",
    "    image[0].numpy(),\n",
    "    tf.cast(mask[0, ..., 0] == False, tf.uint8).numpy(),\n",
    "    multichannel=True,\n",
    ")[None, ...]\n",
    "inpainted_border = dppp.inpaint_border(image_masked, mask)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Half-Quadratic Splitting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the denoiser\n",
    "denoiser, (_, denoiser_stddev) = dppp.load_denoiser(\"models/drugan+_0.0-0.2.h5\")\n",
    "\n",
    "# Callbacks: Print the PSNR every 5th step\n",
    "callbacks = [dppp.callback_print_psnr(\"psnr\", 5, image)]\n",
    "\n",
    "reconstructed = dppp.hqs_inpaint(\n",
    "    degraded=image_masked,\n",
    "    mask=mask,\n",
    "    denoiser=denoiser,\n",
    "    max_denoiser_stddev=denoiser_stddev,\n",
    "    callbacks=callbacks,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize.draw_images(\n",
    "    [image_masked[0], inpainted_skimage[0], inpainted_border[0], reconstructed[0]]\n",
    ")\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": 4
}