{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/jantic/DeOldify/blob/master/ImageColorizerColab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "V-nHLm4jWpnV"
   },
   "source": [
    "# DeOldify on Colab #\n",
    "\n",
    "This notebook allows you to colorize your own images using Google Colab!\n",
    "\n",
    "Special thanks to the that made this possible!\n",
    "\n",
    "Original Author:  Matt Robinson, <matthew67robinson@gmail.com>\n",
    "\n",
    "Additional Contributions: Maria Benavente"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 119
    },
    "colab_type": "code",
    "id": "-dSDNbBNb-N8",
    "outputId": "4ef40df9-46dd-44a4-b54a-e720a3eee232"
   },
   "outputs": [],
   "source": [
    "!git clone -b FastAIv1 --single-branch https://github.com/jantic/DeOldify.git DeOldify"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from os import path\n",
    "import torch\n",
    "print(torch.__version__)\n",
    "print(torch.cuda.is_available())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "k19F34Tsd7CX",
    "outputId": "81828b10-6678-4eec-ec53-5b0b9d645782"
   },
   "outputs": [],
   "source": [
    "cd DeOldify"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 187
    },
    "colab_type": "code",
    "id": "lhejMeOxghBM",
    "outputId": "d063f1dc-1286-4355-f8aa-838a7dfc29ee"
   },
   "outputs": [],
   "source": [
    "!pip install PyDrive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "yODBFi8MgoLZ"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from pydrive.auth import GoogleAuth\n",
    "from pydrive.drive import GoogleDrive\n",
    "from google.colab import auth\n",
    "from oauth2client.client import GoogleCredentials\n",
    "from google.colab import drive\n",
    "from IPython.display import Image\n",
    "import fastai\n",
    "from fastai import *\n",
    "from fasterai.visualize import *\n",
    "from pathlib import Path\n",
    "from itertools import repeat\n",
    "from google.colab import drive\n",
    "torch.backends.cudnn.benchmark=True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Ow9Qhf4jgrgd"
   },
   "outputs": [],
   "source": [
    "auth.authenticate_user()\n",
    "gauth = GoogleAuth()\n",
    "gauth.credentials = GoogleCredentials.get_application_default()\n",
    "drive = GoogleDrive(gauth)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "aVyIYMqrg-SM"
   },
   "source": [
    "Note that the above requires a verification step. It isn't too bad."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Doru7d3rVYr7"
   },
   "source": [
    "With access to your Google Drive, the \"deOldifyImages\" directory will be created. Drop there your personal images, and after the full execution of the notebook find the results at its subdirectory \"results\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "sU2yQAfqhNJv"
   },
   "outputs": [],
   "source": [
    "results_dir=Path('/content/drive/My Drive/deOldifyImages/results')\n",
    "\n",
    "#Adjust this if image doesn't look quite right (max 64 on 11GB GPU).  The default here works for most photos.  \n",
    "#It literally just is a number multiplied by 16 to get the square render resolution.  \n",
    "#Note that this doesn't affect the resolution of the final output- the output is the same resolution as the input.\n",
    "#Example:  render_factor=21 => color is rendered at 16x21 = 336x336 px.  \n",
    "render_factor=36"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#◢ Artistic vs Stable Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "artistic = True #@param {type:\"boolean\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir 'models'\n",
    "if artistic:\n",
    "    !wget https://www.dropbox.com/s/zkehq1uwahhbc2o/ColorizeArtistic_gen.pth?dl=0 -O ./models/ColorizeArtistic_gen.pth\n",
    "else:\n",
    "    !wget https://www.dropbox.com/s/mwjep3vyqk5mkjc/ColorizeStable_gen.pth?dl=0 -O ./models/ColorizeStable_gen.pth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 68
    },
    "colab_type": "code",
    "id": "NlOT9IlBHkk7",
    "outputId": "1bb5cc4d-15a2-4174-a37c-9e11f3a2dce3"
   },
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir \"/content/drive/My Drive/deOldifyImages\"\n",
    "!mkdir \"/content/drive/My Drive/deOldifyImages/results\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vis = get_image_colorizer(render_factor=render_factor, artistic=artistic)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ZpCf0qbxicVK"
   },
   "source": [
    "Here's an example of colorizing an image downloaded from the internet:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 204
    },
    "colab_type": "code",
    "id": "8-hQC-AfiCKq",
    "outputId": "f52e94a1-48bd-4b27-bc21-2306ee85ef9d"
   },
   "outputs": [],
   "source": [
    "!wget \"https://media.githubusercontent.com/media/jantic/DeOldify/master/test_images/TV1930s.jpg\" -O \"family_TV.jpg\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 434
    },
    "colab_type": "code",
    "id": "O6kfUN0GiJsq",
    "outputId": "3e37f84d-6d27-4e4e-b7b6-9feae8be40ae"
   },
   "outputs": [],
   "source": [
    "vis.plot_transformed_image('family_TV.jpg')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "dBY1N_bYaIxq"
   },
   "source": [
    "If you want to colorize pictures from your Google drive, drop them in a directory named deOldifyImages (in the root of your drive) and the next cell will colorize all of them and save the resulting images in deOldifyImages/results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 559
    },
    "colab_type": "code",
    "id": "vL7775DWaFJ4",
    "outputId": "cdb35172-d43a-4d45-f23c-d86c993aaf9b"
   },
   "outputs": [],
   "source": [
    "for img in os.listdir(\"/content/drive/My Drive/deOldifyImages/\"):\n",
    "  img_path = str(\"/content/drive/My Drive/deOldifyImages/\") + img\n",
    "  if os.path.isfile(img_path):\n",
    "    vis.plot_transformed_image(img_path)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "DeOldify_colab.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}