notebooks/example_vsr_hqs.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 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\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Bicubic Downscaling\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scale_factor = 4\n",
    "resize_method = \"bicubic\"\n",
    "builder_kwargs = {\n",
    "    \"resize_method\": resize_method,\n",
    "    \"scale\": scale_factor,\n",
    "    \"antialias\": True,\n",
    "}\n",
    "\n",
    "dataset = tfds.load(\"vid4\", split=\"test\", builder_kwargs=builder_kwargs).map(\n",
    "    datasets.compose(\n",
    "        datasets.map_on_dict(datasets.to_float32),\n",
    "        datasets.map_on_dict(datasets.from_255_to_1_range),\n",
    "    )\n",
    ")\n",
    "\n",
    "example = datasets.get_one_example(dataset, index=0)\n",
    "video_hr = example[\"hr\"]\n",
    "video_lr = example[\"lr\"]\n",
    "ref_index = tf.shape(video_hr)[0] // 2\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reconstruct\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the denoiser\n",
    "denoiser, (denoiser_min, denoiser_max) = dppp.load_denoiser(\"models/drugan+_0.0-0.2.h5\")\n",
    "num_frames = 9\n",
    "resize_fn = dppp.create_resize_fn(resize_method, True)\n",
    "\n",
    "# Crop the video\n",
    "video_frames = video_lr.shape[0]\n",
    "if video_frames > num_frames:\n",
    "    start_index = max(0, ref_index - (num_frames // 2))\n",
    "    end_index = start_index + num_frames\n",
    "    video_lr_cropped = video_lr[start_index:end_index]\n",
    "    ref_index_cropped = ref_index - start_index\n",
    "else:\n",
    "    video_lr_cropped = video_lr\n",
    "    ref_index_cropped = ref_index\n",
    "\n",
    "# Callbacks: Print the PSNR every 5th step\n",
    "callbacks = [dppp.callback_print_psnr(\"psnr\", 5, video_hr[ref_index])]\n",
    "\n",
    "reconstructed = dppp.hqs_video_super_resolve(\n",
    "    degraded=video_lr_cropped,\n",
    "    ref_index=ref_index_cropped,\n",
    "    sr_factor=scale_factor,\n",
    "    denoiser=denoiser,\n",
    "    max_denoiser_stddev=denoiser_max,\n",
    "    resize_fn=resize_fn,\n",
    "    callbacks=callbacks,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print PSNR and LPIPS\n",
    "psnr = dppp.psnr(video_hr[ref_index], reconstructed[0]).numpy()\n",
    "lpips = dppp.lpips_alex(video_hr[ref_index], reconstructed[0]).numpy()\n",
    "print(f\"Reconstructed PSNR: {psnr:0.2f}, LPIPS: {lpips:0.4f}\")\n",
    "\n",
    "# Visualize\n",
    "visualize.draw_images([video_hr[ref_index], video_lr[ref_index], reconstructed[0]])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using a Gaussian Blur Kernel\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scale_factor = 4\n",
    "example = datasets.get_one_example(dataset, index=2)\n",
    "video_hr = example[\"hr\"]\n",
    "\n",
    "kernel = dppp.conv2D_filter_rgb(dppp.ZHANG_GAUSSIAN_KERNELS[1])\n",
    "video_lr = dppp.create_convolve_resize_fn(kernel, mode=\"wrap\")(\n",
    "    video_hr, scale_factor, False\n",
    ")\n",
    "ref_index = tf.shape(video_hr)[0] // 2\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reconstruct\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the denoiser\n",
    "denoiser, (denoiser_min, denoiser_max) = dppp.load_denoiser(\"models/drugan+_0.0-0.2.h5\")\n",
    "num_frames = 9\n",
    "\n",
    "# Crop the video\n",
    "video_frames = video_lr.shape[0]\n",
    "if video_frames > num_frames:\n",
    "    start_index = max(0, ref_index - (num_frames // 2))\n",
    "    end_index = start_index + num_frames\n",
    "    video_lr_cropped = video_lr[start_index:end_index]\n",
    "    ref_index_cropped = ref_index - start_index\n",
    "else:\n",
    "    video_lr_cropped = video_lr\n",
    "    ref_index_cropped = ref_index\n",
    "\n",
    "# Callbacks: Print the PSNR every 5th step\n",
    "callbacks = [dppp.callback_print_psnr(\"psnr\", 5, video_hr[ref_index])]\n",
    "\n",
    "reconstructed = dppp.hqs_video_super_resolve(\n",
    "    degraded=video_lr_cropped,\n",
    "    ref_index=ref_index_cropped,\n",
    "    sr_factor=scale_factor,\n",
    "    denoiser=denoiser,\n",
    "    max_denoiser_stddev=denoiser_max,\n",
    "    kernel=kernel,\n",
    "    callbacks=callbacks,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print PSNR and LPIPS\n",
    "psnr = dppp.psnr(video_hr[ref_index], reconstructed[0]).numpy()\n",
    "lpips = dppp.lpips_alex(video_hr[ref_index], reconstructed[0]).numpy()\n",
    "print(f\"Reconstructed PSNR: {psnr:0.2f}, LPIPS: {lpips:0.4f}\")\n",
    "\n",
    "# Visualize\n",
    "visualize.draw_images([video_hr[ref_index], video_lr[ref_index], reconstructed[0]])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Unkown Downscaling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scale_factor = 4\n",
    "resize_fn = dppp.create_resize_fn(\"bicubic\", True)\n",
    "\n",
    "dataset = tfds.load(\"mdsp_color_sr\", split=\"test\").map(\n",
    "    datasets.compose(\n",
    "        datasets.get_value(\"video\"),\n",
    "        datasets.to_float32,\n",
    "        datasets.from_255_to_1_range,\n",
    "    )\n",
    ")\n",
    "\n",
    "video_lr = datasets.get_one_example(dataset, index=0)\n",
    "ref_index = tf.shape(video_lr)[0] // 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reconstruct\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the denoiser\n",
    "denoiser, (denoiser_min, denoiser_max) = dppp.load_denoiser(\"models/drugan+_0.0-0.2.h5\")\n",
    "num_frames = 9\n",
    "resize_fn = dppp.create_resize_fn(resize_method, True)\n",
    "\n",
    "# Crop the video\n",
    "video_frames = video_lr.shape[0]\n",
    "if video_frames > num_frames:\n",
    "    start_index = max(0, ref_index - (num_frames // 2))\n",
    "    end_index = start_index + num_frames\n",
    "    video_lr_cropped = video_lr[start_index:end_index]\n",
    "    ref_index_cropped = ref_index - start_index\n",
    "else:\n",
    "    video_lr_cropped = video_lr\n",
    "    ref_index_cropped = ref_index\n",
    "\n",
    "reconstructed = dppp.hqs_video_super_resolve(\n",
    "    degraded=video_lr_cropped,\n",
    "    ref_index=ref_index_cropped,\n",
    "    sr_factor=scale_factor,\n",
    "    denoiser=denoiser,\n",
    "    max_denoiser_stddev=denoiser_max,\n",
    "    resize_fn=resize_fn,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize.draw_images([video_lr[ref_index], reconstructed[0]])"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
  },
  "kernelspec": {
   "display_name": "Python 3.6.9 64-bit",
   "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
}