{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Programming Assignment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data pipeline with Keras and tf.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instructions\n",
    "\n",
    "In this notebook, you will implement a data processing pipeline using tools from both Keras and the tf.data module. You will use the `ImageDataGenerator` class in the tf.keras module to feed a network with training and test images from a local directory containing a subset of the LSUN dataset, and train the model both with and without data augmentation. You will then use the `map` and `filter` functions of the `Dataset` class with the CIFAR-100 dataset to train a network to classify a processed subset of the images.\n",
    "\n",
    "Some code cells are provided you in the notebook. You should avoid editing provided code, and make sure to execute the cells in order to avoid unexpected errors. Some cells begin with the line:\n",
    "\n",
    "`#### GRADED CELL ####`\n",
    "\n",
    "Don't move or edit this first line - this is what the automatic grader looks for to recognise graded cells. These cells require you to write your own code to complete them, and are automatically graded when you submit the notebook. Don't edit the function name or signature provided in these cells, otherwise the automatic grader might not function properly. Inside these graded cells, you can use any functions or classes that are imported below, but make sure you don't use any variables that are outside the scope of the function.\n",
    "\n",
    "### How to submit\n",
    "\n",
    "Complete all the tasks you are asked for in the worksheet. When you have finished and are happy with your code, press the **Submit Assignment** button at the top of this notebook.\n",
    "\n",
    "### Let's get started!\n",
    "\n",
    "We'll start running some imports, and loading the dataset. Do not edit the existing imports in the following cell. If you would like to make further Tensorflow imports, you should add them here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### PACKAGE IMPORTS ####\n",
    "\n",
    "# Run this cell first to import all required packages. Do not make any imports elsewhere in the notebook\n",
    "\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras.datasets import cifar100\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import json\n",
    "%matplotlib inline\n",
    "\n",
    "# If you would like to make further imports from tensorflow, add them here\n",
    "from tensorflow.keras.preprocessing.image import ImageDataGenerator\n",
    "\n",
    "from tensorflow.keras.layers import Conv2D,MaxPooling2D,Dense,Flatten,Input\n",
    "from tensorflow.keras.models import Model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Part 1: tf.keras\n",
    "<table><tr>\n",
    "<td> <img src=\"data/lsun/church.png\" alt=\"Church\" style=\"height: 210px;\"/>  </td>\n",
    "<td> <img src=\"data/lsun/classroom.png\" alt=\"Classroom\" style=\"height: 210px;\"/> </td>\n",
    "    <td> <img src=\"data/lsun/conference_room.png\" alt=\"Conference Room\" style=\"height: 210px;\"/> </td>\n",
    "</tr></table>\n",
    "  \n",
    "#### The LSUN Dataset\n",
    "\n",
    "In the first part of this assignment, you will use a subset of the [LSUN dataset](https://www.yf.io/p/lsun). This is a large-scale image dataset with 10 scene and 20 object categories. A subset of the LSUN dataset has been provided, and has already been split into training and test sets. The three classes included in the subset are `church_outdoor`, `classroom` and `conference_room`.\n",
    "\n",
    "* F. Yu, A. Seff, Y. Zhang, S. Song, T. Funkhouser and J. Xia. \"LSUN: Construction of a Large-scale Image Dataset using Deep Learning with Humans in the Loop\". arXiv:1506.03365, 10 Jun 2015 \n",
    "\n",
    "Your goal is to use the Keras preprocessing tools to construct a data ingestion and augmentation pipeline to train a neural network to classify the images into the three classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the directory locations for the training, validation and test sets\n",
    "\n",
    "train_dir = 'data/lsun/train'\n",
    "valid_dir = 'data/lsun/valid'\n",
    "test_dir = 'data/lsun/test'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create a data generator using the ImageDataGenerator class"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should first write a function that creates an `ImageDataGenerator` object, which rescales the image pixel values by a factor of 1/255."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### GRADED CELL ####\n",
    "\n",
    "# Complete the following function. \n",
    "# Make sure to not change the function name or arguments.\n",
    "\n",
    "def get_ImageDataGenerator():\n",
    "    \"\"\"\n",
    "    This function should return an instance of the ImageDataGenerator class.\n",
    "    This instance should be set up to rescale the data with the above scaling factor.\n",
    "    \"\"\"\n",
    "    imagedatagenerator = ImageDataGenerator(rescale =(1/255.))\n",
    "    \n",
    "    return imagedatagenerator\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Call the function to get an ImageDataGenerator as specified\n",
    "\n",
    "image_gen = get_ImageDataGenerator()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should now write a function that returns a generator object that will yield batches of images and labels from the training and test set directories. The generators should:\n",
    "\n",
    "* Generate batches of size 20.\n",
    "* Resize the images to 64 x 64 x 3.\n",
    "* Return one-hot vectors for labels. These should be encoded as follows:\n",
    "    * `classroom` $\\rightarrow$ `[1., 0., 0.]`\n",
    "    * `conference_room` $\\rightarrow$ `[0., 1., 0.]`\n",
    "    * `church_outdoor` $\\rightarrow$ `[0., 0., 1.]`\n",
    "* Pass in an optional random `seed` for shuffling (this should be passed into the `flow_from_directory` method).\n",
    "\n",
    "**Hint:** you may need to refer to the [documentation](https://keras.io/preprocessing/image/#imagedatagenerator-class) for the `ImageDataGenerator`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### GRADED CELL ####\n",
    "\n",
    "# Complete the following function.\n",
    "# Make sure not to change the function name or arguments.\n",
    "\n",
    "def get_generator(image_data_generator, directory, seed=None):\n",
    "    \"\"\"\n",
    "    This function takes an ImageDataGenerator object in the first argument and a \n",
    "    directory path in the second argument.\n",
    "    It should use the ImageDataGenerator to return a generator object according \n",
    "    to the above specifications. \n",
    "    The seed argument should be passed to the flow_from_directory method.\n",
    "    \n",
    "    \"\"\"\n",
    " 
    "    image_data_gen = image_data_generator.flow_from_directory(\n",
    "        directory = directory,\n",
    "        classes=['classroom','conference_room','church_outdoor'],\n",
    "        batch_size = 20,\n",
    "        target_size = (64,64),\n",
    "        seed=seed,\n",
    "        class_mode = \"categorical\"\n",
    "        )\n",
    "    return image_data_gen\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 300 images belonging to 3 classes.\n",
      "Found 120 images belonging to 3 classes.\n"
     ]
    }
   ],
   "source": [
    "# Run this cell to define training and validation generators\n",
    "\n",
    "train_generator = get_generator(image_gen, train_dir)\n",
    "valid_generator = get_generator(image_gen, valid_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are using a small subset of the dataset for demonstrative purposes in this assignment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Display sample images and labels from the training set\n",
    "\n",
    "The following cell depends on your function `get_generator` to be implemented correctly. If it raises an error, go back and check the function specifications carefully."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Display a few images and labels from the training set\n",
    "\n",
    "batch = next(train_generator)\n",
    "batch_images = np.array(batch[0])\n",
    "batch_labels = np.array(batch[1])\n",
    "lsun_classes = ['classroom', 'conference_room', 'church_outdoor']\n",
    "\n",
    "plt.figure(figsize=(16,10))\n",
    "for i in range(20):\n",
    "    ax = plt.subplot(4, 5, i+1)\n",
    "    plt.imshow(batch_images[i])\n",
    "    plt.title(lsun_classes[np.where(batch_labels[i] == 1.)[0][0]])\n",
    "    plt.axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 300 images belonging to 3 classes.\n"
     ]
    }
   ],
   "source": [
    "# Reset the training generator\n",
    "\n",
    "train_generator = get_generator(image_gen, train_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Build the neural network model\n",
    "\n",
    "You will now build and compile a convolutional neural network classifier. Using the functional API, build your model according to the following specifications:\n",
    "\n",
    "* The model should use the `input_shape` in the function argument to define the Input layer.\n",
    "* The first hidden layer should be a Conv2D layer with 8 filters, a 8x8 kernel size.\n",
    "* The second hidden layer should be a MaxPooling2D layer with a 2x2 pooling window size.\n",
    "* The third hidden layer should be a Conv2D layer with 4 filters, a 4x4 kernel size.\n",
    "* The fourth hidden layer should be a MaxPooling2D layer with a 2x2 pooling window size.\n",
    "* This should be followed by a Flatten layer, and then a Dense layer with 16 units and ReLU activation.\n",
    "* The final layer should be a Dense layer with 3 units and softmax activation.\n",
    "* All Conv2D layers should use `\"SAME\"` padding and a ReLU activation function.\n",
    "\n",
    "In total, the network should have 8 layers. The model should then be compiled with the Adam optimizer with learning rate 0.0005, categorical cross entropy loss, and categorical accuracy metric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### GRADED CELL ####\n",
    "\n",
    "# Complete the following function.\n",
    "# Make sure not to change the function name or arguments.\n",
    "\n",
    "def get_model(input_shape):\n",
    "    \"\"\"\n",
    "    This function should build and compile a CNN model according to the above specification,\n",
    "    using the functional API. Your function should return the model.\n",
    "    \"\"\"\n",
    "    inputs = Input(input_shape)\n",
    "    h = Conv2D(8,(8,8),padding = \"SAME\")(inputs)\n",
    "    h = MaxPooling2D((2,2))(h)\n",
    "    h = Conv2D(4,(4,4),padding = \"SAME\")(h)\n",
    "    h = MaxPooling2D((2,2))(h)\n",
    "    h = Flatten()(h)\n",
    "    h = Dense(16, activation = \"relu\")(h)\n",
    "    outputs = Dense(3, activation = \"softmax\")(h)\n",
    "    \n",
    "    model = Model(inputs = inputs,outputs = outputs)\n",
    "    \n",
    "    model.compile(\n",
    "        optimizer = tf.keras.optimizers.Adam(learning_rate = 0.0005),\n",
    "        loss = \"categorical_crossentropy\",\n",
    "        metrics = ['categorical_accuracy']\n",
    "                 )\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model_1\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "input_2 (InputLayer)         [(None, 64, 64, 3)]       0         \n",
      "_________________________________________________________________\n",
      "conv2d_2 (Conv2D)            (None, 64, 64, 8)         1544      \n",
      "_________________________________________________________________\n",
      "max_pooling2d_2 (MaxPooling2 (None, 32, 32, 8)         0         \n",
      "_________________________________________________________________\n",
      "conv2d_3 (Conv2D)            (None, 32, 32, 4)         516       \n",
      "_________________________________________________________________\n",
      "max_pooling2d_3 (MaxPooling2 (None, 16, 16, 4)         0         \n",
      "_________________________________________________________________\n",
      "flatten_1 (Flatten)          (None, 1024)              0         \n",
      "_________________________________________________________________\n",
      "dense_2 (Dense)              (None, 16)                16400     \n",
      "_________________________________________________________________\n",
      "dense_3 (Dense)              (None, 3)                 51        \n",
      "=================================================================\n",
      "Total params: 18,511\n",
      "Trainable params: 18,511\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "# Build and compile the model, print the model summary\n",
    "\n",
    "lsun_model = get_model((64, 64, 3))\n",
    "lsun_model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Train the neural network model\n",
    "\n",
    "You should now write a function to train the model for a specified number of epochs (specified in the `epochs` argument). The function takes a `model` argument, as well as `train_gen` and `valid_gen` arguments for the training and validation generators respectively, which you should use for training and validation data in the training run. You should also use the following callbacks:\n",
    "\n",
    "* An `EarlyStopping` callback that monitors the validation accuracy and has patience set to 10. \n",
    "* A `ReduceLROnPlateau` callback that monitors the validation loss and has the factor set to 0.5 and minimum learning set to 0.0001\n",
    "\n",
    "Your function should return the training history."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#### GRADED CELL ####\n",
    "\n",
    "# Complete the following function.\n",
    "# Make sure not to change the function name or arguments.\n",
    "\n",
    "def train_model(model, train_gen, valid_gen, epochs):\n",
    "    \"\"\"\n",
    "    This function should define the callback objects specified above, and then use the\n",
    "    train_gen and valid_gen generator object arguments to train the model for the (maximum) \n",
    "    number of epochs specified in the function argument, using the defined callbacks.\n",
    "    The function should return the training history.\n",
    "    \"\"\"\n",
    "    history = model.fit(train_gen,\n",
    "               validation_data = valid_gen,\n",
    "               epochs = epochs,\n",
    "               callbacks = [tf.keras.callbacks.EarlyStopping(patience = 10),tf.keras.callbacks.ReduceLROnPlateau(factor = 0.5,min_delta = 0.0001)]\n",
    "              )\n",
    "    \n",
    "    return history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train for 15 steps, validate for 6 steps\n",
      "Epoch 1/50\n",
      "15/15 [==============================] - 11s 743ms/step - loss: 1.0774 - categorical_accuracy: 0.3967 - val_loss: 0.9810 - val_categorical_accuracy: 0.5667\n",
      "Epoch 2/50\n",
      "15/15 [==============================] - 10s 647ms/step - loss: 0.9670 - categorical_accuracy: 0.5367 - val_loss: 0.8666 - val_categorical_accuracy: 0.6000\n",
      "Epoch 3/50\n",
      "15/15 [==============================] - 9s 633ms/step - loss: 0.8358 - categorical_accuracy: 0.6100 - val_loss: 0.8097 - val_categorical_accuracy: 0.6667\n",
      "Epoch 4/50\n",
      "15/15 [==============================] - 9s 633ms/step - loss: 0.7275 - categorical_accuracy: 0.6867 - val_loss: 0.7484 - val_categorical_accuracy: 0.6583\n",
      "Epoch 5/50\n",
      "15/15 [==============================] - 10s 653ms/step - loss: 0.6560 - categorical_accuracy: 0.7033 - val_loss: 0.7666 - val_categorical_accuracy: 0.6417\n",
      "Epoch 6/50\n",
      "15/15 [==============================] - 10s 673ms/step - loss: 0.6627 - categorical_accuracy: 0.7133 - val_loss: 0.7740 - val_categorical_accuracy: 0.6750\n",
      "Epoch 7/50\n",
      "15/15 [==============================] - 10s 640ms/step - loss: 0.6027 - categorical_accuracy: 0.7333 - val_loss: 0.8330 - val_categorical_accuracy: 0.6000\n",
      "Epoch 8/50\n",
      "15/15 [==============================] - 9s 633ms/step - loss: 0.6221 - categorical_accuracy: 0.7500 - val_loss: 0.7647 - val_categorical_accuracy: 0.6917\n",
      "Epoch 9/50\n",
      "15/15 [==============================] - 9s 627ms/step - loss: 0.6372 - categorical_accuracy: 0.7133 - val_loss: 0.8383 - val_categorical_accuracy: 0.6417\n",
      "Epoch 10/50\n",
      "15/15 [==============================] - 9s 625ms/step - loss: 0.6037 - categorical_accuracy: 0.7367 - val_loss: 0.7184 - val_categorical_accuracy: 0.7083\n",
      "Epoch 11/50\n",
      "15/15 [==============================] - 10s 634ms/step - loss: 0.5290 - categorical_accuracy: 0.8000 - val_loss: 0.7152 - val_categorical_accuracy: 0.7000\n",
      "Epoch 12/50\n",
      "15/15 [==============================] - 9s 631ms/step - loss: 0.4615 - categorical_accuracy: 0.8400 - val_loss: 0.7218 - val_categorical_accuracy: 0.7083\n",
      "Epoch 13/50\n",
      "15/15 [==============================] - 9s 627ms/step - loss: 0.4641 - categorical_accuracy: 0.8233 - val_loss: 0.8215 - val_categorical_accuracy: 0.6833\n",
      "Epoch 14/50\n",
      "15/15 [==============================] - 10s 638ms/step - loss: 0.4746 - categorical_accuracy: 0.7767 - val_loss: 0.7595 - val_categorical_accuracy: 0.6750\n",
      "Epoch 15/50\n",
      "15/15 [==============================] - 10s 640ms/step - loss: 0.4272 - categorical_accuracy: 0.8367 - val_loss: 0.7713 - val_categorical_accuracy: 0.6583\n",
      "Epoch 16/50\n",
      "15/15 [==============================] - 10s 640ms/step - loss: 0.3775 - categorical_accuracy: 0.8600 - val_loss: 0.7660 - val_categorical_accuracy: 0.6583\n",
      "Epoch 17/50\n",
      "15/15 [==============================] - 9s 627ms/step - loss: 0.3475 - categorical_accuracy: 0.8733 - val_loss: 0.8566 - val_categorical_accuracy: 0.6167\n",
      "Epoch 18/50\n",
      "15/15 [==============================] - 9s 627ms/step - loss: 0.3619 - categorical_accuracy: 0.8633 - val_loss: 0.7964 - val_categorical_accuracy: 0.6333\n",
      "Epoch 19/50\n",
      "15/15 [==============================] - 9s 627ms/step - loss: 0.3050 - categorical_accuracy: 0.9033 - val_loss: 0.8057 - val_categorical_accuracy: 0.6583\n",
      "Epoch 20/50\n",
      "15/15 [==============================] - 10s 653ms/step - loss: 0.2810 - categorical_accuracy: 0.9033 - val_loss: 0.8329 - val_categorical_accuracy: 0.6750\n",
      "Epoch 21/50\n",
      "15/15 [==============================] - 11s 733ms/step - loss: 0.2612 - categorical_accuracy: 0.9333 - val_loss: 0.8622 - val_categorical_accuracy: 0.6667\n"
     ]
    }
   ],
   "source": [
    "# Train the model for (maximum) 50 epochs\n",
    "\n",
    "history = train_model(lsun_model, train_generator, valid_generator, epochs=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot the learning curves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 50452 characters]\n",
      "text/plain": [
       "<Figure size 1080x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Run this cell to plot accuracy vs epoch and loss vs epoch\n",
    "\n",
    "plt.figure(figsize=(15,5))\n",
    "plt.subplot(121)\n",
    "try:\n",
    "    plt.plot(history.history['accuracy'])\n",
    "    plt.plot(history.history['val_accuracy'])\n",
    "except KeyError:\n",
    "    try:\n",
    "        plt.plot(history.history['acc'])\n",
    "        plt.plot(history.history['val_acc'])\n",
    "    except KeyError:\n",
    "        plt.plot(history.history['categorical_accuracy'])\n",
    "        plt.plot(history.history['val_categorical_accuracy'])\n",
    "plt.title('Accuracy vs. epochs')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.xlabel('Epoch')\n",
    "plt.legend(['Training', 'Validation'], loc='lower right')\n",
    "\n",
    "plt.subplot(122)\n",
    "plt.plot(history.history['loss'])\n",
    "plt.plot(history.history['val_loss'])\n",
    "plt.title('Loss vs. epochs')\n",
    "plt.ylabel('Loss')\n",
    "plt.xlabel('Epoch')\n",
    "plt.legend(['Training', 'Validation'], loc='upper right')\n",
    "plt.show() "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may notice overfitting in the above plots, through a growing discrepancy between the training and validation loss and accuracy. We will aim to mitigate this using data augmentation. Given our limited dataset, we may be able to improve the performance by applying random modifications to the images in the training data, effectively increasing the size of the dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create a new data generator with data augmentation\n",
    "\n",
    "You should now write a function to create a new `ImageDataGenerator` object, which performs the following data preprocessing and augmentation:\n",
    "\n",
    "* Scales the image pixel values by a factor of 1/255.\n",
    "* Randomly rotates images by up to 30 degrees\n",
    "* Randomly alters the brightness (picks a brightness shift value) from the range (0.5, 1.5)\n",
    "* Randomly flips images horizontally\n",
    "\n",
    "Hint: you may need to refer to the [documentation](https://keras.io/preprocessing/image/#imagedatagenerator-class) for the `ImageDataGenerator`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### GRADED CELL ####\n",
    "\n",
    "# Complete the following function. \n",
    "# Make sure to not change the function name or arguments.\n",
    "\n",
    "def get_ImageDataGenerator_augmented():\n",
    "    \"\"\"\n",
    "    This function should return an instance of the ImageDataGenerator class \n",
    "    with the above specifications.\n",
    "    \"\"\"\n",
    "    \n",
    "    image_data_gen = ImageDataGenerator(\n",
    "        rescale = (1/255.),\n",
    "        rotation_range = 30,\n",
    "        brightness_range = (0.5,1.5),\n",
    "        horizontal_flip = True)\n",
    "    \n",
    "    return image_data_gen\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Call the function to get an ImageDataGenerator as specified\n",
    "\n",
    "image_gen_aug = get_ImageDataGenerator_augmented()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 120 images belonging to 3 classes.\n",
      "Found 300 images belonging to 3 classes.\n"
     ]
    }
   ],
   "source": [
    "# Run this cell to define training and validation generators \n",
    "\n",
    "valid_generator_aug = get_generator(image_gen_aug, valid_dir)\n",
    "train_generator_aug = get_generator(image_gen_aug, train_dir, seed=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 300 images belonging to 3 classes.\n"
     ]
    }
   ],
   "source": [
    "# Reset the original train_generator with the same random seed\n",
    "\n",
    "train_generator = get_generator(image_gen, train_dir, seed=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Display sample augmented images and labels from the training set\n",
    "\n",
    "The following cell depends on your function `get_generator` to be implemented correctly. If it raises an error, go back and check the function specifications carefully. \n",
    "\n",
    "The cell will display augmented and non-augmented images (and labels) from the training dataset, using the `train_generator_aug` and `train_generator` objects defined above (if the images do not correspond to each other, check you have implemented the `seed` argument correctly)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 193192 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 10 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 177680 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 10 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Display a few images and labels from the non-augmented and augmented generators\n",
    "\n",
    "batch = next(train_generator)\n",
    "batch_images = np.array(batch[0])\n",
    "batch_labels = np.array(batch[1])\n",
    "\n",
    "aug_batch = next(train_generator_aug)\n",
    "aug_batch_images = np.array(aug_batch[0])\n",
    "aug_batch_labels = np.array(aug_batch[1])\n",
    "\n",
    "plt.figure(figsize=(16,5))\n",
    "plt.suptitle(\"Unaugmented images\", fontsize=16)\n",
    "for n, i in enumerate(np.arange(10)):\n",
    "    ax = plt.subplot(2, 5, n+1)\n",
    "    plt.imshow(batch_images[i])\n",
    "    plt.title(lsun_classes[np.where(batch_labels[i] == 1.)[0][0]])\n",
    "    plt.axis('off')\n",
    "plt.figure(figsize=(16,5))\n",
    "plt.suptitle(\"Augmented images\", fontsize=16)\n",
    "for n, i in enumerate(np.arange(10)):\n",
    "    ax = plt.subplot(2, 5, n+1)\n",
    "    plt.imshow(aug_batch_images[i])\n",
    "    plt.title(lsun_classes[np.where(aug_batch_labels[i] == 1.)[0][0]])\n",
    "    plt.axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 300 images belonging to 3 classes.\n"
     ]
    }
   ],
   "source": [
    "# Reset the augmented data generator\n",
    "\n",
    "train_generator_aug = get_generator(image_gen_aug, train_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Train a new model on the augmented dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build and compile a new model\n",
    "\n",
    "lsun_new_model = get_model((64, 64, 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train for 15 steps, validate for 6 steps\n",
      "Epoch 1/50\n",
      "15/15 [==============================] - 11s 761ms/step - loss: 1.0921 - categorical_accuracy: 0.3900 - val_loss: 1.0440 - val_categorical_accuracy: 0.4500\n",
      "Epoch 2/50\n",
      "15/15 [==============================] - 10s 700ms/step - loss: 0.9937 - categorical_accuracy: 0.5300 - val_loss: 0.9592 - val_categorical_accuracy: 0.6000\n",
      "Epoch 3/50\n",
      "15/15 [==============================] - 10s 700ms/step - loss: 0.9181 - categorical_accuracy: 0.5933 - val_loss: 0.8677 - val_categorical_accuracy: 0.6000\n",
      "Epoch 4/50\n",
      "15/15 [==============================] - 10s 683ms/step - loss: 0.8742 - categorical_accuracy: 0.6167 - val_loss: 0.8353 - val_categorical_accuracy: 0.6333\n",
      "Epoch 5/50\n",
      "15/15 [==============================] - 10s 694ms/step - loss: 0.7762 - categorical_accuracy: 0.6633 - val_loss: 0.7767 - val_categorical_accuracy: 0.6417\n",
      "Epoch 6/50\n",
      "15/15 [==============================] - 10s 700ms/step - loss: 0.7298 - categorical_accuracy: 0.6733 - val_loss: 0.8494 - val_categorical_accuracy: 0.6333\n",
      "Epoch 7/50\n",
      "15/15 [==============================] - 10s 689ms/step - loss: 0.7398 - categorical_accuracy: 0.6733 - val_loss: 0.8149 - val_categorical_accuracy: 0.6417\n",
      "Epoch 8/50\n",
      "15/15 [==============================] - 10s 680ms/step - loss: 0.7302 - categorical_accuracy: 0.6900 - val_loss: 0.7711 - val_categorical_accuracy: 0.6333\n",
      "Epoch 9/50\n",
      "15/15 [==============================] - 10s 693ms/step - loss: 0.6948 - categorical_accuracy: 0.7033 - val_loss: 0.8223 - val_categorical_accuracy: 0.6167\n",
      "Epoch 10/50\n",
      "15/15 [==============================] - 11s 713ms/step - loss: 0.7074 - categorical_accuracy: 0.6967 - val_loss: 0.8049 - val_categorical_accuracy: 0.6833\n",
      "Epoch 11/50\n",
      "15/15 [==============================] - 10s 693ms/step - loss: 0.6856 - categorical_accuracy: 0.6967 - val_loss: 0.7708 - val_categorical_accuracy: 0.6750\n",
      "Epoch 12/50\n",
      "15/15 [==============================] - 10s 693ms/step - loss: 0.6761 - categorical_accuracy: 0.7233 - val_loss: 0.7492 - val_categorical_accuracy: 0.7000\n",
      "Epoch 13/50\n",
      "15/15 [==============================] - 10s 693ms/step - loss: 0.6434 - categorical_accuracy: 0.7100 - val_loss: 0.7391 - val_categorical_accuracy: 0.7083\n",
      "Epoch 14/50\n",
      "15/15 [==============================] - 10s 686ms/step - loss: 0.6191 - categorical_accuracy: 0.7400 - val_loss: 0.8185 - val_categorical_accuracy: 0.6333\n",
      "Epoch 15/50\n",
      "15/15 [==============================] - 10s 660ms/step - loss: 0.6553 - categorical_accuracy: 0.7433 - val_loss: 0.7689 - val_categorical_accuracy: 0.6917\n",
      "Epoch 16/50\n",
      "15/15 [==============================] - 10s 673ms/step - loss: 0.6007 - categorical_accuracy: 0.7433 - val_loss: 0.7495 - val_categorical_accuracy: 0.7000\n",
      "Epoch 17/50\n",
      "15/15 [==============================] - 10s 693ms/step - loss: 0.6047 - categorical_accuracy: 0.7400 - val_loss: 0.7563 - val_categorical_accuracy: 0.6833\n",
      "Epoch 18/50\n",
      "15/15 [==============================] - 10s 680ms/step - loss: 0.5955 - categorical_accuracy: 0.7633 - val_loss: 0.7878 - val_categorical_accuracy: 0.6750\n",
      "Epoch 19/50\n",
      "15/15 [==============================] - 10s 680ms/step - loss: 0.6175 - categorical_accuracy: 0.7400 - val_loss: 0.7690 - val_categorical_accuracy: 0.6917\n",
      "Epoch 20/50\n",
      "15/15 [==============================] - 10s 680ms/step - loss: 0.6363 - categorical_accuracy: 0.7567 - val_loss: 0.7529 - val_categorical_accuracy: 0.6583\n",
      "Epoch 21/50\n",
      "15/15 [==============================] - 10s 673ms/step - loss: 0.6516 - categorical_accuracy: 0.7300 - val_loss: 0.8479 - val_categorical_accuracy: 0.6333\n",
      "Epoch 22/50\n",
      "15/15 [==============================] - 10s 693ms/step - loss: 0.5844 - categorical_accuracy: 0.7667 - val_loss: 0.8102 - val_categorical_accuracy: 0.6500\n",
      "Epoch 23/50\n",
      "15/15 [==============================] - 10s 674ms/step - loss: 0.6111 - categorical_accuracy: 0.7233 - val_loss: 0.7905 - val_categorical_accuracy: 0.7250\n"
     ]
    }
   ],
   "source": [
    "# Train the model\n",
    "\n",
    "history_augmented = train_model(lsun_new_model, train_generator_aug, valid_generator_aug, epochs=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot the learning curves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 54688 characters]\n",
      "text/plain": [
       "<Figure size 1080x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Run this cell to plot accuracy vs epoch and loss vs epoch\n",
    "\n",
    "plt.figure(figsize=(15,5))\n",
    "plt.subplot(121)\n",
    "try:\n",
    "    plt.plot(history_augmented.history['accuracy'])\n",
    "    plt.plot(history_augmented.history['val_accuracy'])\n",
    "except KeyError:\n",
    "    try:\n",
    "        plt.plot(history_augmented.history['acc'])\n",
    "        plt.plot(history_augmented.history['val_acc'])\n",
    "    except KeyError:\n",
    "        plt.plot(history_augmented.history['categorical_accuracy'])\n",
    "        plt.plot(history_augmented.history['val_categorical_accuracy'])\n",
    "plt.title('Accuracy vs. epochs')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.xlabel('Epoch')\n",
    "plt.legend(['Training', 'Validation'], loc='lower right')\n",
    "\n",
    "plt.subplot(122)\n",
    "plt.plot(history_augmented.history['loss'])\n",
    "plt.plot(history_augmented.history['val_loss'])\n",
    "plt.title('Loss vs. epochs')\n",
    "plt.ylabel('Loss')\n",
    "plt.xlabel('Epoch')\n",
    "plt.legend(['Training', 'Validation'], loc='upper right')\n",
    "plt.show() "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Do you see an improvement in the overfitting? This will of course vary based on your particular run and whether you have altered the hyperparameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Get predictions using the trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 300 images belonging to 3 classes.\n"
     ]
    }
   ],
   "source": [
    "# Get model predictions for the first 3 batches of test data\n",
    "\n",
    "num_batches = 3\n",
    "seed = 25\n",
    "test_generator = get_generator(image_gen_aug, test_dir, seed=seed)\n",
    "predictions = lsun_new_model.predict_generator(test_generator, steps=num_batches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 300 images belonging to 3 classes.\n",
      "[26 14 27 55]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 104960 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Run this cell to view randomly selected images and model predictions\n",
    "\n",
    "# Get images and ground truth labels\n",
    "test_generator = get_generator(image_gen_aug, test_dir, seed=seed)\n",
    "batches = []\n",
    "for i in range(num_batches):\n",
    "    batches.append(next(test_generator))\n",
    "    \n",
    "batch_images = np.vstack([b[0] for b in batches])\n",
    "batch_labels = np.concatenate([b[1].astype(np.int32) for b in batches])\n",
    "\n",
    "# Randomly select images from the batch\n",
    "inx = np.random.choice(predictions.shape[0], 4, replace=False)\n",
    "print(inx)\n",
    "\n",
    "fig, axes = plt.subplots(4, 2, figsize=(16, 12))\n",
    "fig.subplots_adjust(hspace=0.4, wspace=-0.2)\n",
    "\n",
    "for n, i in enumerate(inx):\n",
    "    axes[n, 0].imshow(batch_images[i])\n",
    "    axes[n, 0].get_xaxis().set_visible(False)\n",
    "    axes[n, 0].get_yaxis().set_visible(False)\n",
    "    axes[n, 0].text(30., -3.5, lsun_classes[np.where(batch_labels[i] == 1.)[0][0]], \n",
    "                    horizontalalignment='center')\n",
    "    axes[n, 1].bar(np.arange(len(predictions[i])), predictions[i])\n",
    "    axes[n, 1].set_xticks(np.arange(len(predictions[i])))\n",
    "    axes[n, 1].set_xticklabels(lsun_classes)\n",
    "    axes[n, 1].set_title(f\"Categorical distribution. Model prediction: {lsun_classes[np.argmax(predictions[i])]}\")\n",
    "    \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congratulations! This completes the first part of the programming assignment using the tf.keras image data processing tools."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Part 2: tf.data\n",
    "\n",
    "![CIFAR-100 overview image](data/cifar100/cifar100.png)\n",
    "\n",
    "#### The CIFAR-100 Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the second part of this assignment, you will use the [CIFAR-100 dataset](https://www.cs.toronto.edu/~kriz/cifar.html). This image dataset has 100 classes with 500 training images and 100 test images per class. \n",
    "\n",
    "* A. Krizhevsky. \"Learning Multiple Layers of Features from Tiny Images\". April 2009 \n",
    "\n",
    "Your goal is to use the tf.data module preprocessing tools to construct a data ingestion pipeline including filtering and function mapping over the dataset to train a neural network to classify the images."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading data from https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz\n",
      "169009152/169001437 [==============================] - 3s 0us/step\n"
     ]
    }
   ],
   "source": [
    "# Load the data, along with the labels\n",
    "\n",
    "(train_data, train_labels), (test_data, test_labels) = cifar100.load_data(label_mode='fine')\n",
    "with open('data/cifar100/cifar100_labels.json', 'r') as j:\n",
    "    cifar_labels = json.load(j)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Display sample images and labels from the training set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 171532 characters]\n",
      "text/plain": [
       "<Figure size 1080x576 with 32 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Display a few images and labels\n",
    "\n",
    "plt.figure(figsize=(15,8))\n",
    "inx = np.random.choice(train_data.shape[0], 32, replace=False)\n",
    "for n, i in enumerate(inx):\n",
    "    ax = plt.subplot(4, 8, n+1)\n",
    "    plt.imshow(train_data[i])\n",
    "    plt.title(cifar_labels[int(train_labels[i])])\n",
    "    plt.axis('off')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create Dataset objects for the train and test images\n",
    "\n",
    "You should now write a function to create a `tf.data.Dataset` object for each of the training and test images and labels. This function should take a numpy array of images in the first argument and a numpy array of labels in the second argument, and create a `Dataset` object. \n",
    "\n",
    "Your function should then return the `Dataset` object. Do not batch or shuffle the `Dataset` (this will be done later)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### GRADED CELL ####\n",
    "\n",
    "# Complete the following function. \n",
    "# Make sure to not change the function name or arguments.\n",
    "\n",
    "def create_dataset(data, labels):\n",
    "    \"\"\"\n",
    "    This function takes a numpy array batch of images in the first argument, and\n",
    "    a corresponding array containing the labels in the second argument.\n",
    "    The function should then create a tf.data.Dataset object with these inputs\n",
    "    and outputs, and return it.\n",
    "    \"\"\"\n",
    "    \n",
    "    dataset = tf.data.Dataset.from_tensor_slices((data,labels))\n",
    "    return dataset\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the below cell to convert the training and test data and labels into datasets\n",
    "\n",
    "train_dataset = create_dataset(train_data, train_labels)\n",
    "test_dataset = create_dataset(test_data, test_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(TensorSpec(shape=(32, 32, 3), dtype=tf.uint8, name=None), TensorSpec(shape=(1,), dtype=tf.int64, name=None))\n",
      "(TensorSpec(shape=(32, 32, 3), dtype=tf.uint8, name=None), TensorSpec(shape=(1,), dtype=tf.int64, name=None))\n"
     ]
    }
   ],
   "source": [
    "# Check the element_spec of your datasets\n",
    "\n",
    "print(train_dataset.element_spec)\n",
    "print(test_dataset.element_spec)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Filter the Dataset\n",
    "\n",
    "Write a function to filter the train and test datasets so that they only generate images that belong to a specified set of classes. \n",
    "\n",
    "The function should take a `Dataset` object in the first argument, and a list of integer class indices in the second argument. Inside your function you should define an auxiliary function that you will use with the `filter` method of the `Dataset` object. This auxiliary function should take image and label arguments (as in the `element_spec`) for a single element in the batch, and return a boolean indicating if the label is one of the allowed classes. \n",
    "\n",
    "Your function should then return the filtered dataset.\n",
    "\n",
    "**Hint:** you may need to use the [`tf.equal`](https://www.tensorflow.org/api_docs/python/tf/math/equal), [`tf.cast`](https://www.tensorflow.org/api_docs/python/tf/dtypes/cast) and [`tf.math.reduce_any`](https://www.tensorflow.org/api_docs/python/tf/math/reduce_any) functions in your auxiliary function. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### GRADED CELL ####\n",
    "\n",
    "# Complete the following function. \n",
    "# Make sure to not change the function name or arguments.\n",
    "\n",
    "def filter_classes(dataset, classes):\n",
    "    \"\"\"\n",
    "    This function should filter the dataset by only retaining dataset elements whose\n",
    "    label belongs to one of the integers in the classes list.\n",
    "    The function should then return the filtered Dataset object.\n",
    "    \"\"\"\n",
    "    def filterer(image,label):\n",
    "        \n",
    "            \n",
    "        flag = tf.math.reduce_any(tf.equal(label,classes))\n",
    "        return flag\n",
    "        \n",
    "    \n",
    "    dataset = dataset.filter(filterer)\n",
    "    return dataset\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the below cell to filter the datasets using your function\n",
    "\n",
    "cifar_classes = [0, 29, 99] # Your datasets should contain only classes in this list\n",
    "\n",
    "train_dataset = filter_classes(train_dataset, cifar_classes)\n",
    "test_dataset = filter_classes(test_dataset, cifar_classes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Apply map functions to the Dataset\n",
    "\n",
    "You should now write two functions that use the `map` method to process the images and labels in the filtered dataset. \n",
    "\n",
    "The first function should one-hot encode the remaining labels so that we can train the network using a categorical cross entropy loss. \n",
    "\n",
    "The function should take a `Dataset` object as an argument. Inside your function you should define an auxiliary function that you will use with the `map` method of the `Dataset` object. This auxiliary function should take image and label arguments (as in the `element_spec`) for a single element in the batch, and return a tuple of two elements, with the unmodified image in the first element, and a one-hot vector in the second element. The labels should be encoded according to the following:\n",
    "\n",
    "* Class 0 maps to `[1., 0., 0.]`\n",
    "* Class 29 maps to `[0., 1., 0.]`\n",
    "* Class 99 maps to `[0., 0., 1.]`\n",
    "\n",
    "Your function should then return the mapped dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### GRADED CELL ####\n",
    "\n",
    "# Complete the following function. \n",
    "# Make sure to not change the function name or arguments.\n",
    "\n",
    "def map_labels(dataset):\n",
    "    \"\"\"\n",
    "    This function should map over the dataset to convert the label to a \n",
    "    one-hot vector. The encoding should be done according to the above specification.\n",
    "    The function should then return the mapped Dataset object.\n",
    "    \"\"\"\n",
    "    def One_hot(image,label):\n",
    "\n",
    "        if label== 0:\n",
    "\n",
    "            label = [1., 0., 0.]\n",
    "\n",
    "        elif label== 29:\n",
    "\n",
    "            label =[0., 1., 0.]\n",
    "\n",
    "        else:\n",
    "\n",
    "            label = [0., 0., 1.]\n",
    "        \n",
    "\n",
    "       \n",
    "        return (image,label)\n",
    "        \n",
    "    dataset = dataset.map(One_hot)\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the below cell to one-hot encode the training and test labels.\n",
    "\n",
    "train_dataset = map_labels(train_dataset)\n",
    "test_dataset = map_labels(test_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorSpec(shape=(32, 32, 3), dtype=tf.uint8, name=None),\n",
       " TensorSpec(shape=(3,), dtype=tf.float32, name=None))"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_dataset.element_spec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The second function should process the images according to the following specification:\n",
    "\n",
    "* Rescale the image pixel values by a factor of 1/255.\n",
    "* Convert the colour images (3 channels) to black and white images (single channel) by computing the average pixel value across all channels. \n",
    "\n",
    "The function should take a `Dataset` object as an argument. Inside your function you should again define an auxiliary function that you will use with the `map` method of the `Dataset` object. This auxiliary function should take image and label arguments (as in the `element_spec`) for a single element in the batch, and return a tuple of two elements, with the processed image in the first element, and the unmodified label in the second argument.\n",
    "\n",
    "Your function should then return the mapped dataset.\n",
    "\n",
    "**Hint:** you may find it useful to use [`tf.reduce_mean`](https://www.tensorflow.org/api_docs/python/tf/math/reduce_mean?version=stable) since the black and white image is the colour-average of the colour images. You can also use the `keepdims` keyword in `tf.reduce_mean` to retain the single colour channel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### GRADED CELL ####\n",
    "\n",
    "# Complete the following function. \n",
    "# Make sure to not change the function name or arguments.\n",
    "\n",
    "def map_images(dataset):\n",
    "    \"\"\"\n",
    "    This function should map over the dataset to process the image according to the \n",
    "    above specification. The function should then return the mapped Dataset object.\n",
    "    \"\"\"\n",
    "\n",
    "    \n",
    "    def rescale(image, label):\n",
    "\n",
    "        image = tf.cast(image, tf.float32) / 255.0\n",
    "\n",
    "        image = tf.reduce_mean(image,axis = 2,keepdims=True)\n",
    "\n",
    "        return (image, label)\n",
    "        \n",
    "        \n",
    "    \n",
    "    dataset = dataset.map(rescale)\n",
    "    return dataset\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the below cell to apply your mapping function to the datasets\n",
    "\n",
    "train_dataset_bw = map_images(train_dataset)\n",
    "test_dataset_bw = map_images(test_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Display a batch of processed images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 57512 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 10 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 41084 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 10 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Run this cell to view a selection of images before and after processing\n",
    "\n",
    "plt.figure(figsize=(16,5))\n",
    "plt.suptitle(\"Unprocessed images\", fontsize=16)\n",
    "for n, elem in enumerate(train_dataset.take(10)):\n",
    "    images, labels = elem\n",
    "    ax = plt.subplot(2, 5, n+1)\n",
    "    plt.title(cifar_labels[cifar_classes[np.where(labels == 1.)[0][0]]])\n",
    "    plt.imshow(np.squeeze(images), cmap='gray')\n",
    "    plt.axis('off')\n",
    "    \n",
    "plt.figure(figsize=(16,5))\n",
    "plt.suptitle(\"Processed images\", fontsize=16)\n",
    "for n, elem in enumerate(train_dataset_bw.take(10)):\n",
    "    images_bw, labels_bw = elem\n",
    "    ax = plt.subplot(2, 5, n+1)\n",
    "    plt.title(cifar_labels[cifar_classes[np.where(labels_bw == 1.)[0][0]]])\n",
    "    plt.imshow(np.squeeze(images_bw), cmap='gray')\n",
    "    plt.axis('off')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now batch and shuffle the Dataset objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the below cell to batch the training dataset and expand the final dimensinos\n",
    "\n",
    "train_dataset_bw = train_dataset_bw.batch(10)\n",
    "train_dataset_bw = train_dataset_bw.shuffle(100)\n",
    "\n",
    "test_dataset_bw = test_dataset_bw.batch(10)\n",
    "test_dataset_bw = test_dataset_bw.shuffle(100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Train a neural network model\n",
    "\n",
    "Now we will train a model using the `Dataset` objects. We will use the model specification and function from the first part of this assignment, only modifying the size of the input images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build and compile a new model with our original spec, using the new image size\n",
    "    \n",
    "cifar_model = get_model((32, 32, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<ShuffleDataset shapes: ((None, 32, 32, 1), (None, 3)), types: (tf.float32, tf.float32)>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset_bw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/15\n",
      "150/150 [==============================] - 19s 127ms/step - loss: 1.0330 - categorical_accuracy: 0.4607 - val_loss: 0.0000e+00 - val_categorical_accuracy: 0.0000e+00\n",
      "Epoch 2/15\n",
      "150/150 [==============================] - 18s 122ms/step - loss: 0.9124 - categorical_accuracy: 0.5660 - val_loss: 0.8098 - val_categorical_accuracy: 0.6433\n",
      "Epoch 3/15\n",
      "150/150 [==============================] - 19s 124ms/step - loss: 0.8025 - categorical_accuracy: 0.6567 - val_loss: 0.7377 - val_categorical_accuracy: 0.7267\n",
      "Epoch 4/15\n",
      "150/150 [==============================] - 19s 126ms/step - loss: 0.7323 - categorical_accuracy: 0.6940 - val_loss: 0.6823 - val_categorical_accuracy: 0.7367\n",
      "Epoch 5/15\n",
      "150/150 [==============================] - 18s 122ms/step - loss: 0.6936 - categorical_accuracy: 0.7160 - val_loss: 0.6610 - val_categorical_accuracy: 0.7333\n",
      "Epoch 6/15\n",
      "150/150 [==============================] - 18s 123ms/step - loss: 0.6527 - categorical_accuracy: 0.7307 - val_loss: 0.6424 - val_categorical_accuracy: 0.7533\n",
      "Epoch 7/15\n",
      "150/150 [==============================] - 18s 121ms/step - loss: 0.6353 - categorical_accuracy: 0.7353 - val_loss: 0.6299 - val_categorical_accuracy: 0.7567\n",
      "Epoch 8/15\n",
      "150/150 [==============================] - 18s 121ms/step - loss: 0.5887 - categorical_accuracy: 0.7553 - val_loss: 0.6124 - val_categorical_accuracy: 0.7533\n",
      "Epoch 9/15\n",
      "150/150 [==============================] - 18s 121ms/step - loss: 0.5996 - categorical_accuracy: 0.7613 - val_loss: 0.6354 - val_categorical_accuracy: 0.7500\n",
      "Epoch 10/15\n",
      "150/150 [==============================] - 18s 121ms/step - loss: 0.5592 - categorical_accuracy: 0.7700 - val_loss: 0.6084 - val_categorical_accuracy: 0.7533\n",
      "Epoch 11/15\n",
      "150/150 [==============================] - 18s 123ms/step - loss: 0.5527 - categorical_accuracy: 0.7720 - val_loss: 0.6201 - val_categorical_accuracy: 0.7433\n",
      "Epoch 12/15\n",
      "150/150 [==============================] - 18s 121ms/step - loss: 0.5486 - categorical_accuracy: 0.7787 - val_loss: 0.5873 - val_categorical_accuracy: 0.7633\n",
      "Epoch 13/15\n",
      "150/150 [==============================] - 18s 122ms/step - loss: 0.5131 - categorical_accuracy: 0.7913 - val_loss: 0.5851 - val_categorical_accuracy: 0.7800\n",
      "Epoch 14/15\n",
      "150/150 [==============================] - 18s 123ms/step - loss: 0.5064 - categorical_accuracy: 0.7987 - val_loss: 0.5622 - val_categorical_accuracy: 0.7767\n",
      "Epoch 15/15\n",
      "150/150 [==============================] - 18s 122ms/step - loss: 0.4948 - categorical_accuracy: 0.8007 - val_loss: 0.5729 - val_categorical_accuracy: 0.7900\n"
     ]
    }
   ],
   "source": [
    "# Train the model for 15 epochs\n",
    "\n",
    "history = cifar_model.fit(train_dataset_bw, validation_data=test_dataset_bw, epochs=15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot the learning curves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 45812 characters]\n",
      "text/plain": [
       "<Figure size 1080x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Run this cell to plot accuracy vs epoch and loss vs epoch\n",
    "\n",
    "plt.figure(figsize=(15,5))\n",
    "plt.subplot(121)\n",
    "try:\n",
    "    plt.plot(history.history['accuracy'])\n",
    "    plt.plot(history.history['val_accuracy'])\n",
    "except KeyError:\n",
    "    try:\n",
    "        plt.plot(history.history['acc'])\n",
    "        plt.plot(history.history['val_acc'])\n",
    "    except KeyError:\n",
    "        plt.plot(history.history['categorical_accuracy'])\n",
    "        plt.plot(history.history['val_categorical_accuracy'])\n",
    "plt.title('Accuracy vs. epochs')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.xlabel('Epoch')\n",
    "plt.legend(['Training', 'Validation'], loc='lower right')\n",
    "\n",
    "plt.subplot(122)\n",
    "plt.plot(history.history['loss'])\n",
    "plt.plot(history.history['val_loss'])\n",
    "plt.title('Loss vs. epochs')\n",
    "plt.ylabel('Loss')\n",
    "plt.xlabel('Epoch')\n",
    "plt.legend(['Training', 'Validation'], loc='upper right')\n",
    "plt.show() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an iterable from the batched test dataset\n",
    "\n",
    "test_dataset = test_dataset.batch(10)\n",
    "iter_test_dataset = iter(test_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 62220 characters]\n",
      "text/plain": [
       "<Figure size 1080x576 with 10 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Display model predictions for a sample of test images\n",
    "\n",
    "plt.figure(figsize=(15,8))\n",
    "inx = np.random.choice(test_data.shape[0], 18, replace=False)\n",
    "images, labels = next(iter_test_dataset)\n",
    "probs = cifar_model(tf.reduce_mean(tf.cast(images, tf.float32), axis=-1, keepdims=True) / 255.)\n",
    "preds = np.argmax(probs, axis=1)\n",
    "for n in range(10):\n",
    "    ax = plt.subplot(2, 5, n+1)\n",
    "    plt.imshow(images[n])\n",
    "    plt.title(cifar_labels[cifar_classes[np.where(labels[n].numpy() == 1.0)[0][0]]])\n",
    "    plt.text(0, 35, \"Model prediction: {}\".format(cifar_labels[cifar_classes[preds[n]]]))\n",
    "    plt.axis('off')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congratulations for completing this programming assignment! In the next week of the course we will learn to develop models for sequential data."
   ]
  }
 ],
 "metadata": {
  "coursera": {
   "course_slug": "tensor-flow-2-2",
   "graded_item_id": "3hWzU",
   "launcher_item_id": "AStQh"
  },
  "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
