BeautifyOldPhoto.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Mark1708/beautify-old-photo/blob/main/BeautifyOldPhoto.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xHY6kHTmp2xP"
      },
      "source": [
        "# Libraries used:\n",
        "\n",
        "*   [Coloize](https://github.com/jantic/DeOldify)\n",
        "*   [Restoration of faces](https://github.com/TencentARC/GFPGAN)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PE-SXjdPAYxu"
      },
      "source": [
        "# Download libs and install requirements \n",
        "(Don't forget restart)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6wyYmvcG-dRX"
      },
      "outputs": [],
      "source": [
        "# If the link does not work, you can download the archive from the link - https://drive.google.com/file/d/1O-neWA-asLI1c1-LTxdpTtzzmkMkw4cF/view?usp=sharing\n",
        "!wget \"https://www.dropbox.com/s/2yrzr98wj1faxi3/libs_beautify_photo.zip\"\n",
        "!unzip \"/content/libs_beautify_photo.zip\" && rm \"/content/libs_beautify_photo.zip\"\n",
        "!rm -r \"/content/__MACOSX\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1Vz-IKn-Aqrs"
      },
      "outputs": [],
      "source": [
        "!pip install -r /content/colab_requirements.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u32PXA-CBRWf"
      },
      "source": [
        "# Import libs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "18VHWF_FBQc7"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "import os\n",
        "\n",
        "sys.path.append('/content/libs/GFPGAN')\n",
        "sys.path.append('/content/libs/DeOldify')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KwiqKVwIBnqu"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import fastai\n",
        "from deoldify import device\n",
        "from deoldify.visualize import *\n",
        "from deoldify.device_id import DeviceId\n",
        "from pathlib import Path"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J4iIvvnxMUf5"
      },
      "source": [
        "# Prepare image folders\n",
        "(Don't forget upload images in `input/` folder)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uZ4jL17AMT4C"
      },
      "outputs": [],
      "source": [
        "!mkdir '/content/input' \n",
        "!mkdir '/content/high_resolution_output' \n",
        "!mkdir '/content/final_colorized_output'\n",
        "!mkdir '/content/final_colorized_output/from_original' \n",
        "!mkdir '/content/final_colorized_output/from_recontraction'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bQZB_NGaIZUw"
      },
      "source": [
        "# Improve quality"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Uw9WVFtfCCXz"
      },
      "outputs": [],
      "source": [
        "!cd /content/libs/GFPGAN && python setup.py develop\n",
        "!cd /content/libs/GFPGAN && python inference_gfpgan.py -i ../../input -o ../../high_resolution_output -v 1.3 -s 2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wA_stvK6KG_q"
      },
      "source": [
        "# Colorize images"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2tvvTDdxLsKW"
      },
      "outputs": [],
      "source": [
        "device.set(device=DeviceId.GPU0)\n",
        "if not torch.cuda.is_available():\n",
        "    print('GPU not available.')\n",
        "torch.backends.cudnn.benchmark = True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hz2FGV1tLymL"
      },
      "outputs": [],
      "source": [
        "colorizer = get_image_colorizer(artistic=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IGM9PciOKFlw"
      },
      "outputs": [],
      "source": [
        "input = \"Both\" #@param [\"Original images\", \"Images after Face Restoration\", \"Both\"]\n",
        "# result_dir_path = '/content/result_images' #@param {type:\"string\"}\n",
        "render_factor = 17  #@param {type: \"slider\", min: 7, max: 40}\n",
        "# watermarked = False #@param {type:\"boolean\"}\n",
        "# display_render_factor = False #@param {type:\"boolean\"}\n",
        "image_pathes = []\n",
        "result_pathes = []\n",
        "if input == \"Original images\":\n",
        "  image_pathes.append('/content/input')\n",
        "  result_pathes.append('/content/final_colorized_output/from_original')\n",
        "if input == \"Images after Face Restoration\":\n",
        "  image_pathes.append('/content/high_resolution_output/restored_imgs')\n",
        "  result_pathes.append('/content/final_colorized_output/from_recontraction')\n",
        "if input == \"Both\":\n",
        "  image_pathes.append('/content/input')\n",
        "  result_pathes.append('/content/final_colorized_output/from_original')\n",
        "  image_pathes.append('/content/high_resolution_output/restored_imgs')\n",
        "  result_pathes.append('/content/final_colorized_output/from_recontraction')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bc4QjhHFhRiz"
      },
      "source": [
        "If you don't need to draw the result, you need to comment out lines 108-116 in the `plot_transformed_image` method (when there are too many photos, rendering slows down and causes a block overflow error)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VLtUE3QWLheD"
      },
      "outputs": [],
      "source": [
        "def colorize_photo(image_path, result_dir_path, render_factor, watermarked=False, display_render_factor=False):\n",
        "  colorizer.plot_transformed_image(image_path, results_dir=Path(result_dir_path), render_factor=render_factor, display_render_factor=display_render_factor, figsize=(8,8), watermarked=watermarked)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KtDrnz5MPMkF"
      },
      "outputs": [],
      "source": [
        "import imghdr"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LRKUyCWzNVu2"
      },
      "outputs": [],
      "source": [
        "for i in range(len(image_pathes)):\n",
        "  # os.mkdir(result_pathes[i])\n",
        "  for file in os.listdir(image_pathes[i]):\n",
        "    full_path = image_pathes[i] + \"/\" + file\n",
        "    if os.path.isfile(full_path) and imghdr.what(full_path) in ['jpg', 'png', 'jpeg']:\n",
        "      colorize_photo(full_path, result_pathes[i], render_factor) "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WyBcVjkErevS"
      },
      "source": [
        "# Ziped result"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OskLh8nkUMJL"
      },
      "outputs": [],
      "source": [
        "!zip -r \"colorized_output.zip\" \"/content/final_colorized_output\" "
      ]
    }
  ],
  "metadata": {
    "colab": {
      "authorship_tag": "ABX9TyPyoh8+t9BbR8DC2TU5ujR0",
      "collapsed_sections": [
        "PE-SXjdPAYxu"
      ],
      "include_colab_link": true,
      "name": "BeautifyOldPhoto.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}