image_superres.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "hires_image_superres.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "1QtGpdSrrgAU"
      ],
      "toc_visible": true,
      "authorship_tag": "ABX9TyNID28AjVQ9BF14ALBCns/9",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/olaviinha/NeuralImageSuperResolution/blob/master/image_superres.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X7JkCX2wuHQb",
        "colab_type": "text"
      },
      "source": [
        "# Neural Image Super-Resolution\n",
        "\n",
        "Colab for [superresolution_gan](https://github.com/fukumame/superresolution_gan), implementation of [Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network](https://arxiv.org/abs/1609.04802).\n",
        "\n",
        "Seems to work best on very small images. Images larger than 200x200 (in total number of pixels) will be sliced, processed in tiles and combined."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Tvx2wS76UC8Y",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Mount Drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "drive_root = \"/content/drive/My Drive/\"\n",
        "\n",
        "!pip -q install import-ipynb\n",
        "!pip -q install image_slicer\n",
        "!curl -q -s -O https://raw.githubusercontent.com/olaviinha/inhagcutils/master/inhagcutils.ipynb\n",
        "import import_ipynb\n",
        "from inhagcutils import *"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pgHQ8yl2qXjD",
        "colab_type": "text"
      },
      "source": [
        "### Settings"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "58KGRDniox7A",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Path to an image file in Google Drive \n",
        "input_image = \"ai/superres/doggo-goodboy.png\" #@param {type:\"string\"}\n",
        "input_image = drive_root+input_image\n",
        "\n",
        "#@markdown <font color=\"#999\">Output will be saved in the same directory as `superres_<filename>`</font>\n",
        "\n",
        "dir_out = path_dir(input_image)\n",
        "output_image = path_dir(input_image)+\"superres_\"+path_leaf(input_image)\n",
        "\n",
        "target_grid  = 200*200\n",
        "\n",
        "dir_input_slices = dir_tmp+\"input-slices/\"\n",
        "dir_output_slices = dir_tmp+\"output-slices/\"\n",
        "create_dirs([dir_tmp, dir_input_slices, dir_output_slices])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1QtGpdSrrgAU",
        "colab_type": "text"
      },
      "source": [
        "### Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hx1C5bwEUDqr",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Git clone & get model\n",
        "%cd /content/\n",
        "!git clone {git_q} https://github.com/fukumame/superresolution_gan.git\n",
        "%cd /content/superresolution_gan\n",
        "!wget {wget_q} https://www.dropbox.com/s/l4s5a6v4licks62/generator_model_3008000.npz\n",
        "import image_slicer, math\n",
        "from PIL import Image"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rYQkdvi5qY2P",
        "colab_type": "text"
      },
      "source": [
        "### **Action**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pNf-kqahNEuW",
        "colab_type": "code",
        "colab": {},
        "cellView": "form"
      },
      "source": [
        "#@title Superres\n",
        "im = Image.open(input_image)\n",
        "width, height = im.size\n",
        "\n",
        "# Tile & save tiles\n",
        "def slice_image(input_image, dir):\n",
        "  tiles = closestDiv(math.floor((width*height)/target_grid), 4)\n",
        "  slicery = image_slicer.slice(input_image, tiles, save=False)\n",
        "  image_slicer.save_tiles(slicery, directory=dir, prefix='slice')\n",
        "\n",
        "# Superres tiles\n",
        "def superres(dir_in, dir_out):\n",
        "  slices = glob(dir_in+\"slice*\")\n",
        "  %cd /content/superresolution_gan/\n",
        "  i = 1\n",
        "  for slice in slices:\n",
        "    output = dir_out+path_leaf(slice)\n",
        "    print('> processing slice', str(i)+\"/\"+str(len(slices)))\n",
        "    !python superresolution.py --modelpath=generator_model_3008000.npz --imagepath=\"{slice}\" --outputpath=\"{output}\"\n",
        "    i+=1\n",
        "\n",
        "# Mash\n",
        "def mash(dir_in, img_out):\n",
        "  tiles = image_slicer.open_images_in(dir_in)\n",
        "  image = image_slicer.join(tiles)\n",
        "  tmpImg = dir_tmp+rnd_str(8)\n",
        "  image.save('/content/tmp/joined.png')\n",
        "  !cp \"/content/tmp/joined.png\" \"{output_image}\"\n",
        "\n",
        "# Cleanup\n",
        "def clean():\n",
        "  copy = dir_tmp+path_leaf(input_image)\n",
        "  !rm {dir_input_slices}*\n",
        "  !rm {dir_input_slices}*\n",
        "  !rm {dir_tmp}joined.png\n",
        "  !rm {copy}\n",
        "\n",
        "def closestDiv(a, b):\n",
        "  c1 = a - (a % b)\n",
        "  c2 = (a + b) - (a % b)\n",
        "  if (a - c1 > c2 - a):\n",
        "    return c2\n",
        "  else:\n",
        "    return c1\n",
        "\n",
        "if width*height > target_grid:\n",
        "  !cp \"{input_image}\" \"{dir_tmp}\"\n",
        "  print('Slice', input_image+'...')\n",
        "  slice_image(input_image, dir_input_slices)\n",
        "  print('Done.\\nSuperres slices...')\n",
        "  superres(dir_input_slices, dir_output_slices)\n",
        "  print('Done.\\nMash back into single image...')\n",
        "  mash(dir_output_slices, output_image)\n",
        "  print('.Done\\nFile saved as', output_image)\n",
        "  clean()\n",
        "else:\n",
        "  %cd /content/superresolution_gan\n",
        "  print('Superres', input_image)\n",
        "  !python superresolution.py --modelpath=generator_model_3008000.npz --imagepath=\"{input_image}\" --outputpath=\"{output_image}\"\n",
        "  print('Saved to', output_image)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}