skeleton.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/clemsage/NeuralDocumentClassification/blob/master/skeleton.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": "j5UeZiGlvNoH"
      },
      "source": [
        "# Settting up the computing environment\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "n52k6VoU1brz"
      },
      "source": [
        "## Install and import TensorFlow 2.0 with GPU"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "zZVpUWUluZbV"
      },
      "source": [
        "Select \"GPU\" in the Accelerator drop-down on Notebook Settings through the Edit menu."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "colab": {},
        "colab_type": "code",
        "id": "8XYfp8LNcRD3"
      },
      "outputs": [],
      "source": [
        "# !pip install tensorflow-gpu==2.0\n",
        "import tensorflow as tf\n",
        "print (tf.__version__)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "nkd4MHFQv0jS"
      },
      "source": [
        "## Confirm TensorFlow can see the GPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "nbotnVwUpWBa"
      },
      "outputs": [],
      "source": [
        "device_name = tf.test.gpu_device_name()\n",
        "if device_name != '/device:GPU:0':\n",
        "  raise SystemError('GPU device not found')\n",
        "print('Found GPU at: {}'.format(device_name))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "IRPpWCXA05ka"
      },
      "source": [
        "## Additional information about hardware"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "k8XAtu4u1mXZ"
      },
      "source": [
        "For CPU information and RAM, run:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "4Mr3-8s-1jPB"
      },
      "outputs": [],
      "source": [
        "!cat /proc/cpuinfo\n",
        "!cat /proc/meminfo"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ic-CaNISucO-"
      },
      "source": [
        "## Other useful package imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "gevJulhruagf"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from collections import Counter\n",
        "import os\n",
        "import PIL\n",
        "import sys\n",
        "import importlib"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "0kDf1Kmntpwo"
      },
      "source": [
        "# Working on the dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "pQe3wu4U5kOt"
      },
      "source": [
        "The dataset is a subset of the [RVL-CDIP dataset](https://www.cs.cmu.edu/~aharley/rvl-cdip/). See [Harley et al.](http://scs.ryerson.ca/~aharley/icdar15/harley_convnet_icdar15.pdf) and [Asim et al.](https://www.dfki.de/fileadmin/user_upload/import/10637_Asim_Document_Image_Classification.pdf) papers for recent works on this dataset."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "jH2SxwR_1Rpu"
      },
      "source": [
        "## Information about the dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "lxvfM7Wn_YUQ"
      },
      "source": [
        "This project only considers the following 5 classes among the 16 classes of the original dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "colab": {},
        "colab_type": "code",
        "id": "lGLInZca1Pbg"
      },
      "outputs": [],
      "source": [
        "class_names = ['form', 'email', 'handwritten', 'advertisement', 'invoice']\n",
        "NUM_CLASSES = len(class_names)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "XgN4fpA0uO8n"
      },
      "source": [
        "## Import the dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "AVnJr5tzswrq"
      },
      "source": [
        "First, clone or pull the GitHub repository of the project:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ST3fUpSmqncY"
      },
      "outputs": [],
      "source": [
        "if not os.path.exists('NeuralDocumentClassification'):\n",
        "  !git clone https://github.com/clemsage/NeuralDocumentClassification.git\n",
        "else:\n",
        "  !git -C NeuralDocumentClassification pull\n",
        "sys.path.append('NeuralDocumentClassification')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "9zWbX_GYypxH"
      },
      "source": [
        "Download and extract labels, images and dataset assignments from this [Google Drive](https://drive.google.com/drive/folders/1Pkd6sUkDGBUymWKK93abZx1MQiWmzFgP):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "pQJ8Kqy3sv_v"
      },
      "outputs": [],
      "source": [
        "import download_dataset\n",
        "importlib.reload(download_dataset)\n",
        "for elt in ['label', 'image', 'dataset_assignment']:\n",
        "  download_dataset.download_and_extract(elt)\n",
        "dataset_path = 'dataset'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VYoAg4lOau3h"
      },
      "source": [
        "Parse `dataset_assignment.txt` to retrieve the training and test sets:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "a4G-8jVJauWC"
      },
      "outputs": [],
      "source": [
        "dataset = {\"training\": [], \"test\": []}\n",
        "with open(os.path.join(dataset_path, 'dataset_assignment.txt'), 'r') as f:\n",
        "  for line in f.readlines():\n",
        "    line = line.split('\\n')[0]\n",
        "    file_id, assignment = line.split(',')\n",
        "    file_path = os.path.join(dataset_path, 'image_png', '%s.png' % file_id)\n",
        "    dataset[assignment].append(file_path)\n",
        "\n",
        "print(\"Number of training documents: %d\" % len(dataset[\"training\"]))\n",
        "print(\"Number of test documents: %d\" % len(dataset['test']))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "1NRF5WhuiW3S"
      },
      "source": [
        "List the image files of the training and test dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "vWG346GnPSKY"
      },
      "outputs": [],
      "source": [
        "list_train_ds = tf.data.Dataset.from_tensor_slices(dataset['training'])\n",
        "list_train_ds = list_train_ds.shuffle(100000)\n",
        "list_test_ds = tf.data.Dataset.from_tensor_slices(dataset['test'])\n",
        "list_test_ds = list_test_ds.shuffle(100000)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "93EnbGGRid77"
      },
      "source": [
        "Print 5 image file names of the training set (see [TensorFlow tutorial for loading data](https://www.tensorflow.org/tutorials/load_data/images#load_using_tfdata)):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Lw4UmzI2_FV_"
      },
      "outputs": [],
      "source": [
        "### Insert your code here ###\n",
        "# See the expected solution by clicking on the cell below"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "xlg7AAuoiAPU"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "for f in list_train_ds.take(5):\n",
        "  print(f.numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "pTo64W_682WR"
      },
      "source": [
        "Get the labels for all files of the dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "e7QYsgFr88jO"
      },
      "outputs": [],
      "source": [
        "raw_class_indices = ['1', '2', '3', '4', '11']\n",
        "\n",
        "# Parse the labels.txt file to get all labels\n",
        "file_paths, labels = [], []\n",
        "with open(os.path.join(dataset_path, 'label.txt'), 'r') as f:\n",
        "  for line in f.readlines():\n",
        "    line = line.split('\\n')[0]\n",
        "    file_id, label = line.split(',')\n",
        "    file_path = os.path.join(dataset_path, 'image_png', '%s.png' % file_id)\n",
        "    file_paths.append(file_path)\n",
        "    labels.append(raw_class_indices.index(label))\n",
        "\n",
        "labels_idx = tf.lookup.StaticHashTable(\n",
        "    initializer=tf.lookup.KeyValueTensorInitializer(keys=file_paths, \n",
        "                                                    values=labels),\n",
        "    default_value=tf.constant(-1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Id3yJa4C-bqk"
      },
      "source": [
        "Implement a function that resizes images to the [US Letter format](https://en.wikipedia.org/wiki/Letter_(paper_size)) (8.5 by 11 inches) with 35 pixels by inch ([PPI](https://en.wikipedia.org/wiki/Pixel_density)):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "7YSxaSOaAR4W"
      },
      "outputs": [],
      "source": [
        "PPI = 35  # number of Pixels Per Inch \n",
        "IMG_WIDTH, IMG_HEIGHT = None, None  # define them as a function of PPI\n",
        "\n",
        "def decode_img(img):\n",
        "  # Adapt the function given in the tutorial : \n",
        "  # https://www.tensorflow.org/tutorials/load_data/images#load_using_tfdata \n",
        "  return img"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "rzZAayou-XYx"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "PPI = 35  # Number of Pixels Per Inch \n",
        "IMG_WIDTH, IMG_HEIGHT = int(PPI * 8.5), int(PPI * 11)\n",
        "\n",
        "def decode_img(img):\n",
        "  # convert the compressed string to a uint8 tensor\n",
        "  img = tf.io.decode_png(img, channels=1)\n",
        "  # convert to floats in the [0,1] range\n",
        "  img = tf.image.convert_image_dtype(img, tf.float32)\n",
        "  # resize the image to the desired size\n",
        "  img = tf.image.resize(img, [IMG_HEIGHT, IMG_WIDTH])\n",
        "  return img"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SIRu_CYsi06k"
      },
      "source": [
        "Implement a function to convert a file path to an (image_data, label) pair:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "9U_FnyBjzwbV"
      },
      "outputs": [],
      "source": [
        "def process_path(file_path):\n",
        "  label = labels_idx.lookup(file_path)\n",
        "\n",
        "  img = tf.io.read_file(file_path)\n",
        "  img = decode_img(img)\n",
        "  return img, label"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "PqyoYwtHH97_"
      },
      "source": [
        "Apply image and label retrieval to the training and test datasets:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "i5nmK9JFI-Hv"
      },
      "outputs": [],
      "source": [
        "labeled_train_ds = list_train_ds.map(process_path, \n",
        "                         num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "labeled_test_ds = list_test_ds.map(process_path, \n",
        "                         num_parallel_calls=tf.data.experimental.AUTOTUNE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "_1Eq6TC2wicn"
      },
      "source": [
        "## Explore the data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "MWJshtOvIpfo"
      },
      "source": [
        "Get image shape and label for one element of the training dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "NQ7bN2jc0rNK"
      },
      "outputs": [],
      "source": [
        "for image, label in labeled_train_ds.take(1):\n",
        "  print(\"Image shape (height, width, depth):\", image.numpy().shape)\n",
        "  print(\"Label:\", class_names[label.numpy()])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "XVm4YLGH1Y9_"
      },
      "source": [
        "Plot 3 random training images of each class:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "qwsI8QQkEUMV"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(30, 60))\n",
        "n_images_per_class = 3\n",
        "\n",
        "for image, label in labeled_train_ds:\n",
        "  break # Sample images and labels for each class\n",
        "\n",
        "for class_idx, class_name in enumerate(class_names):\n",
        "  for i in range(n_images_per_class):\n",
        "    plt.subplot(NUM_CLASSES, n_images_per_class, \n",
        "                class_idx*n_images_per_class + i + 1)\n",
        "    plt.xticks([])\n",
        "    plt.yticks([])\n",
        "    plt.grid(False)\n",
        "    # plt.imshow(...)\n",
        "    # plt.xlabel(...)\n",
        "\n",
        "#plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "ET7vBNIU15pQ"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "plt.figure(figsize=(30, 60))\n",
        "n_images_per_class = 3\n",
        "images = {class_name: [] for class_name in class_names}\n",
        "\n",
        "for image, label in labeled_train_ds:\n",
        "  image = image.numpy()\n",
        "  label = label.numpy()\n",
        "\n",
        "  if len(images[class_names[label]]) < n_images_per_class:\n",
        "    images[class_names[label]].append(image)\n",
        "  \n",
        "  if all([len(images[class_name]) == n_images_per_class \n",
        "          for class_name in class_names]):\n",
        "    break\n",
        "\n",
        "for class_idx, class_name in enumerate(class_names):\n",
        "  for i in range(n_images_per_class):\n",
        "    plt.subplot(NUM_CLASSES, n_images_per_class, \n",
        "                class_idx*n_images_per_class + i + 1)\n",
        "    plt.xticks([])\n",
        "    plt.yticks([])\n",
        "    plt.grid(False)\n",
        "    plt.imshow(np.squeeze(images[class_name][i]), cmap='gray')\n",
        "    plt.xlabel(class_name)\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gTiy3OjTzPm7"
      },
      "source": [
        "Print the class distribution in the training set:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "9XaIjc79II9E"
      },
      "outputs": [],
      "source": [
        "cnt_class = Counter()\n",
        "for file_path in dataset['training']:\n",
        "  label = labels_idx.lookup(tf.constant(file_path))\n",
        "  # Update the counter with the label value\n",
        "\n",
        "# Print the class counter"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "MjyrGh3szWs5"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "cnt_class = Counter()\n",
        "for file_path in dataset['training']: \n",
        "  label = labels_idx.lookup(tf.constant(file_path))\n",
        "  cnt_class.update([class_names[label.numpy()]])\n",
        "\n",
        "for key, val in cnt_class.most_common():\n",
        "  print('%s: %d' % (key, val))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "vECBpU3JCxyt"
      },
      "source": [
        "## Prepare training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kxJ3RtsuK3nu"
      },
      "source": [
        "Use a temporary folder for caching elements of the dataset in order to speed up training and testing:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Nx5JO952KSqj"
      },
      "outputs": [],
      "source": [
        "temp_folder = '/tmp/%dx%dx1' % (IMG_HEIGHT, IMG_WIDTH)\n",
        "labeled_train_ds = labeled_train_ds.cache(temp_folder)\n",
        "labeled_test_ds = labeled_test_ds.cache(temp_folder)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "aSD1NL7g3kSt"
      },
      "source": [
        "Shuffle the documents within each subset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "pYalPpQi3oeh"
      },
      "outputs": [],
      "source": [
        "labeled_train_ds = labeled_train_ds.shuffle(2048)\n",
        "labeled_test_ds = labeled_test_ds.shuffle(2048)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "eOXl8eGcJP_J"
      },
      "source": [
        "Batch documents within each subset:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ejTxwX3KJ4Gm"
      },
      "outputs": [],
      "source": [
        "batch_size = 128\n",
        "labeled_train_ds = labeled_train_ds.batch(batch_size)\n",
        "labeled_test_ds = labeled_test_ds.batch(batch_size)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-18lHiag4uiS"
      },
      "source": [
        "Prefetch the subsets in the background while the model is computing:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Je7kuYFh4y_2"
      },
      "outputs": [],
      "source": [
        "labeled_train_ds = labeled_train_ds.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)\n",
        "labeled_test_ds = labeled_test_ds.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "9oQZqkwA5mRU"
      },
      "source": [
        "# Visual classifiers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "JcRvGBPJCPAI"
      },
      "source": [
        "## Fully connected neural network"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hmD69yYP8Z7X"
      },
      "source": [
        "### Set up the layers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ooKloCG47hA4"
      },
      "source": [
        "Build a neural network composed of one fully connected (aka dense) hidden layer with 128 [ReLu](https://en.wikipedia.org/wiki/Rectifier_(neural_networks)) units and one output softmax layer.\n",
        "\n",
        "Each image must be reshaped to a 1 dimensional vector before being fed to the hidden layer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "YCIzEefKKKNK"
      },
      "outputs": [],
      "source": [
        "model = tf.keras.Sequential([\n",
        "  # Insert your layers here, see the following documentation:\n",
        "  # https://www.tensorflow.org/tutorials/quickstart/beginner\n",
        "])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "HYxVCSLQ5tmk"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "model = tf.keras.Sequential([\n",
        "    tf.keras.layers.Flatten(input_shape=(IMG_HEIGHT, IMG_WIDTH, 1)),\n",
        "    tf.keras.layers.Dense(128, activation='relu'),\n",
        "    tf.keras.layers.Dense(NUM_CLASSES, activation='softmax')\n",
        "])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "6JjOGshU8ftf"
      },
      "source": [
        "### Compile the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "pHGhRxVD8wl-"
      },
      "source": [
        "Compile the model by providing the optimizer, the loss function you want to minimize and the metrics to monitor during training:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Lhan11blCaW7"
      },
      "outputs": [],
      "source": [
        "model.compile(\n",
        "    optimizer='adam', # https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adam\n",
        "    loss='sparse_categorical_crossentropy', # Loss used for multi-class classification with integer labels\n",
        "    # https://www.tensorflow.org/api_docs/python/tf/keras/losses/SparseCategoricalCrossentropy\n",
        "    metrics=['sparse_categorical_accuracy'] # https://www.tensorflow.org/api_docs/python/tf/keras/metrics/Accuracy\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "qsnjxEZ8gJ1I"
      },
      "source": [
        "Print a summary of the model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "8UUZ51MmgUhL"
      },
      "outputs": [],
      "source": [
        "print(model.summary())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-GFuLNGh9U5w"
      },
      "source": [
        "### Train the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Ut4NKowCE8Hz"
      },
      "source": [
        "Fit the model on the training set for 20 epochs:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "vO1F5yxIMgY-"
      },
      "outputs": [],
      "source": [
        "EPOCHS = 20\n",
        "# model.fit(...)  # https://www.tensorflow.org/api_docs/python/tf/keras/Model#fit"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "4nT6h3y_9XFo"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "EPOCHS = 10\n",
        "model.fit(labeled_train_ds, epochs=EPOCHS)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "nfLk6pw2M08a"
      },
      "source": [
        "### Evaluation on the test set"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "HNIV9G61P3SB"
      },
      "source": [
        "Get the values of the loss and accuracy: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "vGOSP20hqZ68"
      },
      "outputs": [],
      "source": [
        "# model.evaluate(..., verbose=2) # https://www.tensorflow.org/api_docs/python/tf/keras/Model#evaluate"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "wq93ddH0NDpS"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "model.evaluate(labeled_test_ds, verbose=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "x_gKCPtWVg3M"
      },
      "source": [
        "Are these values different from their training counterparts ?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "XXHM8y9kRW5V"
      },
      "source": [
        "### Prediction on the test set"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "5xHWavjURp-t"
      },
      "source": [
        "Implement a function that gathers the model predictions and the ground truth labels for a random batch of a given dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "hq68PQDJsbXh"
      },
      "outputs": [],
      "source": [
        "def predict_random_batch(model, dataset):\n",
        "  \"\"\"\n",
        "  Sample a random batch of the dataset and return the images of this batch \n",
        "  as well as its labels and the predicted classes of a given model\n",
        "\n",
        "  Parameters\n",
        "  ----------\n",
        "  model : tf.keras.Model\n",
        "  dataset: tf.data.Dataset\n",
        "\n",
        "  Returns\n",
        "  -------\n",
        "  images: np.ndarray or EagerTensor\n",
        "  labels: list of str\n",
        "  predicted_classes: list of str\n",
        "  \"\"\"\n",
        "  images, labels = next(iter(dataset))\n",
        "  \n",
        "  # get label names for the sampled batch\n",
        "\n",
        "  # make predictions\n",
        "  predicted_classes = None\n",
        "\n",
        "  return images, labels, predicted_classes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "xX0MIEONRkp7"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "def predict_random_batch(model, dataset):\n",
        "  \"\"\"\n",
        "  Sample a random batch of the dataset and return the images of this batch \n",
        "  as well as its labels and the predicted classes of a given model\n",
        "\n",
        "  Parameters\n",
        "  ----------\n",
        "  model : tf.keras.Model\n",
        "  dataset: tf.data.Dataset\n",
        "\n",
        "  Returns\n",
        "  -------\n",
        "  images: np.ndarray or EagerTensor\n",
        "  labels: list of str\n",
        "  predicted_classes: list of str\n",
        "  \"\"\"\n",
        "  images, labels = next(iter(dataset))\n",
        "  \n",
        "  # get label names for the sampled batch\n",
        "  labels = [class_names[i] for i in labels]\n",
        "\n",
        "  # make predictions\n",
        "  predictions = model.predict(images)\n",
        "  predicted_classes_idx = np.argmax(predictions, axis=1)\n",
        "  predicted_classes = [class_names[i] for i in predicted_classes_idx]\n",
        "\n",
        "  return images, labels, predicted_classes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "uyZFGQLLVFeC"
      },
      "source": [
        "Plot the first 9 images of this batch, give their labels and predicted classes in the legend:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "wrgwVDo4VFG7"
      },
      "outputs": [],
      "source": [
        "def plot_images_predictions_and_labels(images, labels, predicted_classes):\n",
        "  plt.figure(figsize=(30, 40))\n",
        "\n",
        "  for im_idx in range(9):\n",
        "    plt.subplot(3, 3, im_idx + 1)\n",
        "    plt.xticks([])\n",
        "    plt.yticks([])\n",
        "    plt.grid(False)\n",
        "    plt.imshow(np.squeeze(images[im_idx]), cmap='gray')\n",
        "    plt.xlabel(\"label: %s\\npred: %s\" % (labels[im_idx], \n",
        "                                        predicted_classes[im_idx]))\n",
        "\n",
        "  plt.show()\n",
        "\n",
        "result = predict_random_batch(model, labeled_test_ds)\n",
        "plot_images_predictions_and_labels(*result)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "monzP4ZVVkGr"
      },
      "source": [
        "### Under the Hood"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "61YBF5sLfsrZ"
      },
      "source": [
        "Implement an ReLu dense layer by creating its weights and biases and giving the transformation from inputs to outputs:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "8_-Q8GwoChTv"
      },
      "outputs": [],
      "source": [
        "# https://www.tensorflow.org/guide/keras/custom_layers_and_models#the_layer_class\n",
        "class MyDenseLayer(tf.keras.layers.Layer):\n",
        "\n",
        "  def __init__(self, units, input_dim):\n",
        "    super(MyDenseLayer, self).__init__()\n",
        "    self.w = self.add_weight(\n",
        "        shape=None,  ## Insert the shape of the weight matrix here\n",
        "        initializer='glorot_uniform',  # Default initializer for weights of a tf.keras.layers.Dense layer\n",
        "        trainable=True)\n",
        "\n",
        "    self.b = self.add_weight(\n",
        "        shape=None, ## Insert the shape of the bias vector here\n",
        "        initializer='zeros',  # Default initializer for biases of a tf.keras.layers.Dense layer\n",
        "        trainable=True)\n",
        "\n",
        "  def call(self, inputs):\n",
        "    outputs = None\n",
        "    return outputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "5HWRDQkPfBam"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "# https://www.tensorflow.org/guide/keras/custom_layers_and_models#the_layer_class\n",
        "class MyDenseLayer(tf.keras.layers.Layer):\n",
        "\n",
        "  def __init__(self, units, input_dim):\n",
        "    super(MyDenseLayer, self).__init__()\n",
        "    self.w = self.add_weight(\n",
        "        shape=(input_dim, units),\n",
        "        initializer='glorot_uniform',  # Default initializer for weights of a tf.keras.layers.Dense layer\n",
        "        trainable=True)\n",
        "\n",
        "    self.b = self.add_weight(\n",
        "        shape=(units,),\n",
        "        initializer='zeros',  # Default initializer for biases of a tf.keras.layers.Dense layer\n",
        "        trainable=True)\n",
        "\n",
        "  def call(self, inputs):\n",
        "    return tf.keras.activations.relu(tf.matmul(inputs, self.w) + self.b)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VyPOuiP_5ZGv"
      },
      "source": [
        "Using your custom hidden layer, set up again the layers of the model defined previously:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "8fK9p35mD9eI"
      },
      "outputs": [],
      "source": [
        "model = tf.keras.Sequential([\n",
        "    tf.keras.layers.Flatten(input_shape=(IMG_HEIGHT, IMG_WIDTH, 1)),\n",
        "    ## Insert your custom hidden layer here\n",
        "    tf.keras.layers.Dense(NUM_CLASSES, activation='softmax')\n",
        "])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "CBd68qImh1zH"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "model = tf.keras.Sequential([\n",
        "    tf.keras.layers.Flatten(input_shape=(IMG_HEIGHT, IMG_WIDTH, 1)),\n",
        "    MyDenseLayer(128, IMG_HEIGHT*IMG_WIDTH),\n",
        "    tf.keras.layers.Dense(NUM_CLASSES, activation='softmax')\n",
        "])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "eVgtyKeNhXJR"
      },
      "source": [
        "Lower-level implementation of the model compile step:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "-4IEM4Rhg4ji"
      },
      "outputs": [],
      "source": [
        "loss_object = tf.keras.losses.SparseCategoricalCrossentropy()\n",
        "optimizer = tf.keras.optimizers.Adam()\n",
        "train_loss = tf.keras.metrics.Mean(name='train_loss')\n",
        "train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')\n",
        "\n",
        "@tf.function\n",
        "def train_step(images, labels):\n",
        "  with tf.GradientTape() as tape:\n",
        "    predictions = model(images)\n",
        "    loss = loss_object(labels, predictions)\n",
        "  gradients = tape.gradient(loss, model.trainable_variables)\n",
        "  optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
        "\n",
        "  train_loss(loss)\n",
        "  train_accuracy(labels, predictions)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VxMO3ZB7ic0-"
      },
      "source": [
        "Lower-level implementation of the model fit step:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Mo77GEbdifjm"
      },
      "outputs": [],
      "source": [
        "for epoch in range(EPOCHS):\n",
        "  for images, labels in labeled_train_ds:\n",
        "    train_step(images, labels)\n",
        "    template = 'Epoch {}, Loss: {}, Accuracy: {}'\n",
        "    print(template.format(epoch+1,\n",
        "                          train_loss.result(),\n",
        "                          train_accuracy.result()*100))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "LnMkYilGngXs"
      },
      "source": [
        "## Convolutional Neural Networks (CNN)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Kaz6G-cbotNS"
      },
      "source": [
        "### Training from scratch"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gNifb_-ImoL1"
      },
      "source": [
        "Create and compile a model alterning convolution and max pooling layers. You can add some fully connected layers between the last locally connected layer and the output layer. Start with a shallow network (4 or 5 convolution layers) and progressively move to deeper architectures: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "5VSo8vjiGOUc"
      },
      "outputs": [],
      "source": [
        "from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense\n",
        "model = tf.keras.Sequential([\n",
        "    # Alterning Conv2D and MaxPooling2D layers\n",
        "\n",
        "    # Some dense hidden layer(s)\n",
        "\n",
        "    Dense(NUM_CLASSES, activation='softmax')\n",
        "])\n",
        "\n",
        "print(model.summary())\n",
        "\n",
        "model.compile(optimizer='adam',\n",
        "              loss='sparse_categorical_crossentropy',\n",
        "              metrics=['accuracy'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "R-fLDDzLmjMH"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense\n",
        "shallow_model = tf.keras.Sequential([\n",
        "    Conv2D(16, 3, padding='same', activation='relu', \n",
        "           input_shape=(IMG_HEIGHT, IMG_WIDTH, 1)),\n",
        "    MaxPooling2D(4),\n",
        "    Conv2D(32, 3, padding='same', activation='relu'),\n",
        "    MaxPooling2D(4),\n",
        "    Conv2D(64, 3, padding='same', activation='relu'),\n",
        "    MaxPooling2D(4),\n",
        "    Conv2D(128, 3, padding='same', activation='relu'),\n",
        "    MaxPooling2D(4),\n",
        "    Flatten(),\n",
        "    Dense(128, activation='relu'),\n",
        "    Dense(NUM_CLASSES, activation='softmax')\n",
        "])\n",
        "\n",
        "deep_model = tf.keras.Sequential([\n",
        "    Conv2D(16, 3, padding='same', activation='relu', \n",
        "           input_shape=(IMG_HEIGHT, IMG_WIDTH, 1)),\n",
        "    MaxPooling2D(),\n",
        "    Conv2D(32, 3, padding='same', activation='relu'),\n",
        "    MaxPooling2D(),\n",
        "    Conv2D(64, 3, padding='same', activation='relu'),\n",
        "    MaxPooling2D(),\n",
        "    Conv2D(128, 3, padding='same', activation='relu'),\n",
        "    MaxPooling2D(),\n",
        "    Conv2D(256, 3, padding='same', activation='relu'),\n",
        "    MaxPooling2D(),\n",
        "    Conv2D(256, 3, padding='same', activation='relu'),\n",
        "    MaxPooling2D(),\n",
        "    Conv2D(256, 3, padding='same', activation='relu'),\n",
        "    MaxPooling2D(),\n",
        "    Conv2D(256, 3, padding='same', activation='relu'),\n",
        "    MaxPooling2D(),\n",
        "    Flatten(),\n",
        "    Dense(256, activation='relu'),\n",
        "    Dense(NUM_CLASSES, activation='softmax')\n",
        "])\n",
        "\n",
        "model_with_strides = tf.keras.Sequential([\n",
        "    Conv2D(16, 3, padding='same', activation='relu', \n",
        "           input_shape=(IMG_HEIGHT, IMG_WIDTH, 1)),\n",
        "    MaxPooling2D(),\n",
        "    Conv2D(32, 3, padding='same', activation='relu'),\n",
        "    MaxPooling2D(),\n",
        "    Conv2D(64, 3, padding='same', activation='relu', strides=2),\n",
        "    MaxPooling2D(),\n",
        "    Conv2D(128, 3, padding='same', activation='relu', strides=2),\n",
        "    MaxPooling2D(),\n",
        "    Conv2D(128, 3, padding='same', activation='relu', strides=2),\n",
        "    MaxPooling2D(),\n",
        "    Flatten(),\n",
        "    Dense(128, activation='relu'),\n",
        "    Dense(NUM_CLASSES, activation='softmax')\n",
        "])\n",
        "model = deep_model\n",
        "print(model.summary())\n",
        "\n",
        "model.compile(optimizer='adam',\n",
        "              loss='sparse_categorical_crossentropy',\n",
        "              metrics=['sparse_categorical_accuracy'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "FVzSjxP0qP4D"
      },
      "source": [
        "Fit the CNN on the training data:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "3jWd_WMpqPi8"
      },
      "outputs": [],
      "source": [
        "EPOCHS = 10\n",
        "model.fit(labeled_train_ds, epochs=EPOCHS)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "9mGpqBc7rnFv"
      },
      "source": [
        "Evaluate the trained model on the test set:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "tSmadMoVrmnf"
      },
      "outputs": [],
      "source": [
        "model.evaluate(labeled_test_ds, verbose=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "xK3wwRUPOBEs"
      },
      "source": [
        "You should reach test accuracy greater than 0.99 !\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "7S7XNOyD3pRI"
      },
      "source": [
        "Plot images, predictions and labels for some test documents:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "FZ4le3PQuWYo"
      },
      "outputs": [],
      "source": [
        "plot_images_predictions_and_labels(*predict_random_batch(model, labeled_test_ds))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ArNeT3D1nINR"
      },
      "source": [
        "### Transfer Learning with pre-trained models "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "PS-BCUYYRr_I"
      },
      "source": [
        "The objective is to leverage the knowledge learnt by a pre-trained image classifier. See [TensorFlow Hub](https://tfhub.dev/) to browse available state-of-the art models such as [Inception V3](https://www.cv-foundation.org/openaccess/content_cvpr_2016/papers/Szegedy_Rethinking_the_Inception_CVPR_2016_paper.pdf) or [MobileNet V2](https://arxiv.org/pdf/1801.04381.pdf)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "_Q7LRLglIQMY"
      },
      "source": [
        "Choose a pre-trained model for extracting high level feature vectors of document images:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "5F0WHgxTJX0S"
      },
      "outputs": [],
      "source": [
        "extractor_model = 'inception_v3'\n",
        "if extractor_model == 'inception_v3':\n",
        "  feature_extraction_url = \"https://tfhub.dev/google/tf2-preview/inception_v3/feature_vector/4\"\n",
        "  IMG_HEIGHT, IMG_WIDTH = None, None  ## Insert expected input image shape here\n",
        "elif extractor_model == 'mobilenet_v2':\n",
        "  feature_extraction_url = \"https://tfhub.dev/google/imagenet/mobilenet_v2_100_224/feature_vector/4\"\n",
        "  IMG_HEIGHT, IMG_WIDTH = None, None  ## Insert expected input image shape here"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "ho18UyThppL6"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "extractor_model = 'inception_v3'\n",
        "if extractor_model == 'inception_v3':\n",
        "  feature_extraction_url = \"https://tfhub.dev/google/tf2-preview/inception_v3/feature_vector/4\"\n",
        "  IMG_HEIGHT, IMG_WIDTH = 299, 299\n",
        "elif extractor_model == 'mobilenet_v2':\n",
        "  feature_extraction_url = \"https://tfhub.dev/google/imagenet/mobilenet_v2_100_224/feature_vector/4\"\n",
        "  IMG_HEIGHT, IMG_WIDTH = 224, 224"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "WS2LwuZGPfg_"
      },
      "source": [
        "Reshape images to the format expected by the chosen model, i.e. IMG_HEIGHT x IMG_WIDTH x 3 (RGB) and recreate training dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "D51vnWCPPunJ"
      },
      "outputs": [],
      "source": [
        "def decode_img(img):\n",
        "  # convert the compressed string to a uint8 tensor\n",
        "  img = tf.io.decode_png(img, channels=1)\n",
        "  # convert to floats in the [0,1] range\n",
        "  img = tf.image.convert_image_dtype(img, tf.float32)\n",
        "  # resize the image to the desired size\n",
        "  img = tf.image.resize(img, [IMG_HEIGHT, IMG_WIDTH])\n",
        "  # convert to RGB color scale\n",
        "  img = tf.concat([img for _ in range(3)], axis=-1)  # R = G = B\n",
        "  return img\n",
        "\n",
        "def process_path(file_path):\n",
        "  label = labels_idx.lookup(file_path)\n",
        "\n",
        "  img = tf.io.read_file(file_path)\n",
        "  img = decode_img(img)\n",
        "  return img, label\n",
        "\n",
        "labeled_train_ds = list_train_ds.map(process_path, \n",
        "                         num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "labeled_train_ds = labeled_train_ds.cache('/tmp/%dx%dx3' % (IMG_HEIGHT, \n",
        "                                                            IMG_WIDTH))\n",
        "labeled_train_ds = labeled_train_ds.shuffle(2048).batch(batch_size)\n",
        "labeled_train_ds = labeled_train_ds.prefetch(buffer_size=\n",
        "                                             tf.data.experimental.AUTOTUNE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "1WXZOzEvnINO"
      },
      "source": [
        "Construct our own image classifier by retrieving and freezing the hidden layers of the pre-trained model: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "JRq-5IVoMQFD"
      },
      "outputs": [],
      "source": [
        "import tensorflow_hub as hub\n",
        "# feature_extraction_layer = hub.KerasLayer(...)\n",
        "# https://www.tensorflow.org/tutorials/images/transfer_learning_with_hub\n",
        "\n",
        "model = tf.keras.Sequential([     \n",
        "    feature_extraction_layer,\n",
        "\n",
        "    # Some dense hidden layer(s)\n",
        "\n",
        "    Dense(NUM_CLASSES, activation='softmax')\n",
        "    ])\n",
        "\n",
        "model.compile(optimizer='adam',\n",
        "              loss='sparse_categorical_crossentropy',\n",
        "              metrics=['accuracy'])\n",
        "\n",
        "print(model.summary())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "aDMImRlEnINF"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "import tensorflow_hub as hub\n",
        "feature_extraction_layer = hub.KerasLayer(feature_extraction_url, \n",
        "                                          trainable=False,\n",
        "                                          input_shape=(IMG_HEIGHT, IMG_WIDTH, 3))\n",
        "model = tf.keras.Sequential([          \n",
        "    feature_extraction_layer,\n",
        "    Dense(128, activation='relu'),\n",
        "    Dense(128, activation='relu'),\n",
        "    Dense(NUM_CLASSES, activation='softmax')\n",
        "    ])\n",
        "\n",
        "model.compile(optimizer='adam',\n",
        "              loss='sparse_categorical_crossentropy',\n",
        "              metrics=['sparse_categorical_accuracy'])\n",
        "\n",
        "print(model.summary())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "m0fx9gcSRaKf"
      },
      "source": [
        "Train this new model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "gecTChcfRc-K"
      },
      "outputs": [],
      "source": [
        "EPOCHS = 10\n",
        "model.fit(labeled_train_ds, epochs=EPOCHS)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "j5UeZiGlvNoH",
        "n52k6VoU1brz",
        "nkd4MHFQv0jS",
        "IRPpWCXA05ka",
        "ic-CaNISucO-",
        "XgN4fpA0uO8n",
        "_1Eq6TC2wicn",
        "OdGw-l6TEUiP",
        "nfLk6pw2M08a",
        "XXHM8y9kRW5V"
      ],
      "include_colab_link": true,
      "name": "skeleton.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.9.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}