hexbug_head_detection_resnet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "275adc06",
   "metadata": {},
   "source": [
    "# Training Resnet50 for Head's Position Detection\n",
    "After preparing the `margined_cropped_images`, and `data.csv`, now we are able to feed them to a model to predict Hexbugs' heads' position.\n",
    "The approach that we managed to work with was using ***transfer learning*** to build a better model. For this, we used pre-trained `Resnet50` model, which is trained on `CoCo` dataset. This helps us to extract feature better. Hoever, top layers of the `resnet` model are used for classification tasks. To modify the model to suit our problem, we removed tha last `Dense` layer and added two other `Dense` layers to build a regressor with only ***two*** outputs for `(x, y)`.\n",
    "After modifying the model, we can train it over our dataset, and evaluate its performance on the validation set.\n",
    "At last, we save the model for further use cases."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5cbfa65",
   "metadata": {},
   "source": [
    "### Essential Imports\n",
    "\n",
    "* `cv2` is required to read and write images.\n",
    "* `json` is used to open `json` files.\n",
    "* `random` is used to pick random samples for model evaluation.\n",
    "* `logging` is used to ignore `warning` messages. (not mandatory)\n",
    "* `numoy` is cruicial to work with images.\n",
    "* `pandas` is used to read the `data.csv` file.\n",
    "* `tensorflow` is used to create model. it's used as background for `Keras`.\n",
    "* `sklearn` is used to split dataset to `train` and `test` sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "058346ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import json\n",
    "import random\n",
    "import logging\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "from sklearn.model_selection import train_test_split\n",
    "from tensorflow.keras import layers, Model\n",
    "from tensorflow.keras.models import Sequential\n",
    "from tensorflow.keras.utils import load_img\n",
    "from tensorflow.keras.applications.resnet50 import ResNet50\n",
    "from tensorflow.keras.applications import ResNet101, DenseNet121, DenseNet169, DenseNet201\n",
    "from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout, BatchNormalization\n",
    "\n",
    "from IPython.display import clear_output\n",
    "\n",
    "tf.get_logger().setLevel(logging.ERROR)\n",
    "clear_output()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f247aa3",
   "metadata": {},
   "source": [
    "### Reading data from `data.csv`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "16f1de72",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>CroppedHexBugCoordinationX</th>\n",
       "      <th>CroppedHexBugCoordinationY</th>\n",
       "      <th>OriginalBoxCoordinationX1</th>\n",
       "      <th>OriginalBoxCoordinationY1</th>\n",
       "      <th>Path</th>\n",
       "      <th>ID</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>5011</th>\n",
       "      <td>48</td>\n",
       "      <td>36</td>\n",
       "      <td>218</td>\n",
       "      <td>18</td>\n",
       "      <td>cropped_bugs/training018/frame0.jpg</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2805</th>\n",
       "      <td>125</td>\n",
       "      <td>132</td>\n",
       "      <td>481</td>\n",
       "      <td>381</td>\n",
       "      <td>cropped_bugs/training044/frame0.jpg</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1100</th>\n",
       "      <td>269</td>\n",
       "      <td>119</td>\n",
       "      <td>905</td>\n",
       "      <td>708</td>\n",
       "      <td>cropped_bugs/training083/frame0.jpg</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4911</th>\n",
       "      <td>22</td>\n",
       "      <td>269</td>\n",
       "      <td>536</td>\n",
       "      <td>745</td>\n",
       "      <td>cropped_bugs/training021/frame0.jpg</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3710</th>\n",
       "      <td>222</td>\n",
       "      <td>206</td>\n",
       "      <td>760</td>\n",
       "      <td>1339</td>\n",
       "      <td>cropped_bugs/training050/frame0.jpg</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3578</th>\n",
       "      <td>20</td>\n",
       "      <td>36</td>\n",
       "      <td>318</td>\n",
       "      <td>82</td>\n",
       "      <td>cropped_bugs/training022/frame100.jpg</td>\n",
       "      <td>100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1170</th>\n",
       "      <td>19</td>\n",
       "      <td>80</td>\n",
       "      <td>43</td>\n",
       "      <td>11</td>\n",
       "      <td>cropped_bugs/training041/frame100.jpg</td>\n",
       "      <td>100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3679</th>\n",
       "      <td>30</td>\n",
       "      <td>82</td>\n",
       "      <td>105</td>\n",
       "      <td>1466</td>\n",
       "      <td>cropped_bugs/training050/frame100.jpg</td>\n",
       "      <td>100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3871</th>\n",
       "      <td>31</td>\n",
       "      <td>73</td>\n",
       "      <td>520</td>\n",
       "      <td>400</td>\n",
       "      <td>cropped_bugs/training025/frame100.jpg</td>\n",
       "      <td>100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1271</th>\n",
       "      <td>244</td>\n",
       "      <td>34</td>\n",
       "      <td>332</td>\n",
       "      <td>630</td>\n",
       "      <td>cropped_bugs/training048/frame100.jpg</td>\n",
       "      <td>100</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5012 rows × 6 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      CroppedHexBugCoordinationX  CroppedHexBugCoordinationY  \\\n",
       "5011                          48                          36   \n",
       "2805                         125                         132   \n",
       "1100                         269                         119   \n",
       "4911                          22                         269   \n",
       "3710                         222                         206   \n",
       "...                          ...                         ...   \n",
       "3578                          20                          36   \n",
       "1170                          19                          80   \n",
       "3679                          30                          82   \n",
       "3871                          31                          73   \n",
       "1271                         244                          34   \n",
       "\n",
       "      OriginalBoxCoordinationX1  OriginalBoxCoordinationY1  \\\n",
       "5011                        218                         18   \n",
       "2805                        481                        381   \n",
       "1100                        905                        708   \n",
       "4911                        536                        745   \n",
       "3710                        760                       1339   \n",
       "...                         ...                        ...   \n",
       "3578                        318                         82   \n",
       "1170                         43                         11   \n",
       "3679                        105                       1466   \n",
       "3871                        520                        400   \n",
       "1271                        332                        630   \n",
       "\n",
       "                                       Path   ID  \n",
       "5011    cropped_bugs/training018/frame0.jpg    0  \n",
       "2805    cropped_bugs/training044/frame0.jpg    0  \n",
       "1100    cropped_bugs/training083/frame0.jpg    0  \n",
       "4911    cropped_bugs/training021/frame0.jpg    0  \n",
       "3710    cropped_bugs/training050/frame0.jpg    0  \n",
       "...                                     ...  ...  \n",
       "3578  cropped_bugs/training022/frame100.jpg  100  \n",
       "1170  cropped_bugs/training041/frame100.jpg  100  \n",
       "3679  cropped_bugs/training050/frame100.jpg  100  \n",
       "3871  cropped_bugs/training025/frame100.jpg  100  \n",
       "1271  cropped_bugs/training048/frame100.jpg  100  \n",
       "\n",
       "[5012 rows x 6 columns]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = pd.read_csv('data.csv').drop(['Unnamed: 0'], axis=1).sort_values('ID', ascending=True)\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "481aa283",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(5012, 300, 300, 3)\n",
      "(5012, 2)\n"
     ]
    }
   ],
   "source": [
    "# Load the images\n",
    "images = []\n",
    "y = []\n",
    "\n",
    "# Iterates over the `data.csv` file's rows to read images and annotatopns.\n",
    "for idx, row in data.iterrows():\n",
    "#     img_path = f\"{row.Path.replace('cropped_bugs', 'image_sharpning/sharpened_images')}\"\n",
    "    img_path = f\"{row.Path}\"\n",
    "    annotations = [int(row.CroppedHexBugCoordinationX), int(row.CroppedHexBugCoordinationY)]\n",
    "    # Reads the image\n",
    "    img = load_img(img_path, color_mode='rgb')\n",
    "    # Converts the image to a numpy array\n",
    "    np_img = np.array(img)\n",
    "    \n",
    "#     extended_np_img = np.expand_dims(np_img, axis=2)\n",
    "    # Adds the image to a list of images.\n",
    "    images.append(np_img)\n",
    "    # Adds the annotations to y.\n",
    "    y.append(annotations)\n",
    "    \n",
    "X = np.array(images)\n",
    "y = np.array(y)\n",
    "\n",
    "print(X.shape)\n",
    "print(y.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "65f0e5a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Splits the dataset to `train` and `test` sets with the given ratio.\n",
    "# Also, shuffles the data to prevent feeding the model same data over different runs.\n",
    "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "495ae437",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X train shape: (3358, 300, 300, 3)\n",
      "X test shape: (1654, 300, 300, 3)\n",
      "y train shape: (3358, 2)\n",
      "y test shape: (1654, 2)\n"
     ]
    }
   ],
   "source": [
    "print(f'X train shape: {X_train.shape}')\n",
    "print(f'X test shape: {X_test.shape}')\n",
    "print(f'y train shape: {y_train.shape}')\n",
    "print(f'y test shape: {y_test.shape}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b263b73e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = tf.keras.models.load_model(\"resnet50_trained_model_data/v1.5/checkpoint_callback\")\n",
    "clear_output()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d18e66d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dense_model = DenseNet169(\n",
    "#     include_top=False,\n",
    "#     weights=\"imagenet\",\n",
    "#     input_shape=X_train[0].shape)\n",
    "\n",
    "# model = Sequential(dense_model, name='DenseNet169')\n",
    "\n",
    "# x = model.output\n",
    "# x = Flatten()(x)\n",
    "# x = Dense(16, activation='relu')(x)\n",
    "# x = Dense(8, activation='relu')(x)\n",
    "\n",
    "# # The last layer should be the layer that computes the (x, y) coordinations.\n",
    "# output_layer = Dense(2, activation='linear')(x)\n",
    "\n",
    "# # Create the model based on the input and output layer.\n",
    "# model = Model(inputs=model.input, outputs=output_layer)\n",
    "\n",
    "# clear_output()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7c3140ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "# vgg_model = tf.keras.applications.VGG19(\n",
    "#     include_top=False,\n",
    "#     weights=\"imagenet\",\n",
    "#     input_shape=X_train[0].shape,\n",
    "# )\n",
    "\n",
    "# for layer in vgg_model.layers:\n",
    "#     layer.trainable = True\n",
    "\n",
    "# # Adds two dense layers for regression problem.\n",
    "# # It basically uses the output of the `resnet50` model to feed the dense layers.\n",
    "# x = vgg_model.output\n",
    "# x = Flatten()(x)\n",
    "# x = Dense(128, activation='relu')(x)\n",
    "# x = Dense(64, activation='relu')(x)\n",
    "\n",
    "# # The last layer should be the layer that computes the (x, y) coordinations.\n",
    "# output_layer = Dense(2, activation='linear')(x)\n",
    "\n",
    "# # Create the model based on the input and output layer.\n",
    "# model = Model(inputs=vgg_model.input, outputs=output_layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ac597437",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model = Sequential()\n",
    "# model.add(Conv2D(64, (3, 3), activation='relu', input_shape=X_train[0].shape))\n",
    "# # model.add(Dropout(0.25))\n",
    "# # model.add(MaxPooling2D((2, 2)))\n",
    "# model.add(Conv2D(128, (3, 3), activation='relu', kernel_regularizer='L2'))\n",
    "# # model.add(Dropout(0.25))\n",
    "# model.add(MaxPooling2D((2, 2)))\n",
    "\n",
    "# model.add(Conv2D(256, (3, 3), activation='relu', kernel_regularizer='L2'))\n",
    "# # model.add(Dropout(0.25))\n",
    "# model.add(MaxPooling2D((2, 2)))\n",
    "\n",
    "# model.add(Conv2D(512, (3, 3), activation='relu', kernel_regularizer='L2'))\n",
    "# # model.add(Dropout(0.25))\n",
    "# # model.add(MaxPooling2D((2, 2)))\n",
    "\n",
    "# model.add(Conv2D(256, (3, 3), activation='relu', kernel_regularizer='L2'))\n",
    "# # model.add(Dropout(0.25))\n",
    "# model.add(MaxPooling2D((2, 2)))\n",
    "\n",
    "# model.add(Conv2D(128, (3, 3), activation='relu', kernel_regularizer='L2'))\n",
    "# # model.add(Dropout(0.25))\n",
    "# model.add(MaxPooling2D((2, 2)))\n",
    "# model.add(Conv2D(64, (3, 3), activation='relu', kernel_regularizer='L2'))\n",
    "# # model.add(Dropout(0.25))\n",
    "# model.add(MaxPooling2D((2, 2)))\n",
    "# model.add(Conv2D(32, (3, 3), activation='relu', kernel_regularizer='L2'))\n",
    "# # model.add(Dropout(0.25))\n",
    "# model.add(MaxPooling2D((2, 2)))\n",
    "# model.add(Flatten())\n",
    "# model.add(Dense(16, activation='relu'))\n",
    "# model.add(Dense(2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "603e0fd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Creates a pre-trained model instance from class `ResNet50`.\n",
    "# # This model doesn't include the last classification dense layer.\n",
    "# # The given `input_shape` should be the same as our dataset images.\n",
    "# resnet50 = ResNet50(weights='imagenet', include_top=False, input_shape=(X_train[0].shape))\n",
    "\n",
    "# # Iterates over each layer in the model and makes them to be trainable.\n",
    "# for layer in resnet50.layers:\n",
    "#     layer.trainable = True\n",
    "\n",
    "# # Adds two dense layers for regression problem.\n",
    "# # It basically uses the output of the `resnet50` model to feed the dense layers.\n",
    "# x = resnet50.output\n",
    "# x = Flatten()(x)\n",
    "# x = Dense(32, activation='relu')(x)\n",
    "# x = Dense(16, activation='relu')(x)\n",
    "\n",
    "# # The last layer should be the layer that computes the (x, y) coordinations.\n",
    "# output_layer = Dense(2, activation='linear')(x)\n",
    "\n",
    "# # Create the model based on the input and output layer.\n",
    "# model = Model(inputs=resnet50.input, outputs=output_layer)\n",
    "\n",
    "# clear_output()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0bfa0a5",
   "metadata": {},
   "source": [
    "### Visualization\n",
    "For better visualization, we use [tensorboard](https://www.tensorflow.org/tensorboard/get_started). It helps us to plot the figures real-time and analyze the model performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f8ad8d65",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/usr/bin/zsh: /home/farzam/miniconda3/envs/tf/lib/libtinfo.so.6: no version information available (required by /usr/bin/zsh)\r\n"
     ]
    }
   ],
   "source": [
    "# Load the TensorBoard notebook extension\n",
    "%load_ext tensorboard\n",
    "\n",
    "# We need datetime to name our checkpoints.\n",
    "import datetime\n",
    "\n",
    "# Clear any logs from previous runs\n",
    "!rm -rf ./logs/\n",
    "\n",
    "log_dir = \"logs/fit/\" + datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "\n",
    "# Creates a callback, which then will be called by the model during training process.\n",
    "tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6374b0c",
   "metadata": {},
   "source": [
    "### Saving the best model\n",
    "During the training process, we need to track the model's performance and make sure that at the end of the training process, we have the best model.\n",
    "To do so, we used a custom `callback` that checks if the `val_loss` has been decreased over the past epoch. In that case, it will simply save the model to the provided directory, which can be used to load and evaluate later.\n",
    "This call back is then called by the model during the training process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a682ef32",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specifies the checkpoint directory\n",
    "checkpoint_filepath = 'checkpoint_callback'\n",
    "\n",
    "# Creates an instance from the ModelCheckpoint class.\n",
    "model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(\n",
    "    filepath=checkpoint_filepath,\n",
    "    save_weights_only=False,\n",
    "    monitor='val_loss',\n",
    "    mode='min',\n",
    "    save_best_only=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e185478e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Since we are dealing with regression problem, it's recommended to use `mse`.\n",
    "# Also we use `Adam` as our optimizer.\n",
    "model.compile(loss='mse', optimizer='adam')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1369f20d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "_________________________________________________________________\n",
      " Layer (type)                Output Shape              Param #   \n",
      "=================================================================\n",
      " densenet169_input (InputLay  [(None, 300, 300, 3)]    0         \n",
      " er)                                                             \n",
      "                                                                 \n",
      " densenet169 (Functional)    (None, 9, 9, 1664)        12642880  \n",
      "                                                                 \n",
      " flatten (Flatten)           (None, 134784)            0         \n",
      "                                                                 \n",
      " dense (Dense)               (None, 16)                2156560   \n",
      "                                                                 \n",
      " dense_1 (Dense)             (None, 8)                 136       \n",
      "                                                                 \n",
      " dense_2 (Dense)             (None, 2)                 18        \n",
      "                                                                 \n",
      "=================================================================\n",
      "Total params: 14,799,594\n",
      "Trainable params: 14,641,194\n",
      "Non-trainable params: 158,400\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "# Prints a summary of the compiled model\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b69ae2d4",
   "metadata": {},
   "source": [
    "### Training the model\n",
    "Now that we have everything ready, we can train the model. we pass the `validation_data` containing the `X_test` and `y_test` to the model as well.\n",
    "Also, for above mentioned checkopints, we need to call those callbacks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "449d63cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "history = model.fit(X_train, y_train, epochs=20, batch_size=8, validation_data=(X_test, y_test), callbacks=[tensorboard_callback, model_checkpoint_callback])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea0873cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plotting the tensorboard data in jupyter notebook\n",
    "%tensorboard --logdir logs/fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b3d7255",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_error(y, y_pred):\n",
    "    \"\"\"\n",
    "    Since the competion computes the error based on Euclidian distances, this function computes\n",
    "    the error for predicted and ground truth results.\n",
    "    \n",
    "    @param: list y. a (1, 2) list with original coordinations.\n",
    "    @param: list y_pred. a (1, 2) list containing the predicted coordinations.\n",
    "    \n",
    "    @return: float error. The Euclidian distance between these two given datapoints.\n",
    "    \"\"\"\n",
    "    error = np.sqrt(np.abs(y[0] - y_pred[0]) ** 2 + np.abs(y[1] - y_pred[1]) ** 2)\n",
    "    return error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fba6964b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_original_coordination(predicted_coordination, original_box_coordination):\n",
    "    \"\"\"\n",
    "    Computes the coordinations of the predicted Hexbug's head's coordination in the original image\n",
    "    using the bounding box coordinations.\n",
    "    \n",
    "    @param: list predicted_coordination. a (1, 2) list with (x, y) format.\n",
    "    @param: list original_box_coordination. a (1, 2) list with (x, y) format of bounding box.\n",
    "    \n",
    "    @return: list. Coordinations of the predicted Hexbug's head's coordination in the original image\n",
    "    \"\"\"\n",
    "    return (predicted_coordination[0] + original_box_coordination[0], predicted_coordination[1] + original_box_coordination[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35a457d0",
   "metadata": {},
   "source": [
    "### Loading the best model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4afd61d",
   "metadata": {},
   "outputs": [],
   "source": [
    "reconstructed_model = tf.keras.models.load_model(\"checkpoint_callback\")\n",
    "# reconstructed_model = tf.keras.models.load_model(\"resnet50_trained_model_data/v1.5/checkpoint_callback\")\n",
    "clear_output()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89e47066",
   "metadata": {},
   "source": [
    "### Sample Evaluation\n",
    "Now that we have successfully trained the model, let's see its performance on unseen data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a268e22",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_samples = 10\n",
    "\n",
    "# Chooses random indices\n",
    "random_samples_indices = random.sample(range(0, len(X_test)), num_samples)\n",
    "\n",
    "# Iterates over each index and retrieves the image and annotaion\n",
    "for sample_index in random_samples_indices:\n",
    "    sample_img = X_test[sample_index]\n",
    "    \n",
    "    # Configs for drawing a circle on the image.\n",
    "    center = y_test[sample_index]\n",
    "    radius = 5\n",
    "    # Color map is BGR\n",
    "    color = (0, 255, 0)\n",
    "    color_pred = (0, 0, 255)\n",
    "    thickness = 5\n",
    "    \n",
    "    # Since we predict only one image at a time, we need to expand its dimention\n",
    "    # to fit the input layer of our model.\n",
    "    expanded_img = np.expand_dims(sample_img, axis=0)\n",
    "    \n",
    "    # Predicts the coordinations\n",
    "    predicted = list(map(int, reconstructed_model.predict(expanded_img)[0]))\n",
    "    \n",
    "    # Computes the error\n",
    "    error = round(compute_error(center, predicted), 3)\n",
    "    \n",
    "    print(f'y_predicted: {predicted}, y: {center}, Error: {error}')\n",
    "    \n",
    "    # Draws a circle centered in the correct coordination.\n",
    "    img_lv1 = cv2.circle(sample_img, center, radius, color, thickness)\n",
    "    # Draws a circle centered in the predicted coordination.\n",
    "    img_lv2 = cv2.circle(img_lv1, predicted, radius, color_pred, thickness)\n",
    "    \n",
    "    # Shows the output image.\n",
    "    cv2.imshow(f'Sample No. {sample_index}, error: {error}', img_lv2)\n",
    "    cv2.waitKey(0)\n",
    "    cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5219f9d1",
   "metadata": {},
   "source": [
    "### Sample for coordination convertion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50fe6cd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_samples = 10\n",
    "\n",
    "random_samples_indices = random.sample(range(0, len(data)), num_samples)\n",
    "\n",
    "for sample_id in random_samples_indices:\n",
    "    sample_img_path = data.iloc[sample_id].Path\n",
    "\n",
    "    # Constructs the path to the original frame\n",
    "    original_image_path = f\"Videos/{sample_img_path.split('/')[1]}/{sample_img_path.split('/')[2]}\"\n",
    "\n",
    "    sample_img = np.array(load_img(sample_img_path))\n",
    "\n",
    "    expanded_img = np.expand_dims(sample_img, axis=0)\n",
    "\n",
    "    predicted = list(map(int, reconstructed_model.predict(expanded_img)[0]))\n",
    "\n",
    "    radius = 10\n",
    "    color = (0, 255, 0)\n",
    "    thickness = -1\n",
    "\n",
    "    # Computes the coordination of the Hexbug's head in the original image.\n",
    "    scaled_coordinations = compute_original_coordination(predicted, (data.iloc[sample_id].OriginalBoxCoordinationX1, data.iloc[sample_id].OriginalBoxCoordinationY1))\n",
    "\n",
    "    # Reads and converts the original image into numpy array.\n",
    "    original_sample_img = np.array(load_img(original_image_path))\n",
    "\n",
    "    annotated_img = cv2.circle(sample_img, predicted, radius, color, thickness)\n",
    "\n",
    "    cv2.circle(original_sample_img, scaled_coordinations, radius, color, thickness)\n",
    "    \n",
    "    center = [data.iloc[sample_id].CroppedHexBugCoordinationX, data.iloc[sample_id].CroppedHexBugCoordinationY]\n",
    "    \n",
    "    error = round(compute_error(center, predicted), 3)\n",
    "    \n",
    "    file_name = f'Sample ID: {sample_id}\\nTruth: {center}, Predicted: {predicted}, Error: {error}'\n",
    "    \n",
    "    print(file_name)\n",
    "    print('------------------------------------------------------------------------')\n",
    "    \n",
    "    if not os.path.exists('annotated_images'):\n",
    "        os.mkdir('annotated_images')\n",
    "    \n",
    "    cv2.imshow(f'Sample ID. {sample_id}', original_sample_img)\n",
    "    cv2.imwrite(f'annotated_images/{file_name}.jpg', original_sample_img)\n",
    "    cv2.waitKey(0)\n",
    "    cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7b3d725",
   "metadata": {},
   "outputs": [],
   "source": [
    "# reconstructed_model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e462f764",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model = Sequential()\n",
    "# model.add(Conv2D(64, (3, 3), activation='relu', input_shape=X_train[0].shape))\n",
    "# # model.add(Dropout(0.25))\n",
    "# model.add(MaxPooling2D((2, 2)))\n",
    "# model.add(Conv2D(128, (3, 3), activation='relu', kernel_regularizer='L2'))\n",
    "# # model.add(Dropout(0.25))\n",
    "# model.add(MaxPooling2D((2, 2)))\n",
    "# model.add(Conv2D(512, (3, 3), activation='relu', kernel_regularizer='L2'))\n",
    "# # model.add(Dropout(0.25))\n",
    "# model.add(MaxPooling2D((2, 2)))\n",
    "# model.add(Conv2D(1028, (3, 3), activation='relu', kernel_regularizer='L2'))\n",
    "# # model.add(Dropout(0.25))\n",
    "# model.add(MaxPooling2D((2, 2)))\n",
    "# model.add(Conv2D(256, (3, 3), activation='relu', kernel_regularizer='L2'))\n",
    "# # model.add(Dropout(0.25))\n",
    "# model.add(MaxPooling2D((2, 2)))\n",
    "# model.add(Flatten())\n",
    "# model.add(Dense(128, activation='relu'))\n",
    "# model.add(Dense(2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea52e413",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69b6ed69",
   "metadata": {},
   "outputs": [],
   "source": [
    "# reconstructed_model.compile(loss='mse', optimizer='adam')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6f34908",
   "metadata": {},
   "outputs": [],
   "source": [
    "# history = reconstructed_model.fit(X_train, y_train, epochs=100, batch_size=16, validation_data=(X_test, y_test), callbacks=[tensorboard_callback, model_checkpoint_callback])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89ebe43e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}