Data Training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cf5c2594",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import cv2\n",
    "import json\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3e0bef07",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_image(x):\n",
    "    byte_img = tf.io.read_file(x)\n",
    "    img = tf.io.decode_jpeg(byte_img)\n",
    "    return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b6088983",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_images = tf.data.Dataset.list_files('aug_data\\\\train\\\\images\\\\*.jpg' , shuffle=False)\n",
    "train_images = train_images.map(load_image)\n",
    "train_images = train_images.map(lambda x: tf.image.resize(x , (120,120)))\n",
    "train_images = train_images.map(lambda x: x/255)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bb9d1057",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_images = tf.data.Dataset.list_files('aug_data\\\\test\\\\images\\\\*.jpg' , shuffle=False)\n",
    "test_images = test_images.map(load_image)\n",
    "test_images = test_images.map(lambda x: tf.image.resize(x , (120,120)))\n",
    "test_images = test_images.map(lambda x: x/255)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d0eda875",
   "metadata": {},
   "outputs": [],
   "source": [
    "val_images = tf.data.Dataset.list_files('aug_data\\\\val\\\\images\\\\*.jpg' , shuffle=False)\n",
    "val_images = val_images.map(load_image)\n",
    "val_images = val_images.map(lambda x: tf.image.resize(x , (120,120)))\n",
    "val_images = val_images.map(lambda x: x/255)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "888227fb",
   "metadata": {},
   "source": [
    "##### Prepare label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2f13d575",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_labels(label_path):\n",
    "    with open(label_path.numpy() , 'r' , encoding='utf-8') as f:\n",
    "        label = json.load(f)\n",
    "        \n",
    "    return [label['class']] , label['bbox']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f54ed1c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_labels = tf.data.Dataset.list_files('aug_data\\\\train\\\\labels\\\\*.json' , shuffle=False)\n",
    "train_labels = train_labels.map(lambda x:tf.py_function(load_labels , [x] , [tf.uint8 , tf.float16]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9eb38b1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_labels = tf.data.Dataset.list_files('aug_data\\\\test\\\\labels\\\\*.json' , shuffle=False)\n",
    "test_labels = test_labels.map(lambda x:tf.py_function(load_labels , [x] , [tf.uint8 , tf.float16]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b247e6f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "val_labels = tf.data.Dataset.list_files('aug_data\\\\val\\\\labels\\\\*.json' , shuffle=False)\n",
    "val_labels = val_labels.map(lambda x:tf.py_function(load_labels , [x] , [tf.uint8 , tf.float16]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79e348ad",
   "metadata": {},
   "source": [
    "#### Create final dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0a5b3e0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "train = tf.data.Dataset.zip((train_images , train_labels))\n",
    "train = train.shuffle(5000)\n",
    "train = train.batch(8)\n",
    "train = train.prefetch(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d4cec774",
   "metadata": {},
   "outputs": [],
   "source": [
    "test = tf.data.Dataset.zip((train_images , train_labels))\n",
    "test = test.shuffle(5000)\n",
    "test = test.batch(8)\n",
    "test = test.prefetch(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b6cf4e5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "val = tf.data.Dataset.zip((train_images , train_labels))\n",
    "val = val.shuffle(5000)\n",
    "val = val.batch(8)\n",
    "val = val.prefetch(4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9024eb2",
   "metadata": {},
   "source": [
    "# Strat with model training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "bcee9aed",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.layers import Input, Conv2D, Dense, GlobalMaxPooling2D\n",
    "from tensorflow.keras.applications import VGG16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "30ac61d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "vgg = VGG16(include_top=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "f6ae1e34",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_model():\n",
    "    input_layer = Input(shape=(120 , 120 , 3))\n",
    "    \n",
    "    vgg = VGG16(include_top=False)(input_layer)\n",
    "    \n",
    "    # Classification model\n",
    "    f1 = GlobalMaxPooling2D()(vgg)\n",
    "    class1 = Dense(2048 , activation='relu')(f1)\n",
    "    class2 = Dense(1 , activation='sigmoid')(class1)\n",
    "    \n",
    "    # Bounding box model\n",
    "    f2 = GlobalMaxPooling2D()(vgg)\n",
    "    regress1 = Dense(2048 , activation='relu')(f2)\n",
    "    regress2 = Dense(4 , activation='sigmoid')(regress1)\n",
    "    \n",
    "    facetracker = Model(inputs=input_layer , outputs=[class2 , regress2])\n",
    "    return facetracker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "250050a2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"vgg16\"\n",
      "_________________________________________________________________\n",
      " Layer (type)                Output Shape              Param #   \n",
      "=================================================================\n",
      " input_1 (InputLayer)        [(None, None, None, 3)]   0         \n",
      "                                                                 \n",
      " block1_conv1 (Conv2D)       (None, None, None, 64)    1792      \n",
      "                                                                 \n",
      " block1_conv2 (Conv2D)       (None, None, None, 64)    36928     \n",
      "                                                                 \n",
      " block1_pool (MaxPooling2D)  (None, None, None, 64)    0         \n",
      "                                                                 \n",
      " block2_conv1 (Conv2D)       (None, None, None, 128)   73856     \n",
      "                                                                 \n",
      " block2_conv2 (Conv2D)       (None, None, None, 128)   147584    \n",
      "                                                                 \n",
      " block2_pool (MaxPooling2D)  (None, None, None, 128)   0         \n",
      "                                                                 \n",
      " block3_conv1 (Conv2D)       (None, None, None, 256)   295168    \n",
      "                                                                 \n",
      " block3_conv2 (Conv2D)       (None, None, None, 256)   590080    \n",
      "                                                                 \n",
      " block3_conv3 (Conv2D)       (None, None, None, 256)   590080    \n",
      "                                                                 \n",
      " block3_pool (MaxPooling2D)  (None, None, None, 256)   0         \n",
      "                                                                 \n",
      " block4_conv1 (Conv2D)       (None, None, None, 512)   1180160   \n",
      "                                                                 \n",
      " block4_conv2 (Conv2D)       (None, None, None, 512)   2359808   \n",
      "                                                                 \n",
      " block4_conv3 (Conv2D)       (None, None, None, 512)   2359808   \n",
      "                                                                 \n",
      " block4_pool (MaxPooling2D)  (None, None, None, 512)   0         \n",
      "                                                                 \n",
      " block5_conv1 (Conv2D)       (None, None, None, 512)   2359808   \n",
      "                                                                 \n",
      " block5_conv2 (Conv2D)       (None, None, None, 512)   2359808   \n",
      "                                                                 \n",
      " block5_conv3 (Conv2D)       (None, None, None, 512)   2359808   \n",
      "                                                                 \n",
      " block5_pool (MaxPooling2D)  (None, None, None, 512)   0         \n",
      "                                                                 \n",
      "=================================================================\n",
      "Total params: 14,714,688\n",
      "Trainable params: 14,714,688\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "vgg.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "7accb347",
   "metadata": {},
   "outputs": [],
   "source": [
    "facetracker = build_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "e440e31e",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                   Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      " input_2 (InputLayer)           [(None, 120, 120, 3  0           []                               \n",
      "                                )]                                                                \n",
      "                                                                                                  \n",
      " vgg16 (Functional)             (None, None, None,   14714688    ['input_2[0][0]']                \n",
      "                                512)                                                              \n",
      "                                                                                                  \n",
      " global_max_pooling2d (GlobalMa  (None, 512)         0           ['vgg16[0][0]']                  \n",
      " xPooling2D)                                                                                      \n",
      "                                                                                                  \n",
      " global_max_pooling2d_1 (Global  (None, 512)         0           ['vgg16[0][0]']                  \n",
      " MaxPooling2D)                                                                                    \n",
      "                                                                                                  \n",
      " dense (Dense)                  (None, 2048)         1050624     ['global_max_pooling2d[0][0]']   \n",
      "                                                                                                  \n",
      " dense_2 (Dense)                (None, 2048)         1050624     ['global_max_pooling2d_1[0][0]'] \n",
      "                                                                                                  \n",
      " dense_1 (Dense)                (None, 1)            2049        ['dense[0][0]']                  \n",
      "                                                                                                  \n",
      " dense_3 (Dense)                (None, 4)            8196        ['dense_2[0][0]']                \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 16,826,181\n",
      "Trainable params: 16,826,181\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "facetracker.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "8be7d14c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Since we've a classification model and a regression model , we need seperate loss function for both.\n",
    "batches_per_epoch = len(train)\n",
    "lr_decay = (1./0.75-1)/batches_per_epoch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "d7f5dfab",
   "metadata": {},
   "outputs": [],
   "source": [
    "opt = tf.keras.optimizers.Adam(learning_rate=0.001 , decay=lr_decay)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "999ce27b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def localization_loss(y_true, yhat):            \n",
    "    delta_coord = tf.reduce_sum(tf.square(y_true[:,:2] - yhat[:,:2]))\n",
    "                  \n",
    "    h_true = y_true[:,3] - y_true[:,1] \n",
    "    w_true = y_true[:,2] - y_true[:,0] \n",
    "\n",
    "    h_pred = yhat[:,3] - yhat[:,1] \n",
    "    w_pred = yhat[:,2] - yhat[:,0] \n",
    "    \n",
    "    delta_size = tf.reduce_sum(tf.square(w_true - w_pred) + tf.square(h_true-h_pred))\n",
    "    \n",
    "    return delta_coord + delta_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "e6dae174",
   "metadata": {},
   "outputs": [],
   "source": [
    "classloss = tf.keras.losses.BinaryCrossentropy()\n",
    "regressloss = localization_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "6da2271f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FaceTracker(Model): \n",
    "    def __init__(self, eyetracker,  **kwargs): \n",
    "        super().__init__(**kwargs)\n",
    "        self.model = eyetracker\n",
    "\n",
    "    def compile(self, opt, classloss, localizationloss, **kwargs):\n",
    "        super().compile(**kwargs)\n",
    "        self.closs = classloss\n",
    "        self.lloss = localizationloss\n",
    "        self.opt = opt\n",
    "    \n",
    "    def train_step(self, batch, **kwargs): \n",
    "        \n",
    "        X, y = batch\n",
    "        \n",
    "        with tf.GradientTape() as tape: \n",
    "            classes, coords = self.model(X, training=True)\n",
    "            \n",
    "            batch_classloss = self.closs(y[0], classes)\n",
    "            batch_localizationloss = self.lloss(tf.cast(y[1], tf.float32), coords)\n",
    "            \n",
    "            total_loss = batch_localizationloss+0.5*batch_classloss\n",
    "            \n",
    "            grad = tape.gradient(total_loss, self.model.trainable_variables)\n",
    "        \n",
    "        opt.apply_gradients(zip(grad, self.model.trainable_variables))\n",
    "        \n",
    "        return {\"total_loss\":total_loss, \"class_loss\":batch_classloss, \"regress_loss\":batch_localizationloss}\n",
    "    \n",
    "    def test_step(self, batch, **kwargs): \n",
    "        X, y = batch\n",
    "        \n",
    "        classes, coords = self.model(X, training=False)\n",
    "        \n",
    "        batch_classloss = self.closs(y[0], classes)\n",
    "        batch_localizationloss = self.lloss(tf.cast(y[1], tf.float32), coords)\n",
    "        total_loss = batch_localizationloss+0.5*batch_classloss\n",
    "        \n",
    "        return {\"total_loss\":total_loss, \"class_loss\":batch_classloss, \"regress_loss\":batch_localizationloss}\n",
    "        \n",
    "    def call(self, X, **kwargs): \n",
    "        return self.model(X, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "cdceb7bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = FaceTracker(facetracker)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "0b836ced",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.compile(opt , classloss , regressloss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "8003f31e",
   "metadata": {},
   "outputs": [],
   "source": [
    "logdir = 'logs'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "dda94d89",
   "metadata": {},
   "outputs": [],
   "source": [
    "tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=logdir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2666ce9e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/5\n",
      " 91/180 [==============>...............] - ETA: 10:39 - total_loss: 4.7546 - class_loss: 0.0083 - regress_loss: 4.7504"
     ]
    }
   ],
   "source": [
    "hist = model.fit(train , epochs=40 , validation_data=val , callbacks=[tensorboard_callback])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8d685ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "hist.history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a9f7894",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(ncols=3, figsize=(20,5))\n",
    "\n",
    "ax[0].plot(hist.history['total_loss'], color='teal', label='loss')\n",
    "ax[0].plot(hist.history['val_total_loss'], color='orange', label='val loss')\n",
    "ax[0].title.set_text('Loss')\n",
    "ax[0].legend()\n",
    "\n",
    "ax[1].plot(hist.history['class_loss'], color='teal', label='class loss')\n",
    "ax[1].plot(hist.history['val_class_loss'], color='orange', label='val class loss')\n",
    "ax[1].title.set_text('Classification Loss')\n",
    "ax[1].legend()\n",
    "\n",
    "ax[2].plot(hist.history['regress_loss'], color='teal', label='regress loss')\n",
    "ax[2].plot(hist.history['val_regress_loss'], color='orange', label='val regress loss')\n",
    "ax[2].title.set_text('Regression Loss')\n",
    "ax[2].legend()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50cf7f94",
   "metadata": {},
   "source": [
    "##### Testing data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c43552ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data = test.as_numpy_iterator()\n",
    "test_sample = test_data.next()\n",
    "yhat = facetracker.predict(test_sample[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d02e29b",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig , ax = plt.subplots(ncols=4 , figsize=(20,20))\n",
    "for idx in range(4):\n",
    "  sample_image = test_sample[0][idx]\n",
    "  sample_coords = yhat[1][idx]\n",
    "\n",
    "  if yhat[0][idx]>0.5:\n",
    "        cv2.rectangle(sample_image, \n",
    "                      tuple(np.multiply(sample_coords[:2], [120,120]).astype(int)),\n",
    "                      tuple(np.multiply(sample_coords[2:], [120,120]).astype(int)), \n",
    "                            (255,0,0), 2)\n",
    "  ax[idx].imshow(sample_image)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f81d942",
   "metadata": {},
   "source": [
    "##### Save model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76dcf3aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.models import load_model\n",
    "facetracker.save('facetracker.h5')"
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}