08_09_Star_Wars_Project.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import math\n",
    "import random\n",
    "import shutil\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sn\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# download the data from kaggle: \n",
    "# https://www.kaggle.com/ihelon/lego-minifigures-tensorflow-tutorial\n",
    "# move the folder into your project folder and create a backup of\n",
    "# the star-wars images at 'lego/star-wars-images/'\n",
    "BASE_DIR = 'lego/star-wars-images/'\n",
    "names = [\"YODA\", \"LUKE SKYWALKER\", \"R2-D2\", \"MACE WINDU\", \"GENERAL GRIEVOUS\"]\n",
    "\n",
    "tf.random.set_seed(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reorganize the folder structure:\n",
    "if not os.path.isdir(BASE_DIR + 'train/'):\n",
    "    for name in names:\n",
    "        os.makedirs(BASE_DIR + 'train/' + name)\n",
    "        os.makedirs(BASE_DIR + 'val/' + name)\n",
    "        os.makedirs(BASE_DIR + 'test/' + name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Moce the image files\n",
    "orig_folders = [\"0001/\", \"0002/\", \"0003/\", \"0004/\", \"0005/\"]\n",
    "for folder_idx, folder in enumerate(orig_folders):\n",
    "    files = os.listdir(BASE_DIR + folder)\n",
    "    number_of_images = len([name for name in files])\n",
    "    n_train = int((number_of_images * 0.6) + 0.5)\n",
    "    n_valid = int((number_of_images*0.25) + 0.5)\n",
    "    n_test = number_of_images - n_train - n_valid\n",
    "    print(number_of_images, n_train, n_valid, n_test)\n",
    "    for idx, file in enumerate(files):\n",
    "        file_name = BASE_DIR + folder + file\n",
    "        if idx < n_train:\n",
    "            shutil.move(file_name, BASE_DIR + \"train/\" + names[folder_idx])\n",
    "        elif idx < n_train + n_valid:\n",
    "            shutil.move(file_name, BASE_DIR + \"val/\" + names[folder_idx])\n",
    "        else:\n",
    "            shutil.move(file_name, BASE_DIR + \"test/\" + names[folder_idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate batches of tensor image data with\n",
    "# optional real-time data augmentation.\n",
    "\n",
    "# preprocessing_function\n",
    "# rescale=1./255 -> [0,1]\n",
    "train_gen = keras.preprocessing.image.ImageDataGenerator(rescale=1./255)\n",
    "#    rotation_range=20,\n",
    "#    horizontal_flip=True,\n",
    "#    width_shift_range=0.2, height_shift_range=0.2,\n",
    "#    shear_range=0.2, zoom_range=0.2)\n",
    "\n",
    "valid_gen = keras.preprocessing.image.ImageDataGenerator(rescale=1./255)\n",
    "test_gen = keras.preprocessing.image.ImageDataGenerator(rescale=1./255)\n",
    "\n",
    "train_batches = train_gen.flow_from_directory(\n",
    "    'lego/star-wars-images/train',\n",
    "    target_size=(256, 256),\n",
    "    class_mode='sparse',\n",
    "    batch_size=4,\n",
    "    shuffle=True,\n",
    "    color_mode=\"rgb\",\n",
    "    classes=names   \n",
    ")\n",
    "\n",
    "val_batches = valid_gen.flow_from_directory(\n",
    "    'lego/star-wars-images/val',\n",
    "    target_size=(256, 256),\n",
    "    class_mode='sparse',\n",
    "    batch_size=4,\n",
    "    shuffle=False,\n",
    "    color_mode=\"rgb\",\n",
    "    classes=names\n",
    ")\n",
    "\n",
    "test_batches = test_gen.flow_from_directory(\n",
    "    'lego/star-wars-images/test',\n",
    "    target_size=(256, 256),\n",
    "    class_mode='sparse',\n",
    "    batch_size=4,\n",
    "    shuffle=False,\n",
    "    color_mode=\"rgb\",\n",
    "    classes=names\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_batch = train_batches[0]\n",
    "print(train_batch[0].shape)\n",
    "print(train_batch[1])\n",
    "test_batch = test_batches[0]\n",
    "print(test_batch[0].shape)\n",
    "print(test_batch[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show(batch, pred_labels=None):\n",
    "    plt.figure(figsize=(10,10))\n",
    "    for i in range(4):\n",
    "        plt.subplot(2,2,i+1)\n",
    "        plt.xticks([])\n",
    "        plt.yticks([])\n",
    "        plt.grid(False)\n",
    "        plt.imshow(batch[0][i], cmap=plt.cm.binary)\n",
    "        # The CIFAR labels happen to be arrays, \n",
    "        # which is why you need the extra index\n",
    "        lbl = names[int(batch[1][i])]\n",
    "        if pred_labels is not None:\n",
    "            lbl += \"/ Pred:\" + names[int(pred_labels[i])]\n",
    "        plt.xlabel(lbl)\n",
    "    plt.show()\n",
    "\n",
    "show(test_batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(train_batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = keras.models.Sequential()\n",
    "model.add(layers.Conv2D(32, (3,3), strides=(1,1), padding=\"valid\", activation='relu', input_shape=(256, 256,3)))\n",
    "model.add(layers.MaxPool2D((2,2)))\n",
    "model.add(layers.Conv2D(64, 3, activation='relu'))\n",
    "model.add(layers.MaxPool2D((2,2)))\n",
    "model.add(layers.Flatten())\n",
    "model.add(layers.Dense(64, activation='relu'))\n",
    "model.add(layers.Dense(5))\n",
    "print(model.summary())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# loss and optimizer\n",
    "loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "optim = keras.optimizers.Adam(lr=0.001)\n",
    "metrics = [\"accuracy\"]\n",
    "\n",
    "model.compile(optimizer=optim, loss=loss, metrics=metrics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# training\n",
    "epochs = 30\n",
    "\n",
    "# callbacks\n",
    "early_stopping = keras.callbacks.EarlyStopping(\n",
    "    monitor=\"val_loss\",\n",
    "    patience=5,\n",
    "    verbose=2\n",
    ")\n",
    "\n",
    "history = model.fit(train_batches, validation_data=val_batches,\n",
    "                    callbacks=[early_stopping],\n",
    "                      epochs=epochs, verbose=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save(\"lego_model.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot loss and acc\n",
    "plt.figure(figsize=(16, 6))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(history.history['loss'], label='train loss')\n",
    "plt.plot(history.history['val_loss'], label='valid loss')\n",
    "plt.grid()\n",
    "plt.legend(fontsize=15)\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(history.history['accuracy'], label='train acc')\n",
    "plt.plot(history.history['val_accuracy'], label='valid acc')\n",
    "plt.grid()\n",
    "plt.legend(fontsize=15);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluate on test data\n",
    "model.evaluate(test_batches, verbose=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make some predictions\n",
    "predictions = model.predict(test_batches)\n",
    "predictions = tf.nn.softmax(predictions)\n",
    "labels = np.argmax(predictions, axis=1)\n",
    "\n",
    "print(test_batches[0][1])\n",
    "print(labels[0:4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(test_batches[0], labels[0:4])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transfer Learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vgg_model = tf.keras.applications.vgg16.VGG16()\n",
    "print(type(vgg_model))\n",
    "vgg_model.summary()\n",
    "\n",
    "# try out different ones, e.g. MobileNetV2\n",
    "#tl_model = tf.keras.applications.MobileNetV2()\n",
    "#print(type(tl_model))\n",
    "#tl_model.summary()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert to Sequential model, omit the last layer\n",
    "# this works with VGG16 because the structure is linear\n",
    "model = keras.models.Sequential()\n",
    "for layer in vgg_model.layers[0:-1]:\n",
    "    model.add(layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set trainable=False for all layers\n",
    "# we don't want to train them again\n",
    "for layer in model.layers:\n",
    "    layer.trainable = False\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add a last classification layer for our use case with 5 classes\n",
    "model.add(layers.Dense(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# loss and optimizer\n",
    "loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "optim = keras.optimizers.Adam(lr=0.001)\n",
    "metrics = [\"accuracy\"]\n",
    "\n",
    "model.compile(optimizer=optim, loss=loss, metrics=metrics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the preprocessing function of this model\n",
    "preprocess_input = tf.keras.applications.vgg16.preprocess_input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate batches of tensor image data with real-time data augmentation.\n",
    "\n",
    "train_gen = keras.preprocessing.image.ImageDataGenerator(preprocessing_function=preprocess_input)\n",
    "valid_gen = keras.preprocessing.image.ImageDataGenerator(preprocessing_function=preprocess_input)\n",
    "test_gen = keras.preprocessing.image.ImageDataGenerator(preprocessing_function=preprocess_input)\n",
    "\n",
    "train_batches = train_gen.flow_from_directory(\n",
    "    'lego/star-wars-images/train',\n",
    "    target_size=(224, 224),\n",
    "    class_mode='sparse',\n",
    "    batch_size=4,\n",
    "    shuffle=True,\n",
    "    color_mode=\"rgb\",\n",
    "    classes=names   \n",
    ")\n",
    "\n",
    "val_batches = valid_gen.flow_from_directory(\n",
    "    'lego/star-wars-images/val',\n",
    "    target_size=(224, 224),\n",
    "    class_mode='sparse',\n",
    "    batch_size=4,\n",
    "    shuffle=True,\n",
    "    color_mode=\"rgb\",\n",
    "    classes=names\n",
    ")\n",
    "\n",
    "test_batches = test_gen.flow_from_directory(\n",
    "    'lego/star-wars-images/test',\n",
    "    target_size=(224, 224),\n",
    "    class_mode='sparse',\n",
    "    batch_size=4,\n",
    "    shuffle=False,\n",
    "    color_mode=\"rgb\",\n",
    "    classes=names\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 30\n",
    "\n",
    "early_stopping = keras.callbacks.EarlyStopping(\n",
    "    monitor=\"val_loss\",\n",
    "    patience=5,\n",
    "    verbose=2\n",
    ")\n",
    "\n",
    "model.fit(train_batches, validation_data=val_batches,\n",
    "          callbacks=[early_stopping],\n",
    "          epochs=epochs, verbose=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.evaluate(test_batches, verbose=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tf",
   "language": "python",
   "name": "tf"
  },
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}