Updating Colab to use latest model/fastai upgrade

DeOldify_colab.ipynb

@@ -19,17 +19,13 @@
    "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",
+    "This notebook allows you to colorize your own images using Google Colab!\n",
     "\n",
-    "Initially, one must `pip install` a few dependencies, then `wget` is used to download the appropriate picture data.\n",
+    "Special thanks to the that made this possible!\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",
+    "Original Author:  Matt Robinson, <matthew67robinson@gmail.com>\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."
+    "Additional Contributions: Maria Benavente"
    ]
   },
   {
@@ -46,7 +42,7 @@
    },
    "outputs": [],
    "source": [
-    "!git clone https://github.com/jantic/DeOldify.git DeOldify"
+    "!git clone -b FastAIv1 --single-branch https://github.com/jantic/DeOldify.git DeOldify"
    ]
   },
   {
@@ -56,12 +52,6 @@
    "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())"
@@ -84,38 +74,6 @@
     "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": {},
-    "colab_type": "code",
-    "id": "YllLOZKjgbf4"
-   },
-   "outputs": [],
-   "source": [
-    "%matplotlib inline\n",
-    "%reload_ext autoreload\n",
-    "%autoreload 2"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -140,8 +98,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Work around with Pillow being preinstalled on these Colab VMs, causing conflicts otherwise.\n",
-    "!pip install Pillow==4.1.1"
+    "!pip install tensorboardX"
    ]
   },
   {
@@ -159,28 +116,22 @@
     "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 google.colab import drive\n",
+    "from IPython.display import Image\n",
+    "import fastai\n",
+    "from fastai import *\n",
+    "from fastai.vision import *\n",
+    "from fastai.callbacks import *\n",
+    "from fastai.vision.gan import *\n",
     "from fasterai.dataset import *\n",
     "from fasterai.visualize import *\n",
-    "from fasterai.callbacks import *\n",
+    "from fasterai.tensorboard 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 fasterai.generators 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"
    ]
   },
@@ -210,23 +161,6 @@
     "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": {
@@ -237,26 +171,6 @@
     "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,
@@ -267,70 +181,32 @@
    },
    "outputs": [],
    "source": [
-    "weights_path = 'pretrained_weights.h5'\n",
-    "results_dir='/content/drive/My Drive/deOldifyImages/results'\n",
+    "weights_name = 'colorize_gen'\n",
+    "results_dir=Path('/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 certain 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 = [Colorizer34(gpu=0, weights_path=weights_path)]\n",
-    "vis = ModelImageVisualizer(filters, render_factor=render_factor, results_dir=results_dir)"
+    "#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=21"
    ]
   },
   {
    "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": {},
     "colab_type": "code",
-    "id": "8-hQC-AfiCKq",
-    "outputId": "f52e94a1-48bd-4b27-bc21-2306ee85ef9d"
+    "id": "ny3egicahAB4"
    },
    "outputs": [],
    "source": [
-    "!wget \"https://media.githubusercontent.com/media/jantic/DeOldify/master/test_images/TV1930s.jpg\" -O \"family_TV.jpg\""
+    "# 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/3e4dqky91h824ik/colorize_gen.pth\n",
+    "!mkdir \"models\"\n",
+    "download = drive.CreateFile({'id': '1ReNu8QCgi4vfyDovgYGMs7tlUGrnz72V'})\n",
+    "download.GetContentFile('./models/colorize_gen.pth')"
    ]
   },
   {
@@ -339,113 +215,71 @@
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/",
-     "height": 434
+     "height": 68
     },
     "colab_type": "code",
-    "id": "O6kfUN0GiJsq",
-    "outputId": "3e37f84d-6d27-4e4e-b7b6-9feae8be40ae"
+    "id": "NlOT9IlBHkk7",
+    "outputId": "1bb5cc4d-15a2-4174-a37c-9e11f3a2dce3"
    },
    "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:"
+    "from google.colab import drive\n",
+    "drive.mount('/content/drive')"
    ]
   },
   {
    "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"
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "!wget \"https://magnoliaforever.files.wordpress.com/2011/09/wizard-of-oz.jpg\" -O \"Dorothy.jpg\""
+    "!mkdir \"/content/drive/My Drive/deOldifyImages\"\n",
+    "!mkdir \"/content/drive/My Drive/deOldifyImages/results\""
    ]
   },
   {
    "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"
-   },
+   "metadata": {},
    "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."
+    "def colorize_gen_learner_exp(data:ImageDataBunch, gen_loss=FeatureLoss(), arch=models.resnet34):\n",
+    "    return unet_learner3(data, arch, wd=1e-3, blur=True, norm_type=NormType.Spectral,\n",
+    "                        self_attention=True, y_range=(-3.,3.), loss_func=gen_loss)"
    ]
   },
   {
    "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"
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "!wget \"https://i.ebayimg.com/images/g/HqkAAOSwRLZUAwyS/s-l300.jpg\" -O \"butch_and_sundance.jpg\""
+    "data = get_dummy_databunch()"
    ]
   },
   {
    "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"
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "vis.plot_transformed_image('butch_and_sundance.jpg', render_factor=29)"
+    "learn = colorize_gen_learner_exp(data=data)\n",
+    "#switch to read models from proper place\n",
+    "learn.path = Path('./')\n",
+    "learn.load(weights_name)\n",
+    "learn.model.eval()\n",
+    "filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor)\n",
+    "vis = ModelImageVisualizer(filtr, results_dir=results_dir)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {
     "colab_type": "text",
-    "id": "Ra5fs0UWn46V"
+    "id": "ZpCf0qbxicVK"
    },
    "source": [
-    "Let's get a picture of what they were actually wearing:"
+    "Here's an example of colorizing an image downloaded from the internet:"
    ]
   },
   {
@@ -457,12 +291,12 @@
      "height": 204
     },
     "colab_type": "code",
-    "id": "a8zoOYJanxkk",
-    "outputId": "2d85a505-93c7-4289-80cb-a23578b187b7"
+    "id": "8-hQC-AfiCKq",
+    "outputId": "f52e94a1-48bd-4b27-bc21-2306ee85ef9d"
    },
    "outputs": [],
    "source": [
-    "!wget \"https://bethanytompkins.files.wordpress.com/2015/09/freezeframe.jpg\" -O \"butch_and_sundance_color.jpg\""
+    "!wget \"https://media.githubusercontent.com/media/jantic/DeOldify/master/test_images/TV1930s.jpg\" -O \"family_TV.jpg\""
    ]
   },
   {
@@ -471,15 +305,15 @@
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/",
-     "height": 273
+     "height": 434
     },
     "colab_type": "code",
-    "id": "nEULX0Uan-pi",
-    "outputId": "6fe834de-863c-46f9-a740-df75cacb76e7"
+    "id": "O6kfUN0GiJsq",
+    "outputId": "3e37f84d-6d27-4e4e-b7b6-9feae8be40ae"
    },
    "outputs": [],
    "source": [
-    "Image('butch_and_sundance_color.jpg')"
+    "vis.plot_transformed_image('family_TV.jpg')"
    ]
   },
   {
@@ -489,7 +323,7 @@
     "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."
+    "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."
    ]
   },
   {
@@ -536,7 +370,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.5"
+   "version": "3.7.0"
   }
  },
  "nbformat": 4,