{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NXP imx8qm x AWS NEO Object Detection Example\n",
    "\n",
    "1. [Introduction](#Introduction)\n",
    "2. [Compile model using NEO](#Compile-model-using-NEO)\n",
    "3. [Inference on device](#Inference-on-device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "This notebook will demo how to compile pretrained Gluoncv ssd mobilenet model using AWS Neo for NXP imx8qm. First, we need download the pretained model, in this example we use `ssd_512_mobilenet1.0_voc`, then use Neo to cpmpile this model. At last, we will deploy compiled model to device and do inference using the Neo Deep Learning Runtime.\n",
    "\n",
    "To get started, we need to set up the environment for AWS S3 permissions, configurations, and so on. Please refer to [Configuration](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/quickstart.html#configuration) for more information."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use Boto 3, first import it and tell it what service we are going to use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "sm = boto3.client('sagemaker', region_name='us-west-2')\n",
    "s3 = boto3.client('s3', region_name='us-west-2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "account = boto3.client('sts').get_caller_identity().get('Account')\n",
    "print(account)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get the pretrained model \n",
    "\n",
    "To avoid installing frameworks like mxnet, gluoncv on device to download latest pretained models, we store some pretrained models in a S3 bucket with public access. Download pretrained gluoncv ssd model from S3 bucket to device. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'ssd_512_mobilenet1.0_voc'\n",
    "model = model_name + '.tar.gz'\n",
    "model_zoo = 'gluon_cv_object_detection'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3.download_file('neo-ai-dlr-test-artifacts', \n",
    "                 'neo-ai-notebook/{}/{}'.format(model_zoo, model), \n",
    "                 model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Upload model to S3 bucket"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a S3 bucket `imx8qm-demo` to store pretrained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bucket = 'imx8qm-demo'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if boto3.resource('s3').Bucket(bucket) not in boto3.resource('s3').buckets.all():\n",
    "    s3.create_bucket(\n",
    "        Bucket=bucket,\n",
    "        CreateBucketConfiguration={\n",
    "            'LocationConstraint': 'us-west-2'\n",
    "        }\n",
    "    )\n",
    "else:\n",
    "    print('Bucket %s already exists' %bucket)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Upload the pretained model to S3 bucket just created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3.upload_file(model, bucket, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create IAM role\n",
    "In order to use the sagemaker service and have access to S3 bucket, we need to create a IAM role."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iam = boto3.client('iam')\n",
    "role_name = 'imx8qm-demo-test-role'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "policy = {\n",
    "    'Statement': [{\n",
    "        'Action': 'sts:AssumeRole',\n",
    "        'Effect': 'Allow',\n",
    "        'Principal': {'Service': 'sagemaker.amazonaws.com'}},\n",
    "        ],  \n",
    "     'Version': '2012-10-17'}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "roles = iam.list_roles()\n",
    "role_arn = None\n",
    "for role in roles['Roles']:\n",
    "    if role['RoleName'] == role_name:\n",
    "        role_arn = role['Arn']\n",
    "        \n",
    "if role_arn == None:\n",
    "    new_role = iam.create_role(\n",
    "        AssumeRolePolicyDocument=json.dumps(policy),\n",
    "        Path='/',\n",
    "        RoleName=role_name,\n",
    "    )\n",
    "    role_arn = new_role['Role']['Arn']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iam.attach_role_policy(\n",
    "    RoleName=role_name,\n",
    "    PolicyArn='arn:aws:iam::aws:policy/AmazonSageMakerFullAccess'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iam.attach_role_policy(\n",
    "    RoleName=role_name,\n",
    "    PolicyArn='arn:aws:iam::aws:policy/AmazonS3FullAccess'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compile model using NEO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_output_location = 's3://{}/output'.format(bucket)\n",
    "data_shape = '{\"data\":[1,3,512,512]}'\n",
    "framework = 'mxnet'\n",
    "target_device = 'imx8qm'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "compilation_job_name = 'imx8qm-demo'+ str(time.time()).split('.')[0]\n",
    "print('Compilation job for %s started' % compilation_job_name)\n",
    "\n",
    "response = sm.create_compilation_job(\n",
    "        CompilationJobName=compilation_job_name,\n",
    "        RoleArn=role_arn,\n",
    "        InputConfig={\n",
    "            'S3Uri': 's3://{}/{}'.format(bucket, model),\n",
    "            'DataInputConfig': data_shape,\n",
    "            'Framework': framework.upper()\n",
    "        },\n",
    "        OutputConfig={\n",
    "            'S3OutputLocation': s3_output_location,\n",
    "            'TargetDevice': target_device \n",
    "        },\n",
    "        StoppingCondition={\n",
    "            'MaxRuntimeInSeconds': 900\n",
    "        }\n",
    "    )\n",
    "\n",
    "print(response)\n",
    "\n",
    "# Poll every 30 sec\n",
    "while True:\n",
    "    response = sm.describe_compilation_job(CompilationJobName=compilation_job_name)\n",
    "    if response['CompilationJobStatus'] == 'COMPLETED':\n",
    "        break\n",
    "    elif response['CompilationJobStatus'] == 'FAILED':\n",
    "        raise RuntimeError('Compilation failed')\n",
    "    print('Compiling ...')\n",
    "    time.sleep(30)\n",
    "print('Done!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference on device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download compiled model from S3 to device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "object_path = 'output/{}-{}.tar.gz'.format(model_name, target_device)\n",
    "neo_compiled_model = 'compiled-'+ model\n",
    "s3.download_file(bucket, object_path, neo_compiled_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash \n",
    "mkdir compiled_model\n",
    "tar -xf $neo_compiled_model -C ./compiled_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use DLR to read compiled model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dlr import DLRModel\n",
    "import numpy as np\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the model\n",
    "model_path = \"./compiled_model\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cpu'\n",
    "model = DLRModel(model_path, device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download an image to prepare for predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "wget -O test.jpg https://upload.wikimedia.org/wikipedia/commons/c/c6/Newark-broad-street.jpg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_name = \"test.jpg\"\n",
    "\n",
    "# test image\n",
    "from IPython.display import Image\n",
    "Image(file_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Image pre-process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import PIL.Image\n",
    "image = PIL.Image.open(file_name)\n",
    "\n",
    "image = np.asarray(image.resize((512, 512)))\n",
    "\n",
    "# Normalize\n",
    "mean_vec = np.array([0.485, 0.456, 0.406])\n",
    "stddev_vec = np.array([0.229, 0.224, 0.225])\n",
    "image = (image/255- mean_vec)/stddev_vec\n",
    "\n",
    "# Transpose\n",
    "if len(image.shape) == 2:  # for greyscale image\n",
    "    image = np.expand_dims(image, axis=2)\n",
    "    \n",
    "image = np.rollaxis(image, axis=2, start=0)[np.newaxis, :]\n",
    "\n",
    "print(image.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference and prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#flatten within a input array\n",
    "input_data = {'data': image}\n",
    "\n",
    "# dry run\n",
    "for _ in range(5):\n",
    "    model.run(input_data)\n",
    "\n",
    "print('Testing inference...')\n",
    "start_time = time.time()\n",
    "detection = model.run(input_data) #need to be a list of input arrays matching input names\n",
    "print('inference time is ' + str((time.time()-start_time)) + ' seconds')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "object_categories = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', \n",
    "                     'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person', \n",
    "                     'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']\n",
    "\n",
    "obj = detection[0][0]\n",
    "score = detection[1][0] \n",
    "bbox = detection[2][0]\n",
    "\n",
    "# print the first 10 detetction\n",
    "for i in range(10):\n",
    "    print(object_categories[int(obj[i][0])] + ' with score ' + str(score[i][0]))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}