{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fairmot Batch Inference in Amazon SageMaker\n",
    "\n",
    "In this notebook, we will walk through the batch inference with the FairMOT model. Because [SageMaker Batch Transform only can run the batch transform job finishing within 600 seconds](https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-batch-code.html), we use SageMaker processing to run the batch inference.\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-batch-inference\""
   ]
  },
  {
   "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/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Docker Environment Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because the volume size of container may exceed the available size in the 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```. You can skip this step if you have done docker environment preparation in [`fairmot-training.ipynb`](fairmot-training.ipynb)."
   ]
  },
  {
   "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-batch-inference/build_tools/build_and_push.sh`](./container-batch-inference/build_tools/build_and_push.sh) script to build and push the [FairMOT](https://github.com/ifzhang/FairMOT) <b>batch inference</b> container image to Amazon ECR. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "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-batch-inference\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Run Batch Inference\n",
    "\n",
    "Create an instance of SageMaker Processing with the built container above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.processing import ScriptProcessor\n",
    "\n",
    "script_processor = ScriptProcessor(\n",
    "                image_uri=fairmot_image,\n",
    "                role=role,\n",
    "                instance_count=1,\n",
    "                instance_type='ml.p3.2xlarge',\n",
    "                command=['python3'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the S3 URIs for the test data, the trained model and the result data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bucket_name = sagemaker.Session().default_bucket() \n",
    "\n",
    "# Restore the s3 uri of the trained model\n",
    "%store -r s3_model_uri\n",
    "#s3_model_uri=\"s3://{bucket_name}/{predix_model}/model.tar.gz\" you can define the s3 uri of the trained model manually.\n",
    "\n",
    "s3_input=f\"s3://{bucket_name}/fairmot/sagemaker/video\"\n",
    "s3_output=f\"s3://{bucket_name}/fairmot/sagemaker/output/batch-inference\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use [MOT16-03](https://motchallenge.net/sequenceVideos/MOT16-04-raw.webm) from MOT challenge to test the batch inference. First, we download the video to the notebook instance and then upload it to the defined s3 bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/ifzhang/FairMOT/master/videos/MOT16-03.mp4\n",
    "!aws s3 cp MOT16-03.mp4 {s3_input}/MOT16-03.mp4\n",
    "!rm MOT16-03.mp4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we run the sagemaker processing job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from sagemaker.processing import ProcessingInput, ProcessingOutput\n",
    "script_processor.run(\n",
    "    code='./container-batch-inference/processing.py',\n",
    "    inputs=[\n",
    "        ProcessingInput(source=s3_input, destination=\"/opt/ml/processing/input\"),\n",
    "        ProcessingInput(source=s3_model_uri, destination=\"/opt/ml/processing/model\"),\n",
    "    ], \n",
    "    outputs=[\n",
    "        ProcessingOutput(source='/opt/ml/processing/output', destination=s3_output),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can check the result saved in `s3://{bucket-name}/fairmot/sagemaker/output/batch-inference`."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_p39",
   "language": "python",
   "name": "conda_pytorch_p39"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}