notebooks/image_classification/tf_image_classification/Image_Classification_Transfer_Learning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transfer Learning for Image Classification\n",
    "\n",
    "This notebook uses a classifier model that was originally trained using [ImageNet](https://image-net.org) and does transfer learning with either a TF dataset or your own raw images.\n",
    "The notebook performs the following steps:\n",
    "1. [Import dependencies and setup parameters](#1.-Import-dependencies-and-setup-parameters)\n",
    "2. [Prepare the dataset](#2.-Prepare-the-dataset) using either a TF dataset or your own images\n",
    "3. [Predict using the original model](#3.-Predict-using-the-original-model)\n",
    "4. [Transfer learning](#4.-Transfer-Learning)\n",
    "5. [Evaluate the model](#5.-Evaluate-the-model)\n",
    "6. [Export the saved model](#6.-Export-the-saved-model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Import dependencies and setup parameters\n",
    "\n",
    "This notebook assumes that you have already followed the instructions in the [README.md](/notebooks/README.md) to setup a TensorFlow environment with all the dependencies required to run the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_hub as hub\n",
    "import tensorflow_datasets as tfds\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "\n",
    "from model_util import tfhub_model_map\n",
    "from tlt.utils.file_utils import download_and_extract_tar_file\n",
    "\n",
    "print('Supported models:')\n",
    "print('\\n'.join(tfhub_model_map.keys()))\n",
    "\n",
    "# Specify the the parent directory for the custom or tf dataset\n",
    "dataset_directory = os.environ[\"DATASET_DIR\"] if \"DATASET_DIR\" in os.environ else \\\n",
    "    os.path.join(os.environ[\"HOME\"], \"dataset\")\n",
    "    \n",
    "# Specify a directory for output\n",
    "output_directory = os.environ[\"OUTPUT_DIR\"] if \"OUTPUT_DIR\" in os.environ else \\\n",
    "    os.path.join(os.environ[\"HOME\"], \"output\")\n",
    "\n",
    "# Batch size\n",
    "batch_size = 32\n",
    "\n",
    "print(\"\\nDataset directory:\", dataset_directory)\n",
    "print(\"Output directory:\", output_directory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify a model from the list above\n",
    "model_name = \"efficientnet_b0\"\n",
    "\n",
    "if model_name not in tfhub_model_map.keys():\n",
    "    raise ValueError(\"The specified model_name ({}) is invalid. Please select from: {}\".\n",
    "                     format(model_name, tfhub_model_map.keys()))\n",
    "    \n",
    "# Get the info for the specified model from the map\n",
    "model_map_values = tfhub_model_map[model_name]\n",
    "model_handle = tfhub_model_map[model_name][\"imagenet_model\"]\n",
    "feature_vector_handle = tfhub_model_map[model_name][\"feature_vector\"]\n",
    "image_size = tfhub_model_map[model_name][\"image_size\"]\n",
    "print(\"Model:\", model_name)\n",
    "print(\"Classifier model:\", model_handle)\n",
    "print(\"Feature vector:\", feature_vector_handle)\n",
    "print(\"Image size:\", image_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare the dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option A: Use your own image dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use your own image dataset for transfer learning with the rest of this notebook, format your images as `.jpg` files and save them in folders named after the classes that you want the model to predict. To provide a working example using the correct layout, we will download and extract a flower species dataset. This is different from using the TF dataset called `tf_flowers`, although they are the same images, because the download contains image files, not tf_records, and we are not using the TF datasets API. After downloading and extracting, you will have the following  subdirectories in your dataset directory. Each species subfolder will contain numerous `.jpg` files:\n",
    "\n",
    "```\n",
    "dataset_directory\n",
    "└── flower_photos\n",
    "    └── daisy\n",
    "    └── dandelion\n",
    "    └── roses\n",
    "    └── sunflowers\n",
    "    └── tulips\n",
    "```\n",
    "\n",
    "Use this as an example to organize your own image files accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# When you have your own properly organized subdirectory of images, adjust this variable\n",
    "dataset_subdir = os.path.join(dataset_directory, \"flower_photos\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Only run this if you want to use the example flowers dataset\n",
    "if not os.path.exists(dataset_subdir):\n",
    "    os.makedirs(dataset_subdir)\n",
    "    dataset_url = \"https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz\"\n",
    "    download_and_extract_tar_file(dataset_url, dataset_directory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pathlib\n",
    "data_subdir_path = pathlib.Path(dataset_subdir)\n",
    "image_count = len(list(data_subdir_path.glob('*/*.jpg')))\n",
    "print('Images:', image_count)\n",
    "print('Image Size:', image_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an image data generator and partition the data into train and test sets\n",
    "img_gen = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1./255, validation_split=0.25)\n",
    "train_ds = img_gen.flow_from_directory(data_subdir_path, batch_size=batch_size, \n",
    "                                   target_size=(image_size, image_size),\n",
    "                                   class_mode='sparse', subset='training')\n",
    "test_ds = img_gen.flow_from_directory(data_subdir_path, batch_size=batch_size, \n",
    "                                  target_size=(image_size, image_size),\n",
    "                                  class_mode='sparse', subset='validation')\n",
    "\n",
    "# Get class names for the dataset\n",
    "class_names = [k for k in train_ds.class_indices.keys()]\n",
    "print(\"Number of classes:\", len(class_names))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Skip to the next step [3. Predict using the original model](#3.-Predict-using-the-original-model) to continue using your own image dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option B: Use a TF dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use a TF Dataset, specify the name of the dataset to load from the TF Datasets catalog, preprocess the images to convert them to float32, and resize the images. This example uses the [Food-101 dataset using the TensorFlow datasets API](https://www.tensorflow.org/datasets/catalog/food101) dataset, but you can choose from a wide variety of [options](https://www.tensorflow.org/datasets/catalog/overview) (click on the \"Image classification\" section). If the dataset is not found in the dataset directory it is downloaded. Subsequent runs will reuse the already downloaded dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Options include: \"food101\", cats_vs_dogs\", \"rock_paper_scissors\", and \"tf_flowers\"\n",
    "tf_dataset = \"food101\"\n",
    "\n",
    "# Load the dataset using the TensorFlow datasets API\n",
    "[train_ds, test_ds], info = tfds.load(tf_dataset,\n",
    "                       data_dir=dataset_directory,\n",
    "                       split=[\"train[:75%]\", \"train[75%:]\"],\n",
    "                       as_supervised=True,\n",
    "                       shuffle_files=True,\n",
    "                       with_info=True)\n",
    "\n",
    "# Preprocess the images to convert them to float32 and resize the images to match our model\n",
    "def preprocess_image(image, label):\n",
    "    image = tf.image.convert_image_dtype(image, tf.float32)\n",
    "    image = tf.image.resize_with_pad(image, image_size, image_size)\n",
    "    return (image, label)\n",
    "\n",
    "train_ds = train_ds.map(preprocess_image)\n",
    "test_ds = test_ds.map(preprocess_image)\n",
    "\n",
    "print(\"Dataset directory: \", dataset_directory)\n",
    "print(\"Training dataset size:\", len(train_ds))\n",
    "print(\"Validation dataset size:\", len(test_ds))\n",
    "\n",
    "# Training data is shuffled for randomness\n",
    "# https://www.tensorflow.org/datasets/keras_example#build_a_training_pipeline\n",
    "train_ds = train_ds.cache()\n",
    "train_ds = train_ds.shuffle(info.splits['train'].num_examples)\n",
    "train_ds = train_ds.batch(batch_size)\n",
    "train_ds = train_ds.prefetch(tf.data.AUTOTUNE)\n",
    "\n",
    "# Test data does not need to be shuffled, and caching is done after batching\n",
    "# https://www.tensorflow.org/datasets/keras_example#build_an_evaluation_pipeline\n",
    "test_ds = test_ds.batch(batch_size)\n",
    "test_ds = test_ds.cache()\n",
    "test_ds = test_ds.prefetch(tf.data.AUTOTUNE)\n",
    "\n",
    "# Get class names for the dataset\n",
    "class_names = info.features[\"label\"].names\n",
    "print(\"Number of classes:\", len(class_names))\n",
    "print('After processing and batching: ', train_ds.element_spec)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Predict using the original model\n",
    "\n",
    "Use the classifier model that was trained using ImageNet to do predictions with the dataset and view the results for a single batch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get a batch of the dataset to use for testing\n",
    "batch = next(iter(test_ds))\n",
    "image_batch, label_batch = batch\n",
    "\n",
    "# List of the actual labels for this batch\n",
    "actual_label_batch = [class_names[int(id)] for id in label_batch]\n",
    "\n",
    "# Download the ImageNet labels and load them into a list\n",
    "labels_file = \"https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt\"\n",
    "downloaded_file = tf.keras.utils.get_file(\"labels.txt\", origin=labels_file)\n",
    "imagenet_classes = []\n",
    "\n",
    "with open(downloaded_file) as f:\n",
    "    imagenet_labels = f.readlines()\n",
    "    imagenet_classes = [l.strip() for l in imagenet_labels]\n",
    "\n",
    "# Predict using the TF Hub classifier that was trained using ImageNet\n",
    "classifier = tf.keras.Sequential([\n",
    "    hub.KerasLayer(model_handle, input_shape=(image_size, image_size)+(3,))\n",
    "])\n",
    "predicted_batch = classifier.predict(image_batch)\n",
    "predicted_id = np.argmax(predicted_batch, axis=-1)\n",
    "predicted_label_batch = [imagenet_classes[id] for id in predicted_id]\n",
    "\n",
    "# Visualize the results\n",
    "plt.figure(figsize=(16,16))\n",
    "plt.subplots_adjust(hspace=0.5)\n",
    "for n in range(min(batch_size, 30)):\n",
    "    plt.subplot(6,5,n+1)\n",
    "    plt.imshow(image_batch[n])\n",
    "    correct_prediction = actual_label_batch[n] == predicted_label_batch[n]\n",
    "    color = \"darkgreen\" if correct_prediction else \"crimson\"\n",
    "    title = predicted_label_batch[n].title() if correct_prediction else \"{}\\n({})\".format(predicted_label_batch[n], actual_label_batch[n]) \n",
    "    plt.title(title, fontsize=14, color=color)\n",
    "    plt.axis('off')\n",
    "_ = plt.suptitle(\"ImageNet predictions\", fontsize=16)\n",
    "plt.show()\n",
    "\n",
    "print(\"Correct predictions are shown in green\")\n",
    "print(\"Incorrect predictions are shown in red with the actual label in parenthesis\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Transfer Learning\n",
    "\n",
    "Get the feature vector from TF Hub and add on a dense layer based on the number of classes in our dataset. Train the model using the training dataset for the specified number of epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Number of training epochs\n",
    "training_epochs = 1\n",
    "\n",
    "# To reduce training time, the feature extractor layer can remain frozen (do_fine_tuning=False).\n",
    "# Fine-tuning can be enabled to potentially get better accuracy. Note that enabling fine-tuning\n",
    "# will increase training time.\n",
    "do_fine_tuning = False\n",
    "\n",
    "# Optionally add a dropout layer (set to a float between 0 and 1, or None).\n",
    "# If set to None, no dropout layer will be added.\n",
    "dropout_layer_rate = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_extractor_layer = hub.KerasLayer(feature_vector_handle,\n",
    "                                         input_shape=(image_size, image_size, 3),\n",
    "                                         trainable=do_fine_tuning)\n",
    "\n",
    "if dropout_layer_rate == None:\n",
    "    model = tf.keras.Sequential([\n",
    "      feature_extractor_layer,\n",
    "      tf.keras.layers.Dense(len(class_names))\n",
    "    ])\n",
    "else:\n",
    "    model = tf.keras.Sequential([\n",
    "      feature_extractor_layer,\n",
    "      tf.keras.layers.Dropout(dropout_layer_rate),\n",
    "      tf.keras.layers.Dense(len(class_names))\n",
    "    ])\n",
    "\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "model.compile(\n",
    "  optimizer=tf.keras.optimizers.Adam(),\n",
    "  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
    "  metrics=['acc'])\n",
    "\n",
    "class CollectBatchStats(tf.keras.callbacks.Callback):\n",
    "    def __init__(self):\n",
    "        self.batch_losses = []\n",
    "        self.batch_acc = []\n",
    "\n",
    "    def on_epoch_begin(self, epoch, logs=None):\n",
    "        self.batch_losses = []\n",
    "        self.batch_acc = []\n",
    "\n",
    "    def on_train_batch_begin(self, batch, logs=None):\n",
    "        self.model.reset_metrics()\n",
    "\n",
    "    def on_train_batch_end(self, batch, logs=None):\n",
    "        self.batch_losses.append(logs['loss'])\n",
    "        self.batch_acc.append(logs['acc'])\n",
    "\n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        # Using the average over all batches is also common instead of just the last batch\n",
    "        logs['loss'] = self.batch_losses[-1]  # np.mean(self.batch_losses)\n",
    "        logs['acc'] = self.batch_acc[-1]  # np.mean(self.batch_acc)\n",
    "\n",
    "batch_stats_callback = CollectBatchStats()\n",
    "\n",
    "history = model.fit(train_ds, epochs=training_epochs, shuffle=True, callbacks=[batch_stats_callback])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate the model\n",
    "\n",
    "After the training completes, evaluate the model's accuracy using the validation dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "model.evaluate(test_ds, batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Also, predict using the same sample batch that we used earlier with the ImageNet trained classier to visualize the results after training the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Predict using the sample batch\n",
    "predicted_batch = model.predict(image_batch)\n",
    "predicted_id = np.argmax(predicted_batch, axis=-1)\n",
    "predicted_label_batch = [class_names[id] for id in predicted_id]\n",
    "\n",
    "# Display the results\n",
    "plt.figure(figsize=(16,16))\n",
    "plt.subplots_adjust(hspace=0.5)\n",
    "for n in range(min(batch_size, 30)):\n",
    "    plt.subplot(6,5,n+1)\n",
    "    plt.imshow(image_batch[n])\n",
    "    correct_prediction = actual_label_batch[n] == predicted_label_batch[n]\n",
    "    color = \"darkgreen\" if correct_prediction else \"crimson\"\n",
    "    title = predicted_label_batch[n].title() if correct_prediction else \"{}\\n({})\".format(predicted_label_batch[n], actual_label_batch[n]) \n",
    "    plt.title(title, fontsize=14, color=color)\n",
    "    plt.axis('off')\n",
    "_ = plt.suptitle(\"Model predictions\", fontsize=16)\n",
    "plt.show()\n",
    "print(\"Correct predictions are shown in green\")\n",
    "print(\"Incorrect predictions are shown in red with the actual label in parenthesis\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Export the saved model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "saved_model_dir = os.path.join(output_directory, \"{}_saved_model\".format(model_name))\n",
    "model.save(saved_model_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset citations\n",
    "```\n",
    "@inproceedings{bossard14,\n",
    "  title = {Food-101 -- Mining Discriminative Components with Random Forests},\n",
    "  author = {Bossard, Lukas and Guillaumin, Matthieu and Van Gool, Luc},\n",
    "  booktitle = {European Conference on Computer Vision},\n",
    "  year = {2014}\n",
    "}\n",
    "\n",
    "@ONLINE {tfflowers,\n",
    "author = \"The TensorFlow Team\",\n",
    "title = \"Flowers\",\n",
    "month = \"jan\",\n",
    "year = \"2019\",\n",
    "url = \"http://download.tensorflow.org/example_images/flower_photos.tgz\" }\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}