Smart_Bin––Transfer_Learning_Model.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PEw1MaDLCQzq"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "#import keras api needed to implement deep learning techiques\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras.models import Sequential\n",
        "from tensorflow.keras.layers import Activation, Dense, BatchNormalization, Conv2D, MaxPool2D, GlobalAveragePooling2D, Dropout\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "from tensorflow.keras.losses import CategoricalCrossentropy\n",
        "from tensorflow.keras.preprocessing.image import ImageDataGenerator\n",
        "\n",
        "#Allow charts and graphics to display right below the page of browser setup\n",
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Load image data\n",
        "\n",
        "We load our collected image data representing five different classes, plastic, metal, glass, and other. Since we are using transfer learning with imagenet, we load it with preprocessing for this network using ImageDataGenerator"
      ],
      "metadata": {
        "id": "dzwGaOgoCd_A"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Remember to mount drive\n",
        "image_filepath = '/drive/MyDrive/smart_bin_image_data/' # Replace with relevant filepath\n",
        "\n",
        "train_batch = ImageDataGenerator(\n",
        "    preprocessing_function=tf.keras.applications.vgg16.preprocess_input,\n",
        "    shear_range=0.2,\n",
        "    zoom_range=0.2,\n",
        "    width_shit_range=0.2,\n",
        "    height_shift_range=0.2,\n",
        "    validation_split=0.1).flow_from_directory(\n",
        "        directory=image_filepath,\n",
        "        target_size=(500, 500),\n",
        "        classes=['paper', 'plastic', 'glass', 'metal', 'other'],\n",
        "        batch_size=64,\n",
        "    )"
      ],
      "metadata": {
        "id": "-3u6LnpdCcAc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Defining the Architecture\n",
        "\n",
        "We now define the model that we use to train over the images. Since this is transfer learning, we use imagenet and finetune only the last few layers. Note that we had to adjust the "
      ],
      "metadata": {
        "id": "SMF-bgzOEqKg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "DIMS = 500\n",
        "\n",
        "# First import in imagenet model\n",
        "imagenet = tf.keras.applications.VGG16(input_shape=DIMS, \n",
        "                                        include_top=False,\n",
        "                                        weights='imagenet')\n",
        "imagenet.trainable = False # Make sure not to train this part of the model\n",
        "\n",
        "# Now we can define our own model on top of this\n",
        "model = Sequential()\n",
        "model.add(imagenet)\n",
        "model.add(Dense(units=128, activation='relu'))\n",
        "model.add(Dropout(0.2))\n",
        "model.add(Dense(units=64, activation='relu'))\n",
        "model.add(BatchNormalization())\n",
        "model.add(Dropout(0.2))\n",
        "model.add(Dense(units=5, activation='softmax')) # Final layer with softmax for class probabilities\n",
        "\n",
        "model.summary()"
      ],
      "metadata": {
        "id": "9NHpDkZ3Eptj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "First, we train the layers we attached to the model"
      ],
      "metadata": {
        "id": "7Gsu99I_Gwql"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model.compile(optimizer=Adam(learning_rate=0.0001), loss='categorical_crossentropy', metrics=['accuracy'])\n",
        "model.fit(x=train_batches, epochs=10)"
      ],
      "metadata": {
        "id": "taMnJN1fGdXd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Then, we finetune the model for just a few epochs"
      ],
      "metadata": {
        "id": "OQkSMCxSHLo-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "imagenet.trainable=True\n",
        "model.compile(optimizer=Adam(learning_rate=0.0001), loss='categorical_crossentropy', metrics=['accuracy'])\n",
        "model.fit(x=train_batches, epochs=3)"
      ],
      "metadata": {
        "id": "i5eVvmz7HODb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "save_path = '/drive/MyDrive/smart_bin_models/imagenet_transfer.hf5'\n",
        "model.save(save_path)"
      ],
      "metadata": {
        "id": "k_szrKG_umwE"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}