{
 "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/DeOldify_colab.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 shows how to get your own version of [DeOldify](https://github.com/jantic/DeOldify) working on Google Colab. A lot of the initial steps are just installs -- but these are also the steps that can make running the model a tedious exercise.\n",
    "\n",
    "Initially, one must `pip install` a few dependencies, then `wget` is used to download the appropriate picture data.\n",
    "\n",
    "NECESSARY PRELIMINARY STEP: Please make sure you have gone up to the \"Runtime\" menu above and \"Change Runtime Type\" to Python3 and GPU.\n",
    "\n",
    "I hope you have fun, and thanks to Jason Antic for this awesome tool!\n",
    "\n",
    "-Matt Robinson, <matthew67robinson@gmail.com>\n",
    "\n",
    "NEW: You can now load your files from you own Google Drive, check the last cell of the notebook for more information."
   ]
  },
  {
   "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 https://github.com/jantic/DeOldify.git DeOldify"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from os import path\n",
    "from wheel.pep425tags import get_abbr_impl, get_impl_ver, get_abi_tag\n",
    "platform = '{}{}-{}'.format(get_abbr_impl(), get_impl_ver(), get_abi_tag())\n",
    "\n",
    "accelerator = 'cu80' if path.exists('/opt/bin/nvidia-smi') else 'cpu'\n",
    "\n",
    "!pip install -q http://download.pytorch.org/whl/{accelerator}/torch-0.4.1-{platform}-linux_x86_64.whl torchvision\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": 1295
    },
    "colab_type": "code",
    "id": "BMRO_XT6e_ge",
    "outputId": "58cc061e-7547-40d9-9db2-e35cda1495f0"
   },
   "outputs": [],
   "source": [
    "!pip install -e ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "id": "bNDYBcMZePD9",
    "outputId": "81c9ba6a-5875-44bf-9fc4-e2ba8c6d5f00"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "print(torch.__version__)\n",
    "print(torch.cuda.is_available())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "3VUHs2kxfWe8"
   },
   "outputs": [],
   "source": [
    "import fastai\n",
    "from fastai import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YllLOZKjgbf4"
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "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": [
    "# Doing work so I can access data from my google drive\n",
    "!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",
    "import multiprocessing\n",
    "from torch import autograd\n",
    "from fastai.transforms import TfmType\n",
    "from fasterai.transforms import *\n",
    "from fastai.conv_learner import *\n",
    "from fasterai.images import *\n",
    "from fasterai.dataset import *\n",
    "from fasterai.visualize import *\n",
    "from fasterai.callbacks import *\n",
    "from fasterai.loss import *\n",
    "from fasterai.modules import *\n",
    "from fasterai.training import *\n",
    "from fasterai.generators import *\n",
    "from fastai.torch_imports import *\n",
    "from fasterai.filters import *\n",
    "from pathlib import Path\n",
    "from itertools import repeat\n",
    "from google.colab import drive\n",
    "from IPython.display import Image\n",
    "import tensorboardX\n",
    "torch.cuda.set_device(0)\n",
    "plt.style.use('dark_background')\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": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ny3egicahAB4"
   },
   "outputs": [],
   "source": [
    "# Now download the pretrained weights, which I have saved to my google drive\n",
    "# note that the id is the ending part of the shareable link url (after open?id=)\n",
    "# The pretrained weights can be downloaded from https://www.dropbox.com/s/7r2wu0af6okv280/colorize_gen_192.h5\n",
    "download = drive.CreateFile({'id': '1mRRvS3WIHPdp36G0yc1jC0XI6i-Narv6'})\n",
    "download.GetContentFile('pretrained_weights.h5')"
   ]
  },
  {
   "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": {
     "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')\n",
    "!mkdir \"/content/drive/My Drive/deOldifyImages\"\n",
    "!mkdir \"/content/drive/My Drive/deOldifyImages/results\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "sU2yQAfqhNJv"
   },
   "outputs": [],
   "source": [
    "weights_path = 'pretrained_weights.h5'\n",
    "results_dir='/content/drive/My Drive/deOldifyImages/results'\n",
    "\n",
    "#The higher the render_factor, the more GPU memory will be used and generally images will look better.  \n",
    "#11GB can take a factor of 42 max.  Performance generally gracefully degrades with lower factors, \n",
    "#though you may also find that certiain images will actually render better at lower numbers.  \n",
    "#This tends to be the case with the oldest photos.\n",
    "render_factor=42\n",
    "filters = [Colorizer(gpu=0, weights_path=weights_path)]\n",
    "vis = ModelImageVisualizer(filters, render_factor=render_factor, results_dir=results_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 204
    },
    "colab_type": "code",
    "id": "-86zKacjhnDE",
    "outputId": "1438b061-68f5-45af-ad9e-0fb058aa1854"
   },
   "outputs": [],
   "source": [
    "# download an example picture to try.\n",
    "# NOTE: All the jpg files cloned from the git repo are corrupted. Must download yourself.\n",
    "!wget \"https://media.githubusercontent.com/media/jantic/DeOldify/master/test_images/abe.jpg\" -O \"abe2.jpg\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 687
    },
    "colab_type": "code",
    "id": "vz0BmRcGh1Gd",
    "outputId": "91bc0d27-f8ff-476c-ee1f-bdc248517a3a"
   },
   "outputs": [],
   "source": [
    "# %matplotlib inline\n",
    "vis.plot_transformed_image('abe2.jpg', render_factor=25)"
   ]
  },
  {
   "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', render_factor=41)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ZpCf0qbxicVK"
   },
   "source": [
    "Let's see how well it does Dorothy before her world turns to color in the Wizard of Oz:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 204
    },
    "colab_type": "code",
    "id": "xBhxm79riXLc",
    "outputId": "a4e73e62-ab11-4dd0-de79-03fb96974c18"
   },
   "outputs": [],
   "source": [
    "!wget \"https://magnoliaforever.files.wordpress.com/2011/09/wizard-of-oz.jpg\" -O \"Dorothy.jpg\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 391
    },
    "colab_type": "code",
    "id": "yTgTXqx_iil3",
    "outputId": "0cdfb10a-2b64-422a-b9ab-625b7d4cb33e"
   },
   "outputs": [],
   "source": [
    "vis.plot_transformed_image('Dorothy.jpg', render_factor=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GXM77gcunp7M"
   },
   "source": [
    "Let's now try Butch and Sundance. Famously the last scene ends with a black and white still. So we know what the color was beforehand."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 204
    },
    "colab_type": "code",
    "id": "fpoVn8O1nlE1",
    "outputId": "c1cf4a09-3c81-4074-fad0-28dc94847d18"
   },
   "outputs": [],
   "source": [
    "!wget \"https://i.ebayimg.com/images/g/HqkAAOSwRLZUAwyS/s-l300.jpg\" -O \"butch_and_sundance.jpg\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 457
    },
    "colab_type": "code",
    "id": "c_dsWcBjnu_p",
    "outputId": "5bae5428-6896-4434-c586-c22fa7275779"
   },
   "outputs": [],
   "source": [
    "vis.plot_transformed_image('butch_and_sundance.jpg', render_factor=29)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Ra5fs0UWn46V"
   },
   "source": [
    "Let's get a picture of what they were actually wearing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 204
    },
    "colab_type": "code",
    "id": "a8zoOYJanxkk",
    "outputId": "2d85a505-93c7-4289-80cb-a23578b187b7"
   },
   "outputs": [],
   "source": [
    "!wget \"https://bethanytompkins.files.wordpress.com/2015/09/freezeframe.jpg\" -O \"butch_and_sundance_color.jpg\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 273
    },
    "colab_type": "code",
    "id": "nEULX0Uan-pi",
    "outputId": "6fe834de-863c-46f9-a740-df75cacb76e7"
   },
   "outputs": [],
   "source": [
    "Image('butch_and_sundance_color.jpg')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "dBY1N_bYaIxq"
   },
   "source": [
    "If you want to colorise pictures from your drive, drop them in a directory named deOldifyImages (in the root of your drive) and the next cell will save the colorise pictures 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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}