{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Raspberry Pi 3B x AWS NEO Image Classification 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 mobilenetv2 model from gluon imagenet classifier model zoo using AWS Neo for Raspberry Pi 3B. 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 image classification model from S3 bucket to device. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'mobilenetv2_0.25'\n",
    "model = model_name + '.tar.gz'\n",
    "model_zoo = 'gluon_imagenet_classifier'"
   ]
  },
  {
   "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 `pi-demo` to store pretrained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bucket = 'pi-demo'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": "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', region_name='us-west-2')\n",
    "role_name = 'pi-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,224,224]}'\n",
    "framework = 'mxnet'\n",
    "target_device = 'rasp3b'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "compilation_job_name = 'pi-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": [
    "!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": [
    "!wget -O test.jpg http://www.vision.caltech.edu/Image_Datasets/Caltech256/images/080.frog/080_0001.jpg\n",
    "\n",
    "import json\n",
    "import numpy as np\n",
    "    \n",
    "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((224, 224)))\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(10):\n",
    "    model.run(input_data)\n",
    "\n",
    "print('Testing inference...')\n",
    "start_time = time.time()\n",
    "out = model.run(input_data) #need to be a list of input arrays matching input names\n",
    "index = np.argmax(out[0][0,:])\n",
    "prob = np.amax(out[0][0,:])\n",
    "print('inference time is ' + str((time.time()-start_time)) + ' seconds')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load names for ImageNet classes\n",
    "object_categories = {}\n",
    "with open(\"imagenet1000_clsidx_to_labels.txt\", \"r\") as f:\n",
    "    for line in f:\n",
    "        key, val = line.strip().split(':')\n",
    "        object_categories[key] = val\n",
    "print(index)\n",
    "print(\"Result: label - \" + object_categories[str(index)] + \" probability - \" + str(prob))"
   ]
  }
 ],
 "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
}