hexbug_head_detection_train.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": "1babc3a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install --user --proxy http://proxy:80 protobuf==3.19.6\n",
    "# !pip install --user --proxy http://proxy:80 tensorboard==2.11\n",
    "# !pip install --user --proxy http://proxy:80 pip -U scikit-learn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8f3dcd12",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/bin/bash: /apps/python/3.10-anaconda/envs/tensorflow/lib/libtinfo.so.6: no version information available (required by /bin/bash)\n",
      "Mon May 29 11:05:53 2023       \n",
      "+---------------------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 530.41.03              Driver Version: 530.41.03    CUDA Version: 12.1     |\n",
      "|-----------------------------------------+----------------------+----------------------+\n",
      "| GPU  Name                  Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp  Perf            Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                                         |                      |               MIG M. |\n",
      "|=========================================+======================+======================|\n",
      "|   0  NVIDIA A100-SXM4-40GB           On | 00000000:01:00.0 Off |                    0 |\n",
      "| N/A   33C    P0               58W / 400W|      0MiB / 40960MiB |      0%      Default |\n",
      "|                                         |                      |             Disabled |\n",
      "+-----------------------------------------+----------------------+----------------------+\n",
      "|   1  NVIDIA A100-SXM4-40GB           On | 00000000:41:00.0 Off |                    0 |\n",
      "| N/A   33C    P0               59W / 400W|      0MiB / 40960MiB |      0%      Default |\n",
      "|                                         |                      |             Disabled |\n",
      "+-----------------------------------------+----------------------+----------------------+\n",
      "|   2  NVIDIA A100-SXM4-40GB           On | 00000000:81:00.0 Off |                    0 |\n",
      "| N/A   31C    P0               59W / 400W|      0MiB / 40960MiB |      0%      Default |\n",
      "|                                         |                      |             Disabled |\n",
      "+-----------------------------------------+----------------------+----------------------+\n",
      "|   3  NVIDIA A100-SXM4-40GB           On | 00000000:C1:00.0 Off |                    0 |\n",
      "| N/A   30C    P0               56W / 400W|      0MiB / 40960MiB |      0%      Default |\n",
      "|                                         |                      |             Disabled |\n",
      "+-----------------------------------------+----------------------+----------------------+\n",
      "                                                                                         \n",
      "+---------------------------------------------------------------------------------------+\n",
      "| Processes:                                                                            |\n",
      "|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |\n",
      "|        ID   ID                                                             Usage      |\n",
      "|=======================================================================================|\n",
      "|  No running processes found                                                           |\n",
      "+---------------------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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",
    "from tqdm import tqdm\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": 4,
   "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>CroppedImagePath</th>\n",
       "      <th>OriginalImagePath</th>\n",
       "      <th>Xoffset</th>\n",
       "      <th>Yoffset</th>\n",
       "      <th>ID</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>4174</th>\n",
       "      <td>88.267358</td>\n",
       "      <td>26.049741</td>\n",
       "      <td>142</td>\n",
       "      <td>502</td>\n",
       "      <td>cropped_bugs/training01/frame0_0.jpg</td>\n",
       "      <td>samples/training01/frame0.jpg</td>\n",
       "      <td>143</td>\n",
       "      <td>114</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>838</th>\n",
       "      <td>53.343207</td>\n",
       "      <td>176.848746</td>\n",
       "      <td>5</td>\n",
       "      <td>411</td>\n",
       "      <td>cropped_bugs/training0100/frame0_0.jpg</td>\n",
       "      <td>samples/training0100/frame0.jpg</td>\n",
       "      <td>151</td>\n",
       "      <td>100</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4275</th>\n",
       "      <td>79.964068</td>\n",
       "      <td>23.098983</td>\n",
       "      <td>587</td>\n",
       "      <td>447</td>\n",
       "      <td>cropped_bugs/training012/frame0_0.jpg</td>\n",
       "      <td>samples/training012/frame0.jpg</td>\n",
       "      <td>144</td>\n",
       "      <td>94</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2213</th>\n",
       "      <td>36.618670</td>\n",
       "      <td>28.188413</td>\n",
       "      <td>230</td>\n",
       "      <td>26</td>\n",
       "      <td>cropped_bugs/training018/frame0_0.jpg</td>\n",
       "      <td>samples/training018/frame0.jpg</td>\n",
       "      <td>104</td>\n",
       "      <td>96</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1437</th>\n",
       "      <td>183.235621</td>\n",
       "      <td>55.906871</td>\n",
       "      <td>558</td>\n",
       "      <td>935</td>\n",
       "      <td>cropped_bugs/training020/frame0_0.jpg</td>\n",
       "      <td>samples/training020/frame0.jpg</td>\n",
       "      <td>97</td>\n",
       "      <td>150</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",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3503</th>\n",
       "      <td>82.155047</td>\n",
       "      <td>163.824014</td>\n",
       "      <td>426</td>\n",
       "      <td>1020</td>\n",
       "      <td>cropped_bugs/training087/frame100_0.jpg</td>\n",
       "      <td>samples/training087/frame100.jpg</td>\n",
       "      <td>68</td>\n",
       "      <td>98</td>\n",
       "      <td>100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2412</th>\n",
       "      <td>185.219286</td>\n",
       "      <td>35.021525</td>\n",
       "      <td>768</td>\n",
       "      <td>370</td>\n",
       "      <td>cropped_bugs/training090/frame100_0.jpg</td>\n",
       "      <td>samples/training090/frame100.jpg</td>\n",
       "      <td>87</td>\n",
       "      <td>154</td>\n",
       "      <td>100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3016</th>\n",
       "      <td>75.259948</td>\n",
       "      <td>33.841683</td>\n",
       "      <td>138</td>\n",
       "      <td>610</td>\n",
       "      <td>cropped_bugs/training091/frame100_0.jpg</td>\n",
       "      <td>samples/training091/frame100.jpg</td>\n",
       "      <td>148</td>\n",
       "      <td>110</td>\n",
       "      <td>100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2021</th>\n",
       "      <td>98.251275</td>\n",
       "      <td>28.137171</td>\n",
       "      <td>833</td>\n",
       "      <td>277</td>\n",
       "      <td>cropped_bugs/training094/frame100_0.jpg</td>\n",
       "      <td>samples/training094/frame100.jpg</td>\n",
       "      <td>139</td>\n",
       "      <td>106</td>\n",
       "      <td>100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2312</th>\n",
       "      <td>168.684583</td>\n",
       "      <td>172.479272</td>\n",
       "      <td>832</td>\n",
       "      <td>1292</td>\n",
       "      <td>cropped_bugs/training099/frame100_0.jpg</td>\n",
       "      <td>samples/training099/frame100.jpg</td>\n",
       "      <td>96</td>\n",
       "      <td>102</td>\n",
       "      <td>100</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5112 rows × 9 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      CroppedHexBugCoordinationX  CroppedHexBugCoordinationY   \n",
       "4174                   88.267358                   26.049741  \\\n",
       "838                    53.343207                  176.848746   \n",
       "4275                   79.964068                   23.098983   \n",
       "2213                   36.618670                   28.188413   \n",
       "1437                  183.235621                   55.906871   \n",
       "...                          ...                         ...   \n",
       "3503                   82.155047                  163.824014   \n",
       "2412                  185.219286                   35.021525   \n",
       "3016                   75.259948                   33.841683   \n",
       "2021                   98.251275                   28.137171   \n",
       "2312                  168.684583                  172.479272   \n",
       "\n",
       "      OriginalBoxCoordinationX1  OriginalBoxCoordinationY1   \n",
       "4174                        142                        502  \\\n",
       "838                           5                        411   \n",
       "4275                        587                        447   \n",
       "2213                        230                         26   \n",
       "1437                        558                        935   \n",
       "...                         ...                        ...   \n",
       "3503                        426                       1020   \n",
       "2412                        768                        370   \n",
       "3016                        138                        610   \n",
       "2021                        833                        277   \n",
       "2312                        832                       1292   \n",
       "\n",
       "                             CroppedImagePath   \n",
       "4174     cropped_bugs/training01/frame0_0.jpg  \\\n",
       "838    cropped_bugs/training0100/frame0_0.jpg   \n",
       "4275    cropped_bugs/training012/frame0_0.jpg   \n",
       "2213    cropped_bugs/training018/frame0_0.jpg   \n",
       "1437    cropped_bugs/training020/frame0_0.jpg   \n",
       "...                                       ...   \n",
       "3503  cropped_bugs/training087/frame100_0.jpg   \n",
       "2412  cropped_bugs/training090/frame100_0.jpg   \n",
       "3016  cropped_bugs/training091/frame100_0.jpg   \n",
       "2021  cropped_bugs/training094/frame100_0.jpg   \n",
       "2312  cropped_bugs/training099/frame100_0.jpg   \n",
       "\n",
       "                     OriginalImagePath  Xoffset  Yoffset   ID  \n",
       "4174     samples/training01/frame0.jpg      143      114    0  \n",
       "838    samples/training0100/frame0.jpg      151      100    0  \n",
       "4275    samples/training012/frame0.jpg      144       94    0  \n",
       "2213    samples/training018/frame0.jpg      104       96    0  \n",
       "1437    samples/training020/frame0.jpg       97      150    0  \n",
       "...                                ...      ...      ...  ...  \n",
       "3503  samples/training087/frame100.jpg       68       98  100  \n",
       "2412  samples/training090/frame100.jpg       87      154  100  \n",
       "3016  samples/training091/frame100.jpg      148      110  100  \n",
       "2021  samples/training094/frame100.jpg      139      106  100  \n",
       "2312  samples/training099/frame100.jpg       96      102  100  \n",
       "\n",
       "[5112 rows x 9 columns]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = pd.read_csv('data_cropped_bugs.csv').drop(['Unnamed: 0'], axis=1).sort_values(['ID', 'OriginalImagePath'], ascending=True)\n",
    "# data = data[(data['CroppedHexBugCoordinationX'] > 0) &\n",
    "#             (data['CroppedHexBugCoordinationY'] > 0) &\n",
    "#             (data['OriginalBoxCoordinationX1'] > 0) &\n",
    "#             (data['OriginalBoxCoordinationY1'] > 0)]\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "481aa283",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "5112it [00:23, 213.06it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(5112, 400, 400, 3)\n",
      "(5112, 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 tqdm(data.iterrows()):\n",
    "#     img_path = f\"{row.Path.replace('cropped_bugs', 'image_sharpning/sharpened_images')}\"\n",
    "    img_path = f\"{row.CroppedImagePath}\"\n",
    "    annotations = [row.CroppedHexBugCoordinationX + row.Xoffset, row.CroppedHexBugCoordinationY + row.Yoffset]\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": 6,
   "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": 7,
   "id": "495ae437",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X train shape: (3425, 400, 400, 3)\n",
      "X test shape: (1687, 400, 400, 3)\n",
      "y train shape: (3425, 2)\n",
      "y test shape: (1687, 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": 8,
   "id": "d18e66d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dense_model = DenseNet121(\n",
    "#     include_top=False,\n",
    "#     weights='densenet121_weights_tf_dim_ordering_tf_kernels_notop.h5',\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": 9,
   "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": 10,
   "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": 11,
   "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='resnet50_weights_tf_dim_ordering_tf_kernels_notop.h5',\n",
    "#                     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": 12,
   "id": "f8ad8d65",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/bin/bash: /apps/python/3.10-anaconda/envs/tensorflow/lib/libtinfo.so.6: no version information available (required by /bin/bash)\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": 13,
   "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": 14,
   "id": "7549e07c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-05-29 11:06:29.328233: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-05-29 11:06:32.794260: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 38241 MB memory:  -> device: 0, name: NVIDIA A100-SXM4-40GB, pci bus id: 0000:01:00.0, compute capability: 8.0\n",
      "2023-05-29 11:06:32.796424: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 38241 MB memory:  -> device: 1, name: NVIDIA A100-SXM4-40GB, pci bus id: 0000:41:00.0, compute capability: 8.0\n",
      "2023-05-29 11:06:32.797886: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:2 with 38241 MB memory:  -> device: 2, name: NVIDIA A100-SXM4-40GB, pci bus id: 0000:81:00.0, compute capability: 8.0\n",
      "2023-05-29 11:06:32.799270: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:3 with 38241 MB memory:  -> device: 3, name: NVIDIA A100-SXM4-40GB, pci bus id: 0000:c1:00.0, compute capability: 8.0\n"
     ]
    }
   ],
   "source": [
    "# Define the strategy for distributing the training across GPUs\n",
    "strategy = tf.distribute.MirroredStrategy()\n",
    "\n",
    "# Create a distributed training context\n",
    "with strategy.scope():\n",
    "    \n",
    "    # Compile your model\n",
    "    dense_model = DenseNet201(\n",
    "        include_top=False,\n",
    "        weights='imagenet',\n",
    "        input_shape=X_train[0].shape\n",
    "    )\n",
    "\n",
    "    model = Sequential(dense_model, name='DenseNet201')\n",
    "\n",
    "    x = model.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=model.input, outputs=output_layer)\n",
    "    \n",
    "    \"\"\"\n",
    "    resnet50 = ResNet50(weights='imagenet',\n",
    "                        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",
    "    model.compile(loss='mse', optimizer='adam')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "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": 16,
   "id": "1369f20d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "_________________________________________________________________\n",
      " Layer (type)                Output Shape              Param #   \n",
      "=================================================================\n",
      " densenet201_input (InputLay  [(None, 400, 400, 3)]    0         \n",
      " er)                                                             \n",
      "                                                                 \n",
      " densenet201 (Functional)    (None, 12, 12, 1920)      18321984  \n",
      "                                                                 \n",
      " flatten (Flatten)           (None, 276480)            0         \n",
      "                                                                 \n",
      " dense (Dense)               (None, 32)                8847392   \n",
      "                                                                 \n",
      " dense_1 (Dense)             (None, 16)                528       \n",
      "                                                                 \n",
      " dense_2 (Dense)             (None, 2)                 34        \n",
      "                                                                 \n",
      "=================================================================\n",
      "Total params: 27,169,938\n",
      "Trainable params: 26,940,882\n",
      "Non-trainable params: 229,056\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": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/500\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-05-29 11:10:52.278733: I tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:428] Loaded cuDNN version 8100\n",
      "2023-05-29 11:10:53.886445: I tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:428] Loaded cuDNN version 8100\n",
      "2023-05-29 11:10:54.628557: I tensorflow/tsl/platform/default/subprocess.cc:304] Start cannot spawn child process: No such file or directory\n",
      "2023-05-29 11:10:55.030188: I tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:428] Loaded cuDNN version 8100\n",
      "2023-05-29 11:10:57.080546: I tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:428] Loaded cuDNN version 8100\n",
      "2023-05-29 11:11:02.613491: W tensorflow/tsl/framework/bfc_allocator.cc:290] Allocator (GPU_0_bfc) ran out of memory trying to allocate 2.66GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.\n",
      "2023-05-29 11:11:02.613546: W tensorflow/tsl/framework/bfc_allocator.cc:290] Allocator (GPU_0_bfc) ran out of memory trying to allocate 2.66GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.\n",
      "2023-05-29 11:11:02.632867: W tensorflow/tsl/framework/bfc_allocator.cc:290] Allocator (GPU_0_bfc) ran out of memory trying to allocate 2.71GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.\n",
      "2023-05-29 11:11:02.632890: W tensorflow/tsl/framework/bfc_allocator.cc:290] Allocator (GPU_0_bfc) ran out of memory trying to allocate 2.71GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.\n",
      "2023-05-29 11:11:02.651831: W tensorflow/tsl/framework/bfc_allocator.cc:290] Allocator (GPU_0_bfc) ran out of memory trying to allocate 2.75GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.\n",
      "2023-05-29 11:11:02.651855: W tensorflow/tsl/framework/bfc_allocator.cc:290] Allocator (GPU_0_bfc) ran out of memory trying to allocate 2.75GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.\n",
      "2023-05-29 11:11:02.706641: W tensorflow/tsl/framework/bfc_allocator.cc:290] Allocator (GPU_0_bfc) ran out of memory trying to allocate 782.06MiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.\n",
      "2023-05-29 11:11:02.706664: W tensorflow/tsl/framework/bfc_allocator.cc:290] Allocator (GPU_0_bfc) ran out of memory trying to allocate 782.06MiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.\n",
      "2023-05-29 11:11:02.765974: W tensorflow/tsl/framework/bfc_allocator.cc:290] Allocator (GPU_0_bfc) ran out of memory trying to allocate 448.50MiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.\n",
      "2023-05-29 11:11:02.765999: W tensorflow/tsl/framework/bfc_allocator.cc:290] Allocator (GPU_0_bfc) ran out of memory trying to allocate 448.50MiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.\n"
     ]
    }
   ],
   "source": [
    "history = model.fit(X_train, y_train, epochs=500, batch_size=256, 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 = 30\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(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 = 30\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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}