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

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ScitaPqhKtuW"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Hub Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jvztxQ6VsK2k"
      },
      "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": "oXlcl8lqBgAD"
      },
      "source": [
        "# 任意风格的快速风格转换\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>     <a target=\"_blank\" href=\"https://tensorflow.google.cn/hub/tutorials/tf2_arbitrary_image_stylization\"><img src=\"https://tensorflow.google.cn/images/tf_logo_32px.png\">在 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/tf2_arbitrary_image_stylization.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/tf2_arbitrary_image_stylization.ipynb\"><img src=\"https://tensorflow.google.cn/images/GitHub-Mark-32px.png\">在 GitHub 上查看源代码</a>\n",
        "</td>\n",
        "  <td>     <a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/zh-cn/hub/tutorials/tf2_arbitrary_image_stylization.ipynb\"><img src=\"https://tensorflow.google.cn/images/download_logo_32px.png\">下载笔记本</a>\n",
        "</td>\n",
        "  <td>     <a href=\"https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2\"><img src=\"https://tensorflow.google.cn/images/hub_logo_32px.png\">查看 TF Hub 模型 </a>\n",
        "</td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YeeuYzbZcJzs"
      },
      "source": [
        "基于 [magenta](https://github.com/tensorflow/magenta/tree/master/magenta/models/arbitrary_image_stylization) 中的模型代码和以下论文：\n",
        "\n",
        "[Exploring the structure of a real-time, arbitrary neural artistic stylization network](https://arxiv.org/abs/1705.06830). *Golnaz Ghiasi, Honglak Lee, Manjunath Kudlur, Vincent Dumoulin, Jonathon Shlens*, Proceedings of the British Machine Vision Conference (BMVC), 2017.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TaM8BVxrCA2E"
      },
      "source": [
        "## 设置"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J65jog2ncJzt"
      },
      "source": [
        "首先，我们导入 TF2 和所有相关依赖项。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v-KXRY5XBu2u"
      },
      "outputs": [],
      "source": [
        "import functools\n",
        "import os\n",
        "\n",
        "from matplotlib import gridspec\n",
        "import matplotlib.pylab as plt\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "\n",
        "print(\"TF Version: \", tf.__version__)\n",
        "print(\"TF Hub version: \", hub.__version__)\n",
        "print(\"Eager mode enabled: \", tf.executing_eagerly())\n",
        "print(\"GPU available: \", tf.config.list_physical_devices('GPU'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tsoDv_9geoZn"
      },
      "outputs": [],
      "source": [
        "# @title Define image loading and visualization functions  { display-mode: \"form\" }\n",
        "\n",
        "def crop_center(image):\n",
        "  \"\"\"Returns a cropped square image.\"\"\"\n",
        "  shape = image.shape\n",
        "  new_shape = min(shape[1], shape[2])\n",
        "  offset_y = max(shape[1] - shape[2], 0) // 2\n",
        "  offset_x = max(shape[2] - shape[1], 0) // 2\n",
        "  image = tf.image.crop_to_bounding_box(\n",
        "      image, offset_y, offset_x, new_shape, new_shape)\n",
        "  return image\n",
        "\n",
        "@functools.lru_cache(maxsize=None)\n",
        "def load_image(image_url, image_size=(256, 256), preserve_aspect_ratio=True):\n",
        "  \"\"\"Loads and preprocesses images.\"\"\"\n",
        "  # Cache image file locally.\n",
        "  image_path = tf.keras.utils.get_file(os.path.basename(image_url)[-128:], image_url)\n",
        "  # Load and convert to float32 numpy array, add batch dimension, and normalize to range [0, 1].\n",
        "  img = tf.io.decode_image(\n",
        "      tf.io.read_file(image_path),\n",
        "      channels=3, dtype=tf.float32)[tf.newaxis, ...]\n",
        "  img = crop_center(img)\n",
        "  img = tf.image.resize(img, image_size, preserve_aspect_ratio=True)\n",
        "  return img\n",
        "\n",
        "def show_n(images, titles=('',)):\n",
        "  n = len(images)\n",
        "  image_sizes = [image.shape[1] for image in images]\n",
        "  w = (image_sizes[0] * 6) // 320\n",
        "  plt.figure(figsize=(w * n, w))\n",
        "  gs = gridspec.GridSpec(1, n, width_ratios=image_sizes)\n",
        "  for i in range(n):\n",
        "    plt.subplot(gs[i])\n",
        "    plt.imshow(images[i][0], aspect='equal')\n",
        "    plt.axis('off')\n",
        "    plt.title(titles[i] if len(titles) > i else '')\n",
        "  plt.show()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8etHh05-CJHc"
      },
      "source": [
        "加载一些图像来看看效果。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dRc0vat3Alzo"
      },
      "outputs": [],
      "source": [
        "# @title Load example images  { display-mode: \"form\" }\n",
        "\n",
        "content_image_url = 'https://upload.wikimedia.org/wikipedia/commons/thumb/f/fd/Golden_Gate_Bridge_from_Battery_Spencer.jpg/640px-Golden_Gate_Bridge_from_Battery_Spencer.jpg'  # @param {type:\"string\"}\n",
        "style_image_url = 'https://upload.wikimedia.org/wikipedia/commons/0/0a/The_Great_Wave_off_Kanagawa.jpg'  # @param {type:\"string\"}\n",
        "output_image_size = 384  # @param {type:\"integer\"}\n",
        "\n",
        "# The content image size can be arbitrary.\n",
        "content_img_size = (output_image_size, output_image_size)\n",
        "# The style prediction model was trained with image size 256 and it's the \n",
        "# recommended image size for the style image (though, other sizes work as \n",
        "# well but will lead to different results).\n",
        "style_img_size = (256, 256)  # Recommended to keep it at 256.\n",
        "\n",
        "content_image = load_image(content_image_url, content_img_size)\n",
        "style_image = load_image(style_image_url, style_img_size)\n",
        "style_image = tf.nn.avg_pool(style_image, ksize=[3,3], strides=[1,1], padding='SAME')\n",
        "show_n([content_image, style_image], ['Content image', 'Style image'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yL2Bn5ThR1nY"
      },
      "source": [
        "## 导入 TF Hub 模块"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "467AVDSuzBPc"
      },
      "outputs": [],
      "source": [
        "# Load TF Hub module.\n",
        "\n",
        "hub_handle = 'https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2'\n",
        "hub_module = hub.load(hub_handle)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uAR70_3wLEDB"
      },
      "source": [
        "该 Hub 模块用于图像风格化的签名为：\n",
        "\n",
        "```\n",
        "outputs = hub_module(content_image, style_image)\n",
        "stylized_image = outputs[0]\n",
        "```\n",
        "\n",
        "其中，`content_image`、`style_image` 和 `stylized_image` 预期是形状为 `[batch_size, image_height, image_width, 3]` 的四维张量。\n",
        "\n",
        "在当前示例中，我们仅提供单个图像，因此，批次维度为 1，但是您可以使用同一模块同时处理更多图像。\n",
        "\n",
        "图像的输入和输出值应在 [0, 1] 范围内。\n",
        "\n",
        "内容与风格图像的形状不一定要匹配。输出图像形状与内容图像形状相同。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qEhYJno1R7rP"
      },
      "source": [
        "## 演示图像风格化"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lnAv-F3O9fLV"
      },
      "outputs": [],
      "source": [
        "# Stylize content image with given style image.\n",
        "# This is pretty fast within a few milliseconds on a GPU.\n",
        "\n",
        "outputs = hub_module(tf.constant(content_image), tf.constant(style_image))\n",
        "stylized_image = outputs[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OEAPEdq698gs"
      },
      "outputs": [],
      "source": [
        "# Visualize input images and the generated stylized image.\n",
        "\n",
        "show_n([content_image, style_image, stylized_image], titles=['Original content image', 'Style image', 'Stylized image'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v-gYvjTWK-lx"
      },
      "source": [
        "## 试试处理更多图像"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WSMaY0YBNfkK"
      },
      "outputs": [],
      "source": [
        "# @title To Run: Load more images { display-mode: \"form\" }\n",
        "\n",
        "content_urls = dict(\n",
        "  sea_turtle='https://upload.wikimedia.org/wikipedia/commons/d/d7/Green_Sea_Turtle_grazing_seagrass.jpg',\n",
        "  tuebingen='https://upload.wikimedia.org/wikipedia/commons/0/00/Tuebingen_Neckarfront.jpg',\n",
        "  grace_hopper='https://storage.googleapis.com/download.tensorflow.org/example_images/grace_hopper.jpg',\n",
        "  )\n",
        "style_urls = dict(\n",
        "  kanagawa_great_wave='https://upload.wikimedia.org/wikipedia/commons/0/0a/The_Great_Wave_off_Kanagawa.jpg',\n",
        "  kandinsky_composition_7='https://upload.wikimedia.org/wikipedia/commons/b/b4/Vassily_Kandinsky%2C_1913_-_Composition_7.jpg',\n",
        "  hubble_pillars_of_creation='https://upload.wikimedia.org/wikipedia/commons/6/68/Pillars_of_creation_2014_HST_WFC3-UVIS_full-res_denoised.jpg',\n",
        "  van_gogh_starry_night='https://upload.wikimedia.org/wikipedia/commons/thumb/e/ea/Van_Gogh_-_Starry_Night_-_Google_Art_Project.jpg/1024px-Van_Gogh_-_Starry_Night_-_Google_Art_Project.jpg',\n",
        "  turner_nantes='https://upload.wikimedia.org/wikipedia/commons/b/b7/JMW_Turner_-_Nantes_from_the_Ile_Feydeau.jpg',\n",
        "  munch_scream='https://upload.wikimedia.org/wikipedia/commons/c/c5/Edvard_Munch%2C_1893%2C_The_Scream%2C_oil%2C_tempera_and_pastel_on_cardboard%2C_91_x_73_cm%2C_National_Gallery_of_Norway.jpg',\n",
        "  picasso_demoiselles_avignon='https://upload.wikimedia.org/wikipedia/en/4/4c/Les_Demoiselles_d%27Avignon.jpg',\n",
        "  picasso_violin='https://upload.wikimedia.org/wikipedia/en/3/3c/Pablo_Picasso%2C_1911-12%2C_Violon_%28Violin%29%2C_oil_on_canvas%2C_Kr%C3%B6ller-M%C3%BCller_Museum%2C_Otterlo%2C_Netherlands.jpg',\n",
        "  picasso_bottle_of_rum='https://upload.wikimedia.org/wikipedia/en/7/7f/Pablo_Picasso%2C_1911%2C_Still_Life_with_a_Bottle_of_Rum%2C_oil_on_canvas%2C_61.3_x_50.5_cm%2C_Metropolitan_Museum_of_Art%2C_New_York.jpg',\n",
        "  fire='https://upload.wikimedia.org/wikipedia/commons/3/36/Large_bonfire.jpg',\n",
        "  derkovits_woman_head='https://upload.wikimedia.org/wikipedia/commons/0/0d/Derkovits_Gyula_Woman_head_1922.jpg',\n",
        "  amadeo_style_life='https://upload.wikimedia.org/wikipedia/commons/8/8e/Untitled_%28Still_life%29_%281913%29_-_Amadeo_Souza-Cardoso_%281887-1918%29_%2817385824283%29.jpg',\n",
        "  derkovtis_talig='https://upload.wikimedia.org/wikipedia/commons/3/37/Derkovits_Gyula_Talig%C3%A1s_1920.jpg',\n",
        "  amadeo_cardoso='https://upload.wikimedia.org/wikipedia/commons/7/7d/Amadeo_de_Souza-Cardoso%2C_1915_-_Landscape_with_black_figure.jpg'\n",
        ")\n",
        "\n",
        "content_image_size = 384\n",
        "style_image_size = 256\n",
        "content_images = {k: load_image(v, (content_image_size, content_image_size)) for k, v in content_urls.items()}\n",
        "style_images = {k: load_image(v, (style_image_size, style_image_size)) for k, v in style_urls.items()}\n",
        "style_images = {k: tf.nn.avg_pool(style_image, ksize=[3,3], strides=[1,1], padding='SAME') for k, style_image in style_images.items()}\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dqB6aNTLNVkK"
      },
      "outputs": [],
      "source": [
        "#@title Specify the main content image and the style you want to use.  { display-mode: \"form\" }\n",
        "\n",
        "content_name = 'sea_turtle'  # @param ['sea_turtle', 'tuebingen', 'grace_hopper']\n",
        "style_name = 'munch_scream'  # @param ['kanagawa_great_wave', 'kandinsky_composition_7', 'hubble_pillars_of_creation', 'van_gogh_starry_night', 'turner_nantes', 'munch_scream', 'picasso_demoiselles_avignon', 'picasso_violin', 'picasso_bottle_of_rum', 'fire', 'derkovits_woman_head', 'amadeo_style_life', 'derkovtis_talig', 'amadeo_cardoso']\n",
        "\n",
        "stylized_image = hub_module(tf.constant(content_images[content_name]),\n",
        "                            tf.constant(style_images[style_name]))[0]\n",
        "\n",
        "show_n([content_images[content_name], style_images[style_name], stylized_image],\n",
        "       titles=['Original content image', 'Style image', 'Stylized image'])"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "tf2_arbitrary_image_stylization.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}