{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fairmot Model Inference in Amazon SageMaker\n",
    "\n",
    "This notebook will demonstrate how to create an endpoint for real time inference with the trained FairMOT model.\n",
    "\n",
    "## 1. SageMaker Initialization \n",
    "First we upgrade SageMaker to the latest version. If your notebook is already using latest Sagemaker 2.x API, you may skip the next cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "! pip install --upgrade pip\n",
    "! python3 -m pip install --upgrade sagemaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import sagemaker\n",
    "from sagemaker import get_execution_role\n",
    "\n",
    "role = (\n",
    "    get_execution_role()\n",
    ")  # provide a pre-existing role ARN as an alternative to creating a new role\n",
    "print(f\"SageMaker Execution Role:{role}\")\n",
    "\n",
    "client = boto3.client('sts')\n",
    "account = client.get_caller_identity()['Account']\n",
    "print(f'AWS account:{account}')\n",
    "\n",
    "session = boto3.session.Session()\n",
    "aws_region = session.region_name\n",
    "print(f\"AWS region:{aws_region}\")\n",
    "\n",
    "container_name = \"container-serving\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Build and Push Amazon SageMaker Serving Container Images\n",
    "\n",
    "For this step, the [IAM Role](https://docs.aws.amazon.com/IAM/latest/UserGuide/id_roles.html) attached to this notebook instance needs full access to [Amazon ECR service](https://aws.amazon.com/ecr/). We use the implementation of [FairMOT](https://github.com/ifzhang/FairMOT) to create our own container."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Docker Environment Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because the volume size of container may be larger than the available size in root directory of the notebook instance, we need to put the directory of docker data into the ```/home/ec2-user/SageMaker/docker``` directory.\n",
    "\n",
    "By default, the root directory of docker is set as ```/var/lib/docker/```. We need to change the directory of docker to ```/home/ec2-user/SageMaker/docker```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!bash ./prepare-docker.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Build and Push FairMOT Serving Container Image\n",
    "\n",
    "Use [`./container-serving/build_tools/build_and_push.sh`](./container-serving/build_tools/build_and_push.sh) script to build the [FairMOT](https://github.com/ifzhang/FairMOT) <b>serving</b> container image and push it to Amazon ECR. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ./{container_name}/build_tools/build_and_push.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using your *AWS region* as argument, run the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "! ./{container_name}/build_tools/build_and_push.sh {aws_region}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fairmot_image = f\"{account}.dkr.ecr.{aws_region}.amazonaws.com/fairmot-sagemaker:pytorch1.8-serving\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Create Inference Endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1 Define Amazon SageMaker Model\n",
    "Next, we define an Amazon SageMaker Model that we will serve from an Amazon SageMaker Endpoint. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"fairmot-model-1\" # set the name of the model, like fairmot-model-1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can get the S3 URI of the trained model in Training job console once the training job gets finished (note that the training job is launched in [`fairmot-training.ipynb`](fairmot-training.ipynb)), and then set `ModelDataUrl` to the S3 URI of the trained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Restore the s3 uri of trained model\n",
    "%store -r s3_model_uri\n",
    "#s3_model_uri=\"s3://{bucket_name}/{predix_model}/model.tar.gz\" you can define the model URI manually.\n",
    "\n",
    "serving_container_def = {\n",
    "    'Image': fairmot_image,\n",
    "    'ModelDataUrl': s3_model_uri,\n",
    "    'Mode': 'SingleModel',\n",
    "    'Environment': {\n",
    "                    'SM_MODEL_DIR' : '/opt/ml/model',\n",
    "                   }\n",
    "}\n",
    "\n",
    "sagemaker_session = sagemaker.session.Session(boto_session=session)\n",
    "create_model_response = sagemaker_session.create_model(name=model_name, \n",
    "                                                       role=role, \n",
    "                                                       container_defs=serving_container_def)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### 3.2 Create Endpoint Configuration\n",
    " Next, we set the name of the Amaozn SageMaker hosted service endpoint configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint_config_name = f\"{model_name}-endpoint-config\"\n",
    "print(endpoint_config_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then create the Amazon SageMaker hosted service endpoint configuration that uses one instance of `ml.p3.2xlarge` to serve the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epc = sagemaker_session.create_endpoint_config(\n",
    "    name=endpoint_config_name,\n",
    "    model_name=model_name,\n",
    "    initial_instance_count=1,\n",
    "    instance_type=\"ml.p3.2xlarge\",\n",
    ")\n",
    "print(epc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we specify the Amazon SageMaker endpoint name for the endpoint used to serve the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint_name = f\"{model_name}-endpoint\"\n",
    "print(endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3 Create Endpoint\n",
    "In this step, we create the Amazon SageMaker endpoint using the endpoint configuration we created above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ep = sagemaker_session.create_endpoint(\n",
    "    endpoint_name=endpoint_name, config_name=endpoint_config_name, wait=True\n",
    ")\n",
    "print(ep)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Test Endpoint\n",
    "### 4.1 Visualization Helper Functions\n",
    "Draw the bounding box and ID for each tracked object in the raw frames."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_color(idx):\n",
    "    idx = idx * 3\n",
    "    color = ((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255)\n",
    "\n",
    "    return color\n",
    "\n",
    "def draw_res(tracker_dict, frame, frame_id, image_w):\n",
    "    i = 0\n",
    "    indexIDs = []\n",
    "    boxes = []\n",
    "    person_num = 0\n",
    "    conf = None\n",
    "    text_scale = max(1, image_w / 1600.)\n",
    "    text_thickness = 1\n",
    "    line_thickness = max(1, int(image_w/ 500.))\n",
    "    for track_id, tlwh in tracker_dict.items():\n",
    "        indexIDs.append(track_id)\n",
    "        x1, y1, w, h = tlwh\n",
    "        intbox = tuple(map(int, (x1, y1, x1 + w, y1 + h)))\n",
    "        color = get_color(abs(int(track_id)))\n",
    "        cv2.rectangle(frame, intbox[0:2], intbox[2:4], color=color, thickness=line_thickness)\n",
    "        cv2.putText(frame, str(track_id), (intbox[0], intbox[1] + 30), cv2.FONT_HERSHEY_PLAIN, text_scale, (0, 0, 0),thickness=1)\n",
    "        cv2.putText(frame, 'frame:{}'.format(frame_id), (int(25), int(25)),0, text_scale, (0,0,255),1)\n",
    "        i += 1\n",
    "    return frame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2 Invoke endpoint\n",
    "Next, we use [MOT16-03](https://motchallenge.net/sequenceVideos/MOT16-04-raw.webm) from MOT challenge to test our endpoint. We create a directory `datasets` in the root directory of this project for saving the processed result, and then download it to `datasets` directory with MP4 format from [FairMOT](https://raw.githubusercontent.com/ifzhang/FairMOT/master/videos/MOT16-03.mp4)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p datasets\n",
    "!wget https://raw.githubusercontent.com/ifzhang/FairMOT/master/videos/MOT16-03.mp4 -O datasets/MOT16-03.mp4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After preparing the test data, we invoke the endpoint to run the real time inferece on the test video. It takes about 150 seconds to complete all of the inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import boto3\n",
    "import base64\n",
    "import json\n",
    "import cv2\n",
    "import time\n",
    "import os\n",
    "\n",
    "client = boto3.client(\"sagemaker-runtime\")\n",
    "\n",
    "data_path = \"datasets/MOT16-03.mp4\" \n",
    "cap = cv2.VideoCapture(data_path)\n",
    "frame_w  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "frame_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "\n",
    "fourcc = cv2.VideoWriter_fourcc(*'MP4V')\n",
    "file_path = os.path.join('datasets', 'test.mp4')\n",
    "out = cv2.VideoWriter(file_path, fourcc, 25, (frame_w, frame_h))\n",
    "\n",
    "processing_time = 0\n",
    "frame_id = 0\n",
    "\n",
    "while True:\n",
    "    ret, frame = cap.read()\n",
    "    if ret != True:\n",
    "        break\n",
    "    \n",
    "    frame_path = f'datasets/{frame_id}.jpg'\n",
    "    cv2.imwrite(frame_path, frame)\n",
    "    \n",
    "    with open(frame_path, \"rb\") as image_file:\n",
    "        img_data = base64.b64encode(image_file.read())\n",
    "        data = {\"frame_id\": frame_id}\n",
    "        data[\"frame_data\"] = img_data.decode(\"utf-8\")\n",
    "        if frame_id == 0:\n",
    "            data[\"frame_w\"] = frame_w\n",
    "            data[\"frame_h\"] = frame_h\n",
    "            data[\"batch_size\"] = 1 # for multiple stream\n",
    "        body = json.dumps(data).encode(\"utf-8\")\n",
    "    \n",
    "    os.remove(frame_path)\n",
    "    request_time=time.time()\n",
    "    response = client.invoke_endpoint(\n",
    "        EndpointName=endpoint_name, ContentType=\"application/json\", Accept=\"application/json\", Body=body\n",
    "    )\n",
    "    if frame_id > 0:\n",
    "        processing_time += (time.time() - request_time)\n",
    "    print(f'frame-{frame_id} Processing time: {(time.time() - request_time)}')\n",
    "    body = response[\"Body\"].read()\n",
    "    msg = body.decode(\"utf-8\")\n",
    "    data = json.loads(msg)\n",
    "    frame_res = draw_res(data[0], frame, frame_id, frame_w)\n",
    "    out.write(frame_res)\n",
    "    frame_id += 1\n",
    "\n",
    "out.release()\n",
    "cap.release()\n",
    "print('average processing time: ', processing_time/frame_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The response from the endpoint includes the bounding box information and ID for each person. You can download the processed video from `datasets` to check from the local instance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Delete SageMaker Endpoint, Endpoint Config and Model\n",
    "<span style=\"color:red;\">**If you are done testing, delete the deployed Amazon SageMaker endpoint, endpoint config, and the model below.**</span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sagemaker_session.delete_endpoint(endpoint_name=endpoint_name)\n",
    "sagemaker_session.delete_endpoint_config(endpoint_config_name=endpoint_config_name)\n",
    "sagemaker_session.delete_model(model_name=model_name)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_tensorflow_p36",
   "language": "python",
   "name": "conda_tensorflow_p36"
  },
  "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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}