notebooks/4_finetune.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Finetune Visual Object Detection Models Using Pre-trained Sagemaker Models\n",
    "\n",
    "\n",
    "**!! Please run `0_demo.ipynb` and `1_retrain_from_checkpoint.ipynb` before running this notebook !!**\n",
    "* Notebook `0_demo.ipynb` setups input, output directories, data, and other artifacts in both local and s3. \n",
    "* Notebook `1_retrain_from_checkpoint.ipynb` prepares a DDN model checkpoint and you could use to compare with models trained in this notebook in [Section 6](#6.-Inference-and-Model-Comparison). \n",
    "\n",
    "---\n",
    "\n",
    "This notebook introduces finetuning pretrained object detection (OD) models on new dataset. \n",
    "\n",
    "Training a model from scratch in general is time-consuming and requires large compute resources. When the training data is small, we cannot expect to train a very performant model. A better alternative is to finetune a pretrained model on the target dataset. AWS Sagemaker provides high-quality pretrained models that were trained on very large datasets. Finetuning these models on new dataset takes only fractional training time compard to training from scratch.\n",
    "\n",
    "In this notebook, we demonstrate how to use two types of Amazon Sagemaker built-in OD models to finetune on the *[Steel Surface Defect](https://github.com/siddhartamukherjee/NEU-DET-Steel-Surface-Defect-Detection)*  dataset, which is used in this solution. \n",
    "* Type 1 (legacy): uses a built-in legacy [Object Detection algorithm](https://docs.aws.amazon.com/sagemaker/latest/dg/object-detection.html) and uses the *Single Shot multibox Detector* (SSD) model with either VGG or ResNet backbone, and was pretrained on the ImageNet dataset. \n",
    "* Type 2 (latest):  provides [9 pretrained OD models](https://sagemaker.readthedocs.io/en/stable/doc_utils/pretrainedmodels.html?highlight=jumpstart#built-in-algorithms-with-pre-trained-model-table), including 8 SSD models and 1 FasterRCNN model. These models use VGG, ResNet, or MobileNet as backbone, and were pretrained on COCO, VOC, or FPN datasets. \n",
    "\n",
    "\n",
    "For each type of model, besides training with default hyperparameters, we also perform hyperparameter tuning (i.e., HPO) using [Sagemaker Automatic Model Tuning (AMT)](https://docs.aws.amazon.com/sagemaker/latest/dg/automatic-model-tuning.html) to train even better model. \n",
    "\n",
    "Running the whole notebook takes about 8 hours. The most time-consuming part is running HPO jobs for both types of models. You could choose to run more HPO jobs in parallel in order to reduce running time if there are more EC2 instances available.\n",
    "\n",
    "---\n",
    "\n",
    "Content\n",
    "1. [Data Preparation](#1.-Data-Preparation)\n",
    "2. [Training: Finetune Type 1 (Legacy) OD Model](#2.-Training:-Finetune-Type-1-(Legacy)-OD-Model)\n",
    "3. [Training: Finetune Type 1 (Legacy) OD Model with HPO](#3.-Training:-Finetune-Type-1-(Legacy)-OD-model-with-HPO)\n",
    "4. [Training: Finetune Type 2 (Latest) OD Model](#4.-Training:-Finetune-Type-2-(Latest)-OD-Model)\n",
    "5. [Training: Finetune Type 2 (Latest) OD Model with HPO](#5.-Training:-Finetune-Type-2-(Latest)-OD-model-with-HPO)\n",
    "6. [Inference and Model Comparison](#6.-Inference-and-Model-Comparison)\n",
    "7. [Clean Up the Endpoints](#7.-Clean-Up-the-Endpoints)\n",
    "8. [Conclusion](#8.-Conclusion)\n",
    "\n",
    "---\n",
    "\n",
    "** **ATTENTION** ** \n",
    "\n",
    "* Running the notebook end-to-end takes 8~9 hours. We changed some parameter values so that the notebook took much shorter time to finish, at the cost of model trainig non-convergence.\n",
    "* Please change them back when you want to train till convergence. These parameters include `num_epochs=100` for training all models, and `max_jobs=20`, `max_parallel_jobs=10` for hyperparameter tuning.\n",
    "* The shown results in this notebook is for fully-convergent models.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epochs = 10          # change to 100\n",
    "max_jobs = 3             # change to 20 or more\n",
    "max_parallel_jobs = 3    # change to 5 or 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Update to latest Sagemaker package to use Sagemaker APIs for model training and deployment\n",
    "%pip install --upgrade sagemaker \n",
    "%pip install opencv-python-headless opencv-python==4.5.5.64 mxnet xmltodict pycocotools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import json\n",
    "import glob\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import sagemaker\n",
    "from sagemaker import get_execution_role, image_uris, model_uris, script_uris, hyperparameters\n",
    "from sagemaker.s3 import S3Downloader\n",
    "from sagemaker.utils import name_from_base\n",
    "from sagemaker.session import Session\n",
    "from sagemaker.estimator import Estimator\n",
    "from sagemaker.analytics import TrainingJobAnalytics\n",
    "from sagemaker import exceptions\n",
    "from sagemaker.tuner import HyperparameterTuner, ContinuousParameter, IntegerParameter, CategoricalParameter\n",
    "from botocore.exceptions import ClientError\n",
    "\n",
    "# import helper function to convert all xmls to a json file\n",
    "sys.path.append('../src')\n",
    "from xml2json import XML2JSON, convert_to_pycocotools_ground_truth\n",
    "\n",
    "sagemaker_session = sagemaker.Session()\n",
    "sagemaker_config = json.load(open(\"../stack_outputs.json\"))\n",
    "role = Session().get_caller_identity_arn()\n",
    "solution_bucket = sagemaker_config[\"SolutionS3Bucket\"]\n",
    "region = sagemaker_config[\"AWSRegion\"]\n",
    "solution_name = sagemaker_config[\"SolutionName\"]\n",
    "bucket = sagemaker_config[\"S3Bucket\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "prefix = \"neu-det\"\n",
    "neu_det_s3 = f\"s3://{bucket}/{prefix}\"\n",
    "print(neu_det_s3)\n",
    "\n",
    "# Folders for training data and output artifacts\n",
    "s3_input_train = f\"{neu_det_s3}/data_a\"\n",
    "s3_output_location = f\"{neu_det_s3}/output_a\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 1. Data Preparation\n",
    "\n",
    "The two types of OD models require different data formats.\n",
    "The *steel surface dataset* used in this solution contains one xml file for each image as annotation. However, \n",
    "neither model uses xml annotations. The Type 1 (legacy) OD model requires either RecordIO or image format in either [*file mode* or *pipe mode*](https://docs.aws.amazon.com/sagemaker/latest/dg/model-access-training-data.html). The Type 2 (latest) OD model requires the input must be a directory with a sub-directory of images and a `annotations.json` file. Please check Section 3 of this [notebook](https://github.com/aws/amazon-sagemaker-examples/blob/main/introduction_to_amazon_algorithms/jumpstart_object_detection/Amazon_JumpStart_Object_Detection.ipynb) for more explanation.\n",
    "\n",
    "In this notebook, we split the data to be train:val:test = 64:16:20. We allocate 20% data as test data to numerically compare all trained models in the end of the notebook. The steel surface dataset has 1800 images in 6 categories, we randomly allocate 20% images from each category to the test data.\n",
    "\n",
    "We provide a script to convert the remaining 80% xmls to a single `annotations.json` for training the Type 2 (latest) OD model (under the hood, the source code automatically splits the data to be train:val=80:20, equivalent to 64% of all data as train and 16% as val). We provide another script to convert the `annotations.json` and corresponding images to RecordIO data for the Type 1 (legacy) OD model. \n",
    "\n",
    "If your dataset follows the required input format for Type 1 (legacy) or Type 2 (latest) OD model, you do *not* need these conversions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Allocate 20% data for testing different models in the end, convert \n",
    "# their xmls to test_annotations.json and the remaining 80% to annotations.json\n",
    "\n",
    "path = \"raw_neu_det/NEU-DET/ANNOTATIONS/\"\n",
    "trainXMLFiles = glob.glob(os.path.join(path, '*.xml'))\n",
    "trainXMLFiles.sort()\n",
    "XML2JSON(trainXMLFiles, test_ratio=0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Sync the annotations.json and its corresponding train/val images to s3\n",
    "os.makedirs('neu_det_a/images', exist_ok=True)\n",
    "os.makedirs(\"results\", exist_ok=True)\n",
    "\n",
    "src_path = \"raw_neu_det/NEU-DET/IMAGES\"\n",
    "dst_path = \"neu_det_a/images\"\n",
    "\n",
    "with open(\"annotations.json\") as f:\n",
    "    images_annotations = json.loads(f.read())\n",
    "\n",
    "for entry in images_annotations['images']:\n",
    "    image_path = os.path.join(src_path, entry['file_name'])\n",
    "    os.system(f\"cp {image_path} {dst_path}\")\n",
    "\n",
    "!mv annotations.json neu_det_a\n",
    "!aws s3 sync neu_det_a $s3_input_train --quiet # remove the --quiet flag to view the sync logs\n",
    "!aws s3 ls $s3_input_train/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 2. Training: Finetune Type 1 (Legacy) OD Model\n",
    "\n",
    "\n",
    "We start from finetuning the Type 1 (legacy) OD model, which is the SSD model with ResNet as backbone, and pretrained on ImageNet. \n",
    "\n",
    "**Input data**: follow the [instruction](https://docs.aws.amazon.com/sagemaker/latest/dg/object-detection.html), the legacy OD model supports both RecordIO and image types for training in `file` mode, or RecordIO in `pipe` mode. In this notebook, we use RecordIO in file mode.\n",
    "We provide a script for converting the `annotations.json` to RecordIO format. The [document](https://cv.gluon.ai/build/examples_datasets/detection_custom.html#lst-label-for-gluoncv-and-mxnet) and [example](https://sagemaker-examples.readthedocs.io/en/latest/introduction_to_amazon_algorithms/object_detection_birds/object_detection_birds.html#Generate-RecordIO-files) provide some context for understanding the script.\n",
    "\n",
    "This script first splits the data to train:val = 80:20 according to the `train-ratio`. This is equivalent to use 64% of all data for training and 16% for validation. Then converts each partition, including images and annotations, to a .rec file. We will use the validation data for selecting the best job in HPO training in the next section, and will use the test data for numerically comparing all finetuned models. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "!python ../src/prepare_RecordIO.py ./neu_det_a/data ./neu_det_a/images --train-ratio 0.8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Upload the RecordIO files to train and validation channels\n",
    "train_channel = f\"{prefix}/OD_Type1_train\"\n",
    "validation_channel = f\"{prefix}/OD_Type1_validation\"\n",
    "\n",
    "s3_train_data = f\"s3://{bucket}/{train_channel}\"\n",
    "s3_validation_data = f\"s3://{bucket}/{validation_channel}\"\n",
    "s3_output_location = f\"s3://{bucket}/{prefix}/OD_output\"\n",
    "\n",
    "!aws s3 cp neu_det_a/data_train.rec $s3_train_data/\n",
    "!aws s3 cp neu_det_a/data_val.rec $s3_validation_data/\n",
    "!aws s3 ls $s3_train_data/\n",
    "\n",
    "train_data = sagemaker.inputs.TrainingInput(\n",
    "    s3_train_data,\n",
    "    distribution=\"FullyReplicated\",\n",
    "    content_type=\"application/x-recordio\",\n",
    "    s3_data_type=\"S3Prefix\",\n",
    ")\n",
    "validation_data = sagemaker.inputs.TrainingInput(\n",
    "    s3_validation_data,\n",
    "    distribution=\"FullyReplicated\",\n",
    "    content_type=\"application/x-recordio\",\n",
    "    s3_data_type=\"S3Prefix\",\n",
    ")\n",
    "\n",
    "data_channels = {\"train\": train_data, \"validation\": validation_data}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Train Type 1 (legacy) OD model\n",
    "job_name_prefix = sagemaker_config['SolutionPrefix'] + \"-Type1\"\n",
    "\n",
    "num_classes = 6\n",
    "num_training_samples = 1152  # total 1800 images, use 64% for training\n",
    "\n",
    "train_image_uri = image_uris.retrieve(region=sagemaker_session.boto_region_name, framework=\"object-detection\", version=\"latest\")\n",
    "\n",
    "print(\"Train Type 1 (legacy) OD model -------------------\")\n",
    "\n",
    "# In case one type of EC2 instance is not available and fails the training, try another EC2 instance type.\n",
    "# This built-in OD algorithm only accepts GPU instances for training.\n",
    "for instance in [\"ml.p3.2xlarge\", \"ml.g4dn.xlarge\", \"ml.g5.2xlarge\"]:\n",
    "    try:\n",
    "        od_model = Estimator(\n",
    "            train_image_uri,\n",
    "            role,\n",
    "            instance_count=1,\n",
    "            instance_type=instance,\n",
    "            volume_size=50,\n",
    "            max_run=360000,\n",
    "            input_mode=\"File\",\n",
    "            output_path=s3_output_location,\n",
    "            sagemaker_session=sagemaker_session,\n",
    "            base_job_name=job_name_prefix,\n",
    "        )\n",
    "\n",
    "        od_model.set_hyperparameters(\n",
    "            base_network=\"resnet-50\",\n",
    "            use_pretrained_model=1,\n",
    "            num_classes=num_classes,\n",
    "            mini_batch_size=16,\n",
    "            epochs=num_epochs,\n",
    "            learning_rate=0.001,\n",
    "            momentum=0.9,\n",
    "            weight_decay=0.0005,\n",
    "            lr_scheduler_step=\"33,67\",\n",
    "            lr_scheduler_factor=0.1,\n",
    "            optimizer=\"sgd\",\n",
    "            overlap_threshold=0.5,\n",
    "            nms_threshold=0.45,\n",
    "            num_training_samples=num_training_samples,\n",
    "        )\n",
    "\n",
    "        model = od_model.fit(inputs=data_channels, logs=\"All\")\n",
    "\n",
    "    except exceptions.CapacityError as e:\n",
    "        print(\"Training Exception:\", e)\n",
    "        print(f\"{instance} is not available !!\")\n",
    "        continue\n",
    "    except exceptions.UnexpectedStatusException as e:\n",
    "        print(\"Training Exception:\", e)\n",
    "        continue\n",
    "    except ClientError as e:\n",
    "        print(\"Training Exception:\", e)\n",
    "        continue\n",
    "    else:\n",
    "        print(f\"Instance {instance} is available !\")\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Visualize Training Progress\n",
    "\n",
    "During training, the loss function is the sum of CrossEntropy loss and SmoothL1 loss. We will visualize the two losses on the training data as well as the mean Average Precision (mAP) on the validation data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "job_name = od_model.latest_training_job.job_name\n",
    "df = TrainingJobAnalytics(job_name).dataframe()\n",
    "\n",
    "OD_Type1_metrics = list(set(df.metric_name.values))\n",
    "print('All metrics:', OD_Type1_metrics)\n",
    "num_metrics = len(OD_Type1_metrics)\n",
    "\n",
    "# The train:progress shows the N training epochs, use it to index x axis\n",
    "epochs = df[df['metric_name'] == 'train:progress']['value'].values\n",
    "df = df[df['metric_name'] != 'train:progress']\n",
    "\n",
    "\n",
    "plt.figure(figsize=(18, 5))\n",
    "cnt = 1\n",
    "for m in OD_Type1_metrics:\n",
    "    if m != 'train:progress':\n",
    "        d = df[df['metric_name'] == m]\n",
    "        \n",
    "        if m == 'validation:mAP':\n",
    "            v = list(d.value)[-1]\n",
    "            print(f'Final validation:mAP = {v:.4f}')\n",
    "        plt.subplot(1, num_metrics-1, cnt)\n",
    "        \n",
    "        # in case length mismatch\n",
    "        l1, l2 = len(epochs), len(d['value'])\n",
    "        l = min(l1, l2)\n",
    "        plt.plot(epochs[-l:], d['value'][-l:])\n",
    "        \n",
    "        plt.title(m)\n",
    "        plt.xlabel('epoch')\n",
    "        plt.ylabel(m)\n",
    "        cnt += 1\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Deployment\n",
    "\n",
    "The inference will be deferred to the end of the notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "od_type1_endpoint_name = name_from_base(sagemaker_config['SolutionPrefix'] + '-Type1')\n",
    "print(od_type1_endpoint_name)\n",
    "od_type1_predictor = od_model.deploy(endpoint_name=od_type1_endpoint_name, \n",
    "                                     initial_instance_count=1, instance_type=\"ml.m5.xlarge\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 3. Training: Finetune Type 1 (Legacy) OD model with HPO\n",
    "\n",
    "Now we run HPO to find better hyperparameters which lead to better model. You could find all [finetunable hyperparameters](https://docs.aws.amazon.com/sagemaker/latest/dg/object-detection-tuning.html) for the Type 1 (legacy) OD model. In this notebook, we only finetune learning rate, momentum, and weight decay. \n",
    "\n",
    "We will use [Sagemaker Automatic Model Tuning](https://docs.aws.amazon.com/sagemaker/latest/dg/automatic-model-tuning.html) (AMT) to run HPO. We need to provide hyperparameter ranges and objective metrics. AMT will monitor the log and parse the objective metrics. For object detection, we use mean Average Precision (mAP) on the validation dataset as our metric. mAP is the standard evaluation metric used in the [COCO Challenge](https://cocodataset.org/#detection-eval) for object detection tasks. Here is a nice [blog post](https://jonathan-hui.medium.com/map-mean-average-precision-for-object-detection-45c121a31173) explaining mAP for object detection. \n",
    "\n",
    "We run `max_jobs=20` jobs in this HPO. You could run more jobs to find even better hyperparameters, at the cost of more compute resources and training time. This HPO job takes about 1 hour using p3.2xlarge EC2 instance and run `max_parallel_jobs=10` jobs in parallel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "tuning_job_name = sagemaker_config['SolutionPrefix'] + '-Type1-HPO'\n",
    "\n",
    "hyperparameter_ranges = {\n",
    "    \"learning_rate\": ContinuousParameter(1e-4, 1e-2),\n",
    "    \"momentum\": ContinuousParameter(0.8, 0.99),\n",
    "    \"weight_decay\": ContinuousParameter(1e-4, 1e-3),\n",
    "}\n",
    "\n",
    "objective_metric_name = \"validation:mAP\"\n",
    "\n",
    "print(\"Train Type 1 (legacy) OD model with HPO -------------------\")\n",
    "\n",
    "for instance in [\"ml.p3.2xlarge\", \"ml.g4dn.xlarge\", \"ml.g5.2xlarge\"]:\n",
    "    try:\n",
    "        od_model = Estimator(\n",
    "            train_image_uri,\n",
    "            role,\n",
    "            instance_count=1,\n",
    "            instance_type=instance,\n",
    "            volume_size=50,\n",
    "            max_run=360000,\n",
    "            input_mode=\"File\",\n",
    "            output_path=s3_output_location,\n",
    "            sagemaker_session=sagemaker_session,\n",
    "            base_job_name=job_name_prefix,\n",
    "        )\n",
    "\n",
    "        od_model.set_hyperparameters(\n",
    "            base_network=\"resnet-50\",\n",
    "            use_pretrained_model=1,\n",
    "            num_classes=num_classes,\n",
    "            mini_batch_size=16,\n",
    "            epochs=num_epochs,\n",
    "#             learning_rate=0.001,\n",
    "#             momentum=0.9,\n",
    "#             weight_decay=0.0005,\n",
    "            lr_scheduler_step=\"33,67\",\n",
    "            lr_scheduler_factor=0.1,\n",
    "            optimizer=\"sgd\",\n",
    "            overlap_threshold=0.5,\n",
    "            nms_threshold=0.45,\n",
    "            num_training_samples=num_training_samples,\n",
    "        )\n",
    "\n",
    "        tuner = HyperparameterTuner(\n",
    "            od_model,\n",
    "            objective_metric_name,\n",
    "            hyperparameter_ranges,\n",
    "            objective_type=\"Maximize\",\n",
    "            max_jobs=max_jobs,\n",
    "            max_parallel_jobs=max_parallel_jobs,  # <- increase this number if more EC2 instances are available\n",
    "            base_tuning_job_name=tuning_job_name,\n",
    "        )\n",
    "\n",
    "\n",
    "        tuner.fit(inputs=data_channels, logs=\"All\")\n",
    "\n",
    "    except exceptions.CapacityError as e:\n",
    "        print(\"Training exception:\", e)\n",
    "        print(f\"{instance} is not available !!\")\n",
    "        continue\n",
    "    except exceptions.UnexpectedStatusException as e:\n",
    "        print(\"Training exception:\", e)\n",
    "        continue\n",
    "    except ClientError as e:\n",
    "        print(\"Training Exception:\", e)\n",
    "        continue\n",
    "    else:\n",
    "        print(f\"Instance {instance} is available !\")\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "tuning_job_name = tuner.latest_tuning_job.name\n",
    "print('tuning_job_name', tuning_job_name)\n",
    "\n",
    "tuner_analytics = sagemaker.HyperparameterTuningJobAnalytics(tuning_job_name)\n",
    "\n",
    "full_df = tuner_analytics.dataframe()\n",
    "\n",
    "if len(full_df) > 0:\n",
    "    df = full_df[full_df[\"FinalObjectiveValue\"] > -float(\"inf\")]\n",
    "    if len(df) > 0:\n",
    "        df = df.sort_values(\"FinalObjectiveValue\", ascending=False)\n",
    "        # filter out failed jobs.\n",
    "        df = df[df[\"TrainingJobStatus\"] == \"Completed\"] \n",
    "        print(f\"Number of training jobs completed and with valid objective: {len(df)} / {len(full_df)}\")\n",
    "        print({\"lowest\": min(df[\"FinalObjectiveValue\"]), \"highest\": max(df[\"FinalObjectiveValue\"])})\n",
    "        pd.set_option(\"display.max_colwidth\", None)  # Don't truncate TrainingJobName\n",
    "    else:\n",
    "        print(\"No training jobs have reported valid results yet.\")\n",
    "        \n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Visualize the training progress of the best job\n",
    "best_job_name = df.iloc[0]['TrainingJobName']\n",
    "print('best job:', best_job_name)\n",
    "v = df.iloc[0]['FinalObjectiveValue']\n",
    "print(f\"best job final validation:mAP = {v:.6f}\")\n",
    "\n",
    "df_best = TrainingJobAnalytics(best_job_name).dataframe()\n",
    "\n",
    "OD_Type1_metrics = list(set(df_best.metric_name.values))\n",
    "print('All metrics:', OD_Type1_metrics)\n",
    "print('ObjectiveMetric is exactly the same as validation:mAP')\n",
    "num_metrics = len(OD_Type1_metrics)\n",
    "\n",
    "# The train:progress shows the N training epochs, use it to index x axis\n",
    "epochs = df_best[df_best['metric_name'] == 'train:progress']['value'].values\n",
    "df_best = df_best[df_best['metric_name'] != 'train:progress']\n",
    "\n",
    "plt.figure(figsize=(18, 5))\n",
    "cnt = 1\n",
    "for m in OD_Type1_metrics:\n",
    "    if m != 'train:progress' and m != 'ObjectiveMetric':\n",
    "\n",
    "        d = df_best[df_best['metric_name'] == m]\n",
    "        plt.subplot(1, num_metrics-2, cnt)\n",
    "        \n",
    "        # in case length mismatch\n",
    "        l1, l2 = len(epochs), len(d['value'])\n",
    "        l = min(l1, l2)\n",
    "        plt.plot(epochs[-l:], d['value'][-l:])\n",
    "        \n",
    "        plt.title(m)\n",
    "        plt.xlabel('epoch')\n",
    "        plt.ylabel(m)\n",
    "        cnt += 1\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Deploy the best model from HPO\n",
    "\n",
    "The inference will be deferred to the end of the notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "od_type1_hpo_endpoint_name = name_from_base(sagemaker_config['SolutionPrefix'] + \"-Type1-HPO\")\n",
    "print(od_type1_hpo_endpoint_name)\n",
    "od_type1_hpo_predictor = tuner.deploy(endpoint_name=od_type1_hpo_endpoint_name, initial_instance_count=1, instance_type=\"ml.m5.xlarge\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 4. Training: Finetune Type 2 (Latest) OD Model\n",
    "\n",
    "For the Type 2 (latest) OD model, we follow [Fine-tune a Model and Deploy to a SageMaker Endpoint\n",
    "](https://sagemaker.readthedocs.io/en/stable/overview.html#fine-tune-a-model-and-deploy-to-a-sagemaker-endpoint) and use standard Sagemaker APIs. \n",
    "\n",
    "You can find all finetunable Type 2 (latest) OD models in [Built-in Algorithms with pre-trained Model Table](https://sagemaker.readthedocs.io/en/stable/doc_utils/pretrainedmodels.html) by searching with keywords \"object detection\" and set `FineTunable?=True`.\n",
    "Currently there are 9 finetunable OD models:\n",
    "1. mxnet-od-ssd-300-vgg16-atrous-coco\n",
    "2. mxnet-od-ssd-512-vgg16-atrous-voc\n",
    "3. mxnet-od-ssd-512-resnet50-v1-coco\n",
    "4. mxnet-od-ssd-512-mobilenet1-0-coco\n",
    "5. mxnet-od-ssd-300-vgg16-atrous-voc\n",
    "6. mxnet-od-ssd-512-resnet50-v1-voc\n",
    "7. mxnet-od-ssd-512-mobilenet1-0-voc\n",
    "8. mxnet-od-ssd-512-vgg16-atrous-coco\n",
    "9. pytorch-od1-fasterrcnn-resnet50-fpn\n",
    "\n",
    "\n",
    "There are two major differences between training the two types of OD models: \n",
    "1. The entry point `transfer_learning.py` for finetuning a Type 2 (latest) OD model does not accept a validation data channel. Instead, it splits the input data provided through  `estimator.fit({\"training\": s3_input_train})` to be train:val=80:20, corresponding to use 64% of total data for training and 16% for validation. Note, the train/val data will be different from train/val for training Type 1 (legacy) OD model. \n",
    "2. The evaluation metrics are different. While Type 1 (legacy) OD model reports mAP on the validation data, which is standard, the Type 2 (latest) OD model only reports CrossEntropy loss and SmoothL1 loss on the validation data.   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "job_name_prefix = sagemaker_config['SolutionPrefix'] + '-JS'\n",
    "\n",
    "# Choose a pre-trained model and fine-tune on new dataset\n",
    "model_id, model_version = \"mxnet-od-ssd-512-vgg16-atrous-voc\", \"*\"\n",
    "\n",
    "scope = 'training'\n",
    "\n",
    "# Define metrics for visualization\n",
    "# You can visualize in Training / Traning jobs from Sagemaker Dashboard\n",
    "if model_id.startswith('mxnet'):\n",
    "    metric_definitions = [\n",
    "        {'Name': \"CrossEntropy\", 'Regex': \"CrossEntropy=(.*?),\"},\n",
    "        {'Name': \"SmoothL1\",     'Regex': \"SmoothL1=(.*)\"},\n",
    "        {'Name': \"Val_CrossEntropy\", 'Regex': \"Val_CrossEntropy=(.*?),\"},\n",
    "        {'Name': \"Val_SmoothL1\",     'Regex': \"Val_SmoothL1=(.*)\"},        \n",
    "    ]\n",
    "elif model_id.startswith('pytorch'):\n",
    "    metric_definitions = [\n",
    "        {'Name': \"train_loss\",             \"Regex\": \"train_loss:(.*?) \\(\"},\n",
    "        {\"Name\": \"train_loss_classifier\",  \"Regex\": \"train_loss_classifier:(.*?) \\(\"},\n",
    "        {\"Name\": \"train_loss_box_reg\",     \"Regex\": \"train_loss_box_reg:(.*?) \\(\"},\n",
    "        {\"Name\": \"train_loss_objectness\",  \"Regex\": \"train_loss_objectness:(.*?) \\(\"},\n",
    "        {\"Name\": \"train_loss_rpn_box_reg\", \"Regex\": \"train_loss_rpn_box_reg:(.*?) \\(\"},\n",
    "        {\"Name\": \"val_loss\",               \"Regex\": \"val_loss:(.*?) \\(\"}\n",
    "    ]\n",
    "else:\n",
    "    print(\"Incorrect model_id\")\n",
    "\n",
    "# Retrieve base model, training script, and training docker image\n",
    "train_model_uri = model_uris.retrieve(model_id=model_id, model_version=model_version, model_scope=scope)\n",
    "train_script_uri = script_uris.retrieve(model_id=model_id, model_version=model_version, script_scope=scope)\n",
    "\n",
    "# Change default hyperparameter values\n",
    "hyperparameters = hyperparameters.retrieve_default(model_id=model_id, model_version=model_version)\n",
    "hyperparameters['epochs'] = num_epochs     # 100 epochs takes 2 hours\n",
    "hyperparameters['batch-size'] = 8   # larger (e.g. 16) batch-size could lead to insufficient memory issue.\n",
    "print(hyperparameters)\n",
    "\n",
    "print(\"Train Type 2 (latest) OD model -------------------\")\n",
    "\n",
    "for training_instance_type in [\"ml.p3.2xlarge\", \"ml.g4dn.xlarge\", \"ml.g5.xlarge\", \"ml.m5.4xlarge\"]:\n",
    "    try:\n",
    "        train_image_uri = image_uris.retrieve(region=None, framework=None, image_scope=scope, model_id=model_id, model_version=model_version, instance_type=training_instance_type,)\n",
    "\n",
    "        estimator = Estimator(\n",
    "            role=role,\n",
    "            image_uri=train_image_uri,\n",
    "            source_dir=train_script_uri,\n",
    "            model_uri=train_model_uri,\n",
    "            entry_point=\"transfer_learning.py\",\n",
    "            hyperparameters=hyperparameters,\n",
    "            instance_count=1,\n",
    "            instance_type=training_instance_type,\n",
    "            max_run=36000,\n",
    "            output_path=s3_output_location,\n",
    "            base_job_name=job_name_prefix,\n",
    "            metric_definitions=metric_definitions,\n",
    "        )\n",
    "\n",
    "        estimator.fit({\"training\": s3_input_train}, logs=\"All\")\n",
    "        \n",
    "    except exceptions.CapacityError as e:\n",
    "        print(\"Training exception:\", e)\n",
    "        print(f\"{training_instance_type} is not available !!\")\n",
    "        continue\n",
    "    except exceptions.UnexpectedStatusException as e:\n",
    "        print(\"Training exception:\", e)\n",
    "        continue\n",
    "    except ClientError as e:\n",
    "        print(\"Training Exception:\", e)\n",
    "        continue\n",
    "    else:\n",
    "        print(f\"Instance {training_instance_type} is available !\")\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Visualization\n",
    "job_name = estimator.latest_training_job.job_name\n",
    "df = TrainingJobAnalytics(job_name).dataframe()\n",
    "\n",
    "metric_names = [d['Name'] for d in metric_definitions]\n",
    "num_metrics = len(metric_names)\n",
    "\n",
    "plt.figure(figsize=(16, 5))\n",
    "for i in range(num_metrics):\n",
    "    d = df[df['metric_name'] == metric_names[i]]\n",
    "    plt.subplot(1, num_metrics, i+1)\n",
    "    plt.plot(range(len(d)), d['value'])\n",
    "    plt.title(metric_names[i])\n",
    "    plt.xlabel('epoch')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Deployment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "scope = 'inference'\n",
    "inference_instance_type = \"ml.m5.xlarge\"\n",
    "\n",
    "# Retrieve inference docker image and inference script\n",
    "deploy_script_uri = script_uris.retrieve(model_id=model_id, model_version=model_version, script_scope=scope)\n",
    "deploy_image_uri = image_uris.retrieve(region=None, framework=None, image_scope=scope, model_id=model_id, model_version=model_version, instance_type=inference_instance_type)\n",
    "\n",
    "od_type2_endpoint_name = name_from_base(sagemaker_config['SolutionPrefix'] + \"-Type2\")\n",
    "print(od_type2_endpoint_name)\n",
    "\n",
    "od_type2_predictor = estimator.deploy(\n",
    "    initial_instance_count=1,\n",
    "    instance_type=inference_instance_type,\n",
    "    entry_point=\"inference.py\",  # entry point file in source_dir and present in deploy_source_uri\n",
    "    image_uri=deploy_image_uri,\n",
    "    source_dir=deploy_script_uri,\n",
    "    endpoint_name=od_type2_endpoint_name,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 5. Training: Finetune Type 2 (Latest) OD model with HPO\n",
    "\n",
    "The Type 2 (latest) OD model training reports Val_CrossEntropy loss and Val_SmoothL1 loss instead of mAP on the validation dataset. Since we can only specify one evaluation metric for AMT, we choose to minimize Val_CrossEntropy. It is not the standard practice for evaluating OD models, but is the best choice for now. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "tuning_job_name = sagemaker_config['SolutionPrefix'] + '-JS-HPO'\n",
    "\n",
    "scope = \"training\"\n",
    "\n",
    "objective_metric_name = \"Val_CrossEntropy\" if model_id.startswith(\"mxnet\") else \"val_loss\"\n",
    "hyperparameter_ranges = {\n",
    "    \"adam-learning-rate\": ContinuousParameter(1e-4, 1e-1, scaling_type=\"Logarithmic\"),\n",
    "}\n",
    "\n",
    "print(\"Train Type 2 (latest) OD model with HPO -------------------\")\n",
    "\n",
    "for training_instance_type in [\"ml.p3.2xlarge\", \"ml.g4dn.xlarge\", \"ml.g5.xlarge\", \"ml.m5.2xlarge\"]:\n",
    "    try:\n",
    "        train_image_uri = image_uris.retrieve(region=None, framework=None, image_scope=scope, model_id=model_id, model_version=model_version, instance_type=training_instance_type,)\n",
    "\n",
    "        estimator = Estimator(\n",
    "            role=role,\n",
    "            image_uri=train_image_uri,\n",
    "            source_dir=train_script_uri,\n",
    "            model_uri=train_model_uri,\n",
    "            entry_point=\"transfer_learning.py\",\n",
    "            hyperparameters=hyperparameters,\n",
    "            instance_count=1,\n",
    "            instance_type=training_instance_type,\n",
    "            max_run=36000,\n",
    "            output_path=s3_output_location,\n",
    "            base_job_name=job_name_prefix,\n",
    "            metric_definitions=metric_definitions,\n",
    "        )\n",
    "\n",
    "        tuner = HyperparameterTuner(\n",
    "            estimator,\n",
    "            objective_metric_name,\n",
    "            hyperparameter_ranges,\n",
    "            metric_definitions,\n",
    "            objective_type=\"Minimize\",\n",
    "            max_jobs=max_jobs,\n",
    "            max_parallel_jobs=max_parallel_jobs,  # <- increase this number if more EC2 instances are available\n",
    "            base_tuning_job_name=tuning_job_name,\n",
    "        )\n",
    "\n",
    "        tuner.fit({\"training\": s3_input_train}, logs=True)\n",
    "        \n",
    "    except exceptions.CapacityError as e:\n",
    "        print(\"Training exception:\", e)\n",
    "        print(f\"{training_instance_type} is not available !!\")\n",
    "        continue\n",
    "    except exceptions.UnexpectedStatusException as e:\n",
    "        print(\"Training exception:\", e)\n",
    "        continue\n",
    "    except ClientError as e:\n",
    "        print(\"Training Exception:\", e)\n",
    "        continue\n",
    "    else:\n",
    "        print(f\"Instance {training_instance_type} is available !\")\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "tuning_job_name = tuner.latest_tuning_job.name\n",
    "print('tuning_job_name', tuning_job_name)\n",
    "\n",
    "tuner_analytics = sagemaker.HyperparameterTuningJobAnalytics(tuning_job_name)\n",
    "\n",
    "full_df = tuner_analytics.dataframe()\n",
    "\n",
    "if len(full_df) > 0:\n",
    "    df = full_df[full_df[\"FinalObjectiveValue\"] > -float(\"inf\")]\n",
    "    if len(df) > 0:\n",
    "        df = df.sort_values(\"FinalObjectiveValue\", ascending=True)\n",
    "        # filter out failed jobs.\n",
    "        df = df[df[\"TrainingJobStatus\"] == \"Completed\"] \n",
    "        print(f\"Number of training jobs completed and with valid objective: {len(df)} / {len(full_df)}\")\n",
    "        print({\"lowest\": min(df[\"FinalObjectiveValue\"]), \"highest\": max(df[\"FinalObjectiveValue\"])})\n",
    "        pd.set_option(\"display.max_colwidth\", None)  # Don't truncate TrainingJobName\n",
    "    else:\n",
    "        print(\"No training jobs have reported valid results yet.\")\n",
    "        \n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Visualize the training progress of the best job\n",
    "best_job_name = df.iloc[0]['TrainingJobName']\n",
    "print('best job:', best_job_name)\n",
    "v = df.iloc[0]['FinalObjectiveValue']\n",
    "print(f\"best job final Val_CrossEntropy = {v:.6f}\")\n",
    "\n",
    "df_best = TrainingJobAnalytics(best_job_name).dataframe()\n",
    "\n",
    "OD_Type2_metrics = list(set(df_best.metric_name.values))\n",
    "print('All metrics:', OD_Type2_metrics)\n",
    "print('ObjectiveMetric is exactly the same as Val_CrossEntropy')\n",
    "num_metrics = len(OD_Type2_metrics)\n",
    "\n",
    "plt.figure(figsize=(18, 5))\n",
    "cnt = 1\n",
    "for m in OD_Type2_metrics:\n",
    "    if m != 'ObjectiveMetric':\n",
    "\n",
    "        d = df_best[df_best['metric_name'] == m]\n",
    "        \n",
    "        plt.subplot(1, num_metrics-1, cnt)\n",
    "        plt.plot(range(len(d[\"value\"])), d['value'])\n",
    "        plt.title(m)\n",
    "        plt.xlabel('epoch')\n",
    "        plt.ylabel(m)\n",
    "        cnt += 1\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "scope = 'inference'\n",
    "inference_instance_type = \"ml.m5.xlarge\"\n",
    "\n",
    "od_type2_hpo_endpoint_name = name_from_base(sagemaker_config['SolutionPrefix'] + f\"-Type2-HPO-{model_id}\")\n",
    "print(od_type2_hpo_endpoint_name)\n",
    "\n",
    "od_type2_hpo_predictor = tuner.deploy(\n",
    "    initial_instance_count=1,\n",
    "    instance_type=inference_instance_type,\n",
    "    entry_point=\"inference.py\",  # entry point file in source_dir and present in deploy_source_uri\n",
    "    image_uri=deploy_image_uri,\n",
    "    source_dir=deploy_script_uri,\n",
    "    endpoint_name=od_type2_hpo_endpoint_name,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 6. Inference and Model Comparison\n",
    "\n",
    "We will compare model performance both visually and numerically. \n",
    "1. Visually, we sample images from the test data, one image from each category, and show the predicted bounding boxes, their predicted categories, and the confidence scores.\n",
    "2. Numerically, we compute mAP on the pre-allocated test data. This is a fair comparison because we use the same metric and evaluate on the same test data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib.patches as patches\n",
    "from PIL import Image, ImageColor\n",
    "from utils import query_Type2, query_Type1, query_DDN, plot_results\n",
    "\n",
    "categories = {\n",
    "    1: 'crazing',\n",
    "    2: 'inclusion',\n",
    "    3: 'pitted_surface',\n",
    "    4: 'patches',\n",
    "    5: 'rolled-in_scale',\n",
    "    6: 'scratches'\n",
    "}\n",
    "\n",
    "# Obtain all images and their annotations info\n",
    "with open('test_annotations.json', \"r\") as f:\n",
    "    info = json.load(f)\n",
    "\n",
    "# images = {image['file_name']: image['id'] for image in info['images']}\n",
    "images = info[\"images\"]\n",
    "annotations = info['annotations']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Visual comparison\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "image_dir = \"raw_neu_det/NEU-DET/IMAGES/\"\n",
    "\n",
    "image_name = \"patches_116.jpg\"\n",
    "# image_name = \"crazing_103.jpg\"\n",
    "# image_name = \"inclusion_166.jpg\"\n",
    "# image_name = \"pitted_surface_101.jpg\"\n",
    "# image_name = \"rolled-in_scale_1.jpg\"\n",
    "# image_name = \"scratches_158.jpg\"\n",
    "\n",
    "image = os.path.join(image_dir, image_name)\n",
    "\n",
    "\n",
    "# Ground truth\n",
    "image_info = [img for img in images if img[\"file_name\"] == image_name][0]\n",
    "bboxes = [a for a in annotations if a['image_id'] == image_info[\"id\"]]\n",
    "\n",
    "# dictionary of {endpoint_name_i: {'normalized_boxes': xxx, 'classes_names': yyy, 'confidences': zzz}, endpoint_name_2: {...}}\n",
    "d = {}\n",
    "\n",
    "# Inference. Could find all endpoints from Inference / Endpoints in the Sagemaker Dashboard \n",
    "for endpoint_name in [od_type2_endpoint_name, od_type2_hpo_endpoint_name, od_type1_endpoint_name, od_type1_hpo_endpoint_name]:\n",
    "    query_function = query_Type2 if \"Type2\" in endpoint_name else query_Type1\n",
    "    normalized_boxes, classes_names, confidences = query_function(image, endpoint_name=endpoint_name, num_predictions=len(bboxes))    \n",
    "    d[endpoint_name] = {\n",
    "        'normalized_boxes': normalized_boxes,\n",
    "        'classes_names': classes_names,\n",
    "        'confidences': confidences\n",
    "    } \n",
    "\n",
    "# # Assume you have deployed the DDN model in 1_retrain_from_checkpoint.ipynb with name sagemaker-soln-{.*}-finetuned-endpoint\n",
    "# od_ddn_endpoint_name = ...\n",
    "# unnormalized_boxes, classes_names, confidences = query_DDN(image, endpoint_name=od_ddn_endpoint_name, num_predictions=len(bboxes))    \n",
    "# d[od_ddn_endpoint_name] = {\n",
    "#     'normalized_boxes': unnormalized_boxes,  # it is actually unnormalized bbox\n",
    "#     'classes_names': classes_names,\n",
    "#     'confidences': confidences\n",
    "# }\n",
    "\n",
    "plot_results(image, bboxes, categories, d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "###  Numerical comparison\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# 1. Convert test_annotations.json to the ground truth format that pycocotools can consume\n",
    "\n",
    "from pycocotools.coco import COCO\n",
    "from pycocotools.cocoeval import COCOeval\n",
    "\n",
    "annFile = \"./test_annotations.json\"\n",
    "ground_truth_annFile = convert_to_pycocotools_ground_truth(annFile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "# 2. Use one endpoint to predict all test images\n",
    "# run the following two cells for all endpoints\n",
    "\n",
    "endpoint_name = od_type1_endpoint_name\n",
    "# endpoint_name = od_type1_hpo_endpoint_name\n",
    "# endpoint_name = od_type2_endpoint_name\n",
    "# endpoint_name = od_type2_hpo_endpoint_name\n",
    "# endpoint_name = 'sagemaker-soln-...-finetuned-endpoint'  # DDN model\n",
    "\n",
    "print(\"endpoint: \", endpoint_name)\n",
    "\n",
    "if \"Type2\" in endpoint_name:\n",
    "    query_function = query_Type2\n",
    "    output_file = \"results/type2_results.json\" if endpoint_name == od_type2_endpoint_name else \"results/type2_hpo_results.json\"\n",
    "elif \"Type1\" in endpoint_name:\n",
    "    query_function = query_Type1\n",
    "    output_file = \"results/type1_results.json\" if endpoint_name == od_type1_endpoint_name else \"results/type1_hpo_results.json\"\n",
    "elif endpoint_name.endswith('finetuned-endpoint'):\n",
    "    query_function = query_DDN\n",
    "    output_file = \"results/ddn_results.json\"\n",
    "else:\n",
    "    ValueError(\"Un-recognized endpoint\")\n",
    "print(\"output file:\", output_file)\n",
    "\n",
    "coco_results = []\n",
    "\n",
    "for i, img in enumerate(images):\n",
    "    \n",
    "    if i % 50 == 0:\n",
    "        print(f\"{i} / {len(images)} done\")\n",
    "    \n",
    "    image_name = img[\"file_name\"]\n",
    "    image = os.path.join(image_dir, image_name)\n",
    "    \n",
    "    prediction_boxes, classes_names, confidences = query_function(image, endpoint_name=endpoint_name, num_predictions=100)\n",
    "    \n",
    "    if query_function != query_DDN:\n",
    "        # For non-DDN model, need to rescale to original size for each bbox\n",
    "        image_info = [img for img in images if img[\"file_name\"] == image_name][0]\n",
    "        W, H = image_info[\"width\"], image_info[\"height\"]\n",
    "\n",
    "        prediction_boxes = [\n",
    "            [xmin * W, ymin * H, (xmax-xmin) * W, (ymax-ymin) * H] \n",
    "            for (xmin, ymin, xmax, ymax) in prediction_boxes\n",
    "        ]\n",
    "\n",
    "    coco_results.extend(\n",
    "        [\n",
    "            {\n",
    "                \"image_id\": img[\"id\"],\n",
    "                \"category_id\": classes_names[k],\n",
    "                \"bbox\": box,\n",
    "                \"score\": confidences[k],\n",
    "            }\n",
    "            for k, box in enumerate(prediction_boxes)\n",
    "        ]\n",
    "    )\n",
    "    \n",
    "print(f'Total predictions for {len(images)} images:', len(coco_results))\n",
    "\n",
    "jsonString = json.dumps(coco_results)\n",
    "with open(output_file, \"w\") as f:\n",
    "    f.write(jsonString)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# 3. Calculate mean Average Precision (mAP) on the test data\n",
    "# CoCoeval reports a table of metric values, use the first row result to compare models\n",
    "\n",
    "cocoGt = COCO(ground_truth_annFile)\n",
    "cocoDt = cocoGt.loadRes(output_file)\n",
    "cocoEval = COCOeval(cocoGt, cocoDt, \"bbox\")\n",
    "\n",
    "imgIds = sorted(cocoGt.getImgIds())\n",
    "cocoEval.params.imgIds = imgIds\n",
    "\n",
    "cocoEval.evaluate()\n",
    "cocoEval.accumulate()\n",
    "cocoEval.summarize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "![Numerical comparison!](../docs/numerical.png \"Numerical comparison\")\n",
    "\n",
    "If you predict all test images using all endpoints, you end up with this table. The pycocotools package reports more metric values. We wil focus on row 1 - the mAP averaged over all IoU thresholds, all recall thresholds, all region sizes (small, medium, large), and all numbers of predicted bbox (1, 10, and 100), and all object categories. It's the [standard practice](https://cocodataset.org/#detection-eval) to use this metric for evaluating object detection algorithms."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 7. Clean Up the Endpoints\n",
    "\n",
    "When you are done with the endpoint, you should clean it up.\n",
    "\n",
    "All of the training jobs, models and endpoints we created can be viewed through the SageMaker console of your AWS account."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Delete the SageMaker endpoint\n",
    "od_type1_predictor.delete_model()\n",
    "od_type1_predictor.delete_endpoint()\n",
    "\n",
    "od_type1_hpo_predictor.delete_model()\n",
    "od_type1_hpo_predictor.delete_endpoint()\n",
    "\n",
    "od_type2_predictor.delete_model()\n",
    "od_type2_predictor.delete_endpoint()\n",
    "\n",
    "od_type2_hpo_predictor.delete_model()\n",
    "od_type2_hpo_predictor.delete_endpoint()\n",
    "\n",
    "# You should go to the Sagemaker console and manually delete the DDN model endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 8. Conclusion\n",
    "\n",
    "Both visual and numerical comparison confirm that the Type 2 (latest) OD model or Type 2 (latest) OD + HPO performs the best. \n",
    "\n",
    "1. DDN model is the worst. It was trained from scratch using the target dataset, which is a very small dataset of 1,800 images for 6 categories, the training data consists of only 64% of this small dataset.\n",
    "2. The built-in Sagemaker OD models were pre-trained on large-scale dataset, e.g., the ImageNet dataset includes 14,197,122 images for 21,841 categories, and the PASCAL VOC dataset includes 11,530 images for 20 categories. The pre-trained models have learned rich and diverse low level features, and can efficiently transfer knowledge to finetuned models and focus on learning high-level semantic features for the target dataset.\n",
    "3. HPO is extremely effective, especially for models with large hyperparameter search spaces. Since we finetuned on three hyperparameters (learning rate, momentum, and weight decay) for the Type 1 (legacy) OD models and only one hyperparameter (adam learning rate) for the Type 2 (latest) OD model, there is relatively larger room for improvement for the Type 1 (legacy) OD model and we do observe larger performance enhancement. Of course, we need to trade off model performance with budget (compute resource and training time) when running HPO.\n",
    "4. In terms of training time, for the steel surface dataset, training the Type 1 (legacy) OD model took 34 min, Type 2 (latest) OD model took 1 hour, and DDN model took 8+ hours. It indicates finetuning a pre-trained model is much more efficient.\n",
    "5. In summary, finetuning a pretrained model is both more efficient and more performant, we suggest taking advantage of the pre-trained Sagemaker built-in models and finetune on your target datasets.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (PyTorch 1.10 Python 3.8 CPU Optimized)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-west-2:236514542706:image/pytorch-1.10-cpu-py38"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}