site/zh-cn/hub/tutorials/image_enhancing.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GeerbrLA0uju"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Hub Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "\n",
        "Created by @[Adrish Dey](https://github.com/captain-pool) for [Google Summer of Code](https://summerofcode.withgoogle.com/) 2019"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yntM0JbY0uj5"
      },
      "outputs": [],
      "source": [
        "# Copyright 2019 The TensorFlow Hub Authors. All Rights Reserved.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS, \n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# =============================================================================="
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UJeo2a5C0uj2"
      },
      "source": [
        "# 使用 ESRGAN 进行图像超解析"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ps4toA1d_tkc"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>     <a target=\"_blank\" href=\"https://tensorflow.google.cn/hub/tutorials/image_enhancing\"><img src=\"https://tensorflow.google.cn/images/tf_logo_32px.png\">View on TensorFlow.org</a>\n",
        "</td>\n",
        "  <td><a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/zh-cn/hub/tutorials/image_enhancing.ipynb\"><img src=\"https://tensorflow.google.cn/images/colab_logo_32px.png\">在 Google Colab 中运行</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/zh-cn/hub/tutorials/image_enhancing.ipynb\">     <img src=\"https://tensorflow.google.cn/images/GitHub-Mark-32px.png\">     在 GitHub 上查看源代码</a></td>\n",
        "  <td>     <a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/zh-cn/hub/tutorials/image_enhancing.ipynb\"><img src=\"https://tensorflow.google.cn/images/download_logo_32px.png\">下载笔记本</a>   </td>\n",
        "  <td><a href=\"https://tfhub.dev/captain-pool/esrgan-tf2/1\"><img src=\"https://tensorflow.google.cn/images/hub_logo_32px.png\">查看 TF Hub 模型</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LkW9jAmt_zjB"
      },
      "source": [
        "此 Colab 演示了在 Enhanced Super Resolution Generative Adversarial Network（*由 Xintao Wang 等人撰写*）[[论文](https://arxiv.org/pdf/1809.00219.pdf)] [[代码](https://github.com/captain-pool/GSOC/)] 中如何使用 TensorFlow Hub 模块\n",
        "\n",
        "进行图像增强。*（最好使用双三次降采样的图像。）*\n",
        "\n",
        "模型在 DIV2K 数据集（双三次降采样的图像）中大小为 128 x 128 的图像块上进行了训练。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LBGty4O_0ukJ"
      },
      "source": [
        "**准备环境**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lnyLTyUt0ukN"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import time\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "import matplotlib.pyplot as plt\n",
        "os.environ[\"TFHUB_DOWNLOAD_PROGRESS\"] = \"True\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dremsFdh0ukX"
      },
      "outputs": [],
      "source": [
        "!wget \"https://user-images.githubusercontent.com/12981474/40157448-eff91f06-5953-11e8-9a37-f6b5693fa03f.png\" -O original.png"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DXot2kru0ukh"
      },
      "outputs": [],
      "source": [
        "# Declaring Constants\n",
        "IMAGE_PATH = \"original.png\"\n",
        "SAVED_MODEL_PATH = \"https://tfhub.dev/captain-pool/esrgan-tf2/1\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KF_tHde-p3rn"
      },
      "source": [
        "**定义辅助函数**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IslbQmTj0ukz"
      },
      "outputs": [],
      "source": [
        "def preprocess_image(image_path):\n",
        "  \"\"\" Loads image from path and preprocesses to make it model ready\n",
        "      Args:\n",
        "        image_path: Path to the image file\n",
        "  \"\"\"\n",
        "  hr_image = tf.image.decode_image(tf.io.read_file(image_path))\n",
        "  # If PNG, remove the alpha channel. The model only supports\n",
        "  # images with 3 color channels.\n",
        "  if hr_image.shape[-1] == 4:\n",
        "    hr_image = hr_image[...,:-1]\n",
        "  hr_size = (tf.convert_to_tensor(hr_image.shape[:-1]) // 4) * 4\n",
        "  hr_image = tf.image.crop_to_bounding_box(hr_image, 0, 0, hr_size[0], hr_size[1])\n",
        "  hr_image = tf.cast(hr_image, tf.float32)\n",
        "  return tf.expand_dims(hr_image, 0)\n",
        "\n",
        "def save_image(image, filename):\n",
        "  \"\"\"\n",
        "    Saves unscaled Tensor Images.\n",
        "    Args:\n",
        "      image: 3D image tensor. [height, width, channels]\n",
        "      filename: Name of the file to save.\n",
        "  \"\"\"\n",
        "  if not isinstance(image, Image.Image):\n",
        "    image = tf.clip_by_value(image, 0, 255)\n",
        "    image = Image.fromarray(tf.cast(image, tf.uint8).numpy())\n",
        "  image.save(\"%s.jpg\" % filename)\n",
        "  print(\"Saved as %s.jpg\" % filename)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uh1E2rBpnWxV"
      },
      "outputs": [],
      "source": [
        "%matplotlib inline\n",
        "def plot_image(image, title=\"\"):\n",
        "  \"\"\"\n",
        "    Plots images from image tensors.\n",
        "    Args:\n",
        "      image: 3D image tensor. [height, width, channels].\n",
        "      title: Title to display in the plot.\n",
        "  \"\"\"\n",
        "  image = np.asarray(image)\n",
        "  image = tf.clip_by_value(image, 0, 255)\n",
        "  image = Image.fromarray(tf.cast(image, tf.uint8).numpy())\n",
        "  plt.imshow(image)\n",
        "  plt.axis(\"off\")\n",
        "  plt.title(title)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ycrCTvmlqBMD"
      },
      "source": [
        "#### 对从路径加载的图像执行超解析"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L7XpMk8Y0uk7"
      },
      "outputs": [],
      "source": [
        "hr_image = preprocess_image(IMAGE_PATH)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hWgCbUa_0ulG"
      },
      "outputs": [],
      "source": [
        "# Plotting Original Resolution image\n",
        "plot_image(tf.squeeze(hr_image), title=\"Original Image\")\n",
        "save_image(tf.squeeze(hr_image), filename=\"Original Image\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ouwEyKLA0ulO"
      },
      "outputs": [],
      "source": [
        "model = hub.load(SAVED_MODEL_PATH)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dz79ncnT0ulX"
      },
      "outputs": [],
      "source": [
        "start = time.time()\n",
        "fake_image = model(hr_image)\n",
        "fake_image = tf.squeeze(fake_image)\n",
        "print(\"Time Taken: %f\" % (time.time() - start))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ABjkkJHC2jNL"
      },
      "outputs": [],
      "source": [
        "# Plotting Super Resolution Image\n",
        "plot_image(tf.squeeze(fake_image), title=\"Super Resolution\")\n",
        "save_image(tf.squeeze(fake_image), filename=\"Super Resolution\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tuKu18UYptkx"
      },
      "source": [
        "### 评估模型的性能"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Qdz55sxMgiwO"
      },
      "outputs": [],
      "source": [
        "!wget \"https://lh4.googleusercontent.com/-Anmw5df4gj0/AAAAAAAAAAI/AAAAAAAAAAc/6HxU8XFLnQE/photo.jpg64\" -O test.jpg\n",
        "IMAGE_PATH = \"test.jpg\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F6tMNtqy0ukq"
      },
      "outputs": [],
      "source": [
        "# Defining helper functions\n",
        "def downscale_image(image):\n",
        "  \"\"\"\n",
        "      Scales down images using bicubic downsampling.\n",
        "      Args:\n",
        "          image: 3D or 4D tensor of preprocessed image\n",
        "  \"\"\"\n",
        "  image_size = []\n",
        "  if len(image.shape) == 3:\n",
        "    image_size = [image.shape[1], image.shape[0]]\n",
        "  else:\n",
        "    raise ValueError(\"Dimension mismatch. Can work only on single image.\")\n",
        "\n",
        "  image = tf.squeeze(\n",
        "      tf.cast(\n",
        "          tf.clip_by_value(image, 0, 255), tf.uint8))\n",
        "\n",
        "  lr_image = np.asarray(\n",
        "    Image.fromarray(image.numpy())\n",
        "    .resize([image_size[0] // 4, image_size[1] // 4],\n",
        "              Image.BICUBIC))\n",
        "\n",
        "  lr_image = tf.expand_dims(lr_image, 0)\n",
        "  lr_image = tf.cast(lr_image, tf.float32)\n",
        "  return lr_image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r2ANR1XDy77I"
      },
      "outputs": [],
      "source": [
        "hr_image = preprocess_image(IMAGE_PATH)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r_dautO6qbTV"
      },
      "outputs": [],
      "source": [
        "lr_image = downscale_image(tf.squeeze(hr_image))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jRw1x6xY0ulj"
      },
      "outputs": [],
      "source": [
        "# Plotting Low Resolution Image\n",
        "plot_image(tf.squeeze(lr_image), title=\"Low Resolution\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g--yyHg7qXCw"
      },
      "outputs": [],
      "source": [
        "model = hub.load(SAVED_MODEL_PATH)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZX-deZlhqaYz"
      },
      "outputs": [],
      "source": [
        "start = time.time()\n",
        "fake_image = model(lr_image)\n",
        "fake_image = tf.squeeze(fake_image)\n",
        "print(\"Time Taken: %f\" % (time.time() - start))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AmSga6MSq1PB"
      },
      "outputs": [],
      "source": [
        "plot_image(tf.squeeze(fake_image), title=\"Super Resolution\")\n",
        "# Calculating PSNR wrt Original Image\n",
        "psnr = tf.image.psnr(\n",
        "    tf.clip_by_value(fake_image, 0, 255),\n",
        "    tf.clip_by_value(hr_image, 0, 255), max_val=255)\n",
        "print(\"PSNR Achieved: %f\" % psnr)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5YTBKCXPq9UZ"
      },
      "source": [
        "**并排比较输出大小。**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ubdupldDypCy"
      },
      "outputs": [],
      "source": [
        "plt.rcParams['figure.figsize'] = [15, 10]\n",
        "fig, axes = plt.subplots(1, 3)\n",
        "fig.tight_layout()\n",
        "plt.subplot(131)\n",
        "plot_image(tf.squeeze(hr_image), title=\"Original\")\n",
        "plt.subplot(132)\n",
        "fig.tight_layout()\n",
        "plot_image(tf.squeeze(lr_image), \"x4 Bicubic\")\n",
        "plt.subplot(133)\n",
        "fig.tight_layout()\n",
        "plot_image(tf.squeeze(fake_image), \"Super Resolution\")\n",
        "plt.savefig(\"ESRGAN_DIV2K.jpg\", bbox_inches=\"tight\")\n",
        "print(\"PSNR: %f\" % psnr)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "image_enhancing.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}