{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Amazon Augmented AI (Amazon A2I) integration with Amazon Textract's Analyze Document [Example]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visit https://github.com/aws-samples/amazon-a2i-sample-jupyter-notebooks for all A2I Sample Notebooks\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. [Introduction](#Introduction)\n",
    "2. [Prerequisites](#Prerequisites)\n",
    "    1. [Workteam](#Workteam)\n",
    "    2. [Permissions](#Notebook-Permission)\n",
    "3. [Client Setup](#Client-Setup)\n",
    "4. [Sample Data](#Sample-Data)\n",
    "    1. [Download sample images](#Download-sample-images)\n",
    "    2. [Upload images to S3](#Upload-images-to-S3)\n",
    "5. [Create Control Plane Resources](#Create-Control-Plane-Resources)\n",
    "    1. [Create Human Task UI](#Create-Human-Task-UI)\n",
    "    2. [Create Flow Definition](#Create-Flow-Definition)\n",
    "6. [Analyze Document with Textract](#Analyze-Document-with-Textract)\n",
    "6. [Human Loops](#Human-Loops)\n",
    "    1. [Check Status of Human Loop](#Check-Status-of-Human-Loop)\n",
    "    2. [Wait For Workers to Complete Task](#Wait-For-Workers-to-Complete-Task)\n",
    "    3. [Check Status of Human Loop](#Check-Status-of-Human-Loop)\n",
    "    4. [View Task Results](#View-Task-Results)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "Amazon Augmented AI (Amazon A2I) makes it easy to build the workflows required for human review of ML predictions. Amazon A2I brings human review to all developers, removing the undifferentiated heavy lifting associated with building human review systems or managing large numbers of human reviewers. \n",
    "\n",
    "Amazon A2I provides built-in human review workflows for common machine learning use cases, such as content moderation and text extraction from documents, which allows predictions from Amazon Rekognition and Amazon Textract to be reviewed easily. You can also create your own workflows for ML models built on Amazon SageMaker or any other tools. Using Amazon A2I, you can allow human reviewers to step in when a model is unable to make a high confidence prediction or to audit its predictions on an on-going basis. Learn more here: https://aws.amazon.com/augmented-ai/\n",
    "\n",
    "In this tutorial, we will show how you can use Amazon A2I directly within your API calls to Textract's Analyze Document API. \n",
    "\n",
    "For more in depth instructions, visit https://docs.aws.amazon.com/sagemaker/latest/dg/a2i-getting-started.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To incorporate Amazon A2I into your human review workflows, you need three resources:\n",
    "\n",
    "* A **worker task template** to create a worker UI. The worker UI displays your input data, such as documents or images, and instructions to workers. It also provides interactive tools that the worker uses to complete your tasks. For more information, see https://docs.aws.amazon.com/sagemaker/latest/dg/a2i-instructions-overview.html\n",
    "\n",
    "* A **human review workflow**, also referred to as a flow definition. You use the flow definition to configure your human workforce and provide information about how to accomplish the human review task. For built-in task types, you also use the flow definition to identify the conditions under which a review human loop is triggered. For example, with Amazon Textract can analyze text in a document using machine learning. You can use the flow definition to specify that a document will be sent to a human for content moderation review if Amazon Textracts's confidence score output is low for any or all pieces of text returned by Textract. You can create a flow definition in the Amazon Augmented AI console or with the Amazon A2I APIs. To learn more about both of these options, see https://docs.aws.amazon.com/sagemaker/latest/dg/a2i-create-flow-definition.html\n",
    "\n",
    "* A **human loop** to start your human review workflow. When you use one of the built-in task types, the corresponding AWS service creates and starts a human loop on your behalf when the conditions specified in your flow definition are met or for each object if no conditions were specified. When a human loop is triggered, human review tasks are sent to the workers as specified in the flow definition.\n",
    "\n",
    "When using a custom task type, you start a human loop using the Amazon Augmented AI Runtime API. When you call StartHumanLoop in your custom application, a task is sent to human reviewers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install Latest SDK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First, let's get the latest installations of our dependencies\n",
    "!pip install --upgrade pip\n",
    "!pip install botocore --upgrade\n",
    "!pip install boto3 --upgrade\n",
    "!pip install -U botocore"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "We need to set up the following data:\n",
    "* `region` - Region to call A2I\n",
    "* `bucket` - A S3 bucket accessible by the given role\n",
    "    * Used to store the sample images & output results\n",
    "    * Must be within the same region A2I is called from\n",
    "* `role` - The IAM role used as part of StartHumanLoop. By default, this notebook will use the execution role\n",
    "* `workteam` - Group of people to send the work to"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import botocore\n",
    "\n",
    "REGION = boto3.session.Session().region_name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Setup Bucket and Paths\n",
    "\n",
    "**Important**: The bucket you specify for `BUCKET` must have CORS enabled. You can enable CORS by adding a policy similar to the following to your Amazon S3 bucket. To learn how to add CORS to an S3 bucket, see [CORS Permission Requirement](https://docs.aws.amazon.com/sagemaker/latest/dg/a2i-permissions-security.html#a2i-cors-update) in the Amazon A2I documentation. \n",
    "\n",
    "\n",
    "```\n",
    "[{\n",
    "   \"AllowedHeaders\": [],\n",
    "   \"AllowedMethods\": [\"GET\"],\n",
    "   \"AllowedOrigins\": [\"*\"],\n",
    "   \"ExposeHeaders\": []\n",
    "}]\n",
    "```\n",
    "\n",
    "If you do not add a CORS configuration to the S3 buckets that contains your image input data, human review tasks for those input data objects will fail. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Workteam or Workforce\n",
    "BUCKET = \"<YOUR_BUCKET_NAME>\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your bucket, `BUCKET` must be located in the same AWS Region that you are using to run this notebook. This cell checks if they are located in the same Region. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Amazon S3 (S3) client\n",
    "s3 = boto3.client('s3', REGION)\n",
    "bucket_region = s3.head_bucket(Bucket=BUCKET)['ResponseMetadata']['HTTPHeaders']['x-amz-bucket-region']\n",
    "assert bucket_region == REGION, \"Your S3 bucket {} and this notebook need to be in the same region.\".format(BUCKET)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "OUTPUT_PATH = f's3://{BUCKET}/a2i-results'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A workforce is the group of workers that you have selected to label your dataset. You can choose either the Amazon Mechanical Turk workforce, a vendor-managed workforce, or you can create your own private workforce for human reviews. Whichever workforce type you choose, Amazon Augmented AI takes care of sending tasks to workers. \n",
    "\n",
    "When you use a private workforce, you also create work teams, a group of workers from your workforce that are assigned to Amazon Augmented AI human review tasks. You can have multiple work teams and can assign one or more work teams to each job."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To create your Workteam, visit the instructions here: https://docs.aws.amazon.com/sagemaker/latest/dg/sms-workforce-management.html\n",
    "\n",
    "After you have created your workteam, replace YOUR_WORKTEAM_ARN below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "WORKTEAM_ARN= \"<YOUR_WORKTEAM_ARN>\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Permissions\n",
    "\n",
    "The AWS IAM Role used to execute the notebook needs to have the following permissions:\n",
    "\n",
    "* TextractFullAccess\n",
    "* SagemakerFullAccess\n",
    "* S3 Read and Write Access to the bucket you specified in `BUCKET`. You can grant this permission by attaching a policy similar to the following to this role (replace `BUCKET` with your bucket-name).\n",
    "\n",
    "```\n",
    "{\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": [\n",
    "                \"s3:GetObject\"\n",
    "            ],\n",
    "            \"Resource\": [\n",
    "                \"arn:aws:s3:::BUCKET/*\"\n",
    "            ]\n",
    "        },\n",
    "        {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": [\n",
    "                \"s3:PutObject\"\n",
    "            ],\n",
    "            \"Resource\": [\n",
    "                \"arn:aws:s3:::BUCKET/*\"\n",
    "            ]\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "```\n",
    "* AmazonSageMakerMechanicalTurkAccess (if using MechanicalTurk as your Workforce)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker import get_execution_role\n",
    "\n",
    "# Setting Role to the default SageMaker Execution Role\n",
    "ROLE = get_execution_role()\n",
    "display(ROLE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visit: https://docs.aws.amazon.com/sagemaker/latest/dg/a2i-permissions-security.html to add the necessary permissions to your role"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Client Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we are going to setup the clients. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import io\n",
    "import json\n",
    "import uuid\n",
    "import botocore\n",
    "import time\n",
    "import botocore\n",
    "\n",
    "# Amazon SageMaker client\n",
    "sagemaker = boto3.client('sagemaker', REGION)\n",
    "\n",
    "# Amazon Textract client\n",
    "textract = boto3.client('textract', REGION)\n",
    "\n",
    "# S3 client\n",
    "s3 = boto3.client('s3', REGION)\n",
    "\n",
    "# A2I Runtime client\n",
    "a2i_runtime_client = boto3.client('sagemaker-a2i-runtime', REGION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pprint\n",
    "\n",
    "# Pretty print setup\n",
    "pp = pprint.PrettyPrinter(indent=2)\n",
    "\n",
    "# Function to pretty-print AWS SDK responses\n",
    "def print_response(response):\n",
    "    if 'ResponseMetadata' in response:\n",
    "        del response['ResponseMetadata']\n",
    "    pp.pprint(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download sample images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget 'https://github.com/aws-samples/amazon-a2i-sample-jupyter-notebooks/raw/master/document-demo.jpg' -O 'document-demo.jpg'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Upload images to S3\n",
    "\n",
    "Upload the sample images to your S3 bucket. They will be read by Textract and A2I later when the human task is created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "document = 'document-demo.jpg' "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3.upload_file(document, BUCKET, document)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Control Plane Resources"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Human Task UI\n",
    "\n",
    "Create a human task UI resource, giving a UI template in liquid html. This template will be rendered to the human workers whenever human loop is required.\n",
    "\n",
    "We are providing a simple demo template that is compatible with AWS Textract's Analyze Document API input and response.\n",
    "\n",
    "Since we are integrating A2I with Textract, we can create the template in the Console using default templates provided by A2I, to make the process easier (https://docs.aws.amazon.com/sagemaker/latest/dg/a2i-instructions-overview.html). \n",
    "\n",
    "To make things easier, the below template string is copied from the defeault template provided by Amazon A2I (found in the SageMaker Console under Worker task templates)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "template = r\"\"\"\n",
    "<script src=\"https://assets.crowd.aws/crowd-html-elements.js\"></script>\n",
    "{% capture s3_arn %}http://s3.amazonaws.com/{{ task.input.aiServiceRequest.document.s3Object.bucket }}/{{ task.input.aiServiceRequest.document.s3Object.name }}{% endcapture %}\n",
    "<crowd-form>\n",
    "  <crowd-textract-analyze-document \n",
    "      src=\"{{ s3_arn | grant_read_access }}\" \n",
    "      initial-value=\"{{ task.input.selectedAiServiceResponse.blocks }}\" \n",
    "      header=\"Review the key-value pairs listed on the right and correct them if they don't match the following document.\" \n",
    "      no-key-edit=\"\" \n",
    "      no-geometry-edit=\"\" \n",
    "      keys=\"{{ task.input.humanLoopContext.importantFormKeys }}\" \n",
    "      block-types=\"['KEY_VALUE_SET']\">\n",
    "    <short-instructions header=\"Instructions\">\n",
    "        <p>Click on a key-value block to highlight the corresponding key-value pair in the document.\n",
    "        </p><p><br></p>\n",
    "        <p>If it is a valid key-value pair, review the content for the value. If the content is incorrect, correct it.\n",
    "        </p><p><br></p>\n",
    "        <p>The text of the value is incorrect, correct it.</p>\n",
    "        <p><img src=\"https://assets.crowd.aws/images/a2i-console/correct-value-text.png\">\n",
    "        </p><p><br></p>\n",
    "        <p>A wrong value is identified, correct it.</p>\n",
    "        <p><img src=\"https://assets.crowd.aws/images/a2i-console/correct-value.png\">\n",
    "        </p><p><br></p>\n",
    "        <p>If it is not a valid key-value relationship, choose No.</p>\n",
    "        <p><img src=\"https://assets.crowd.aws/images/a2i-console/not-a-key-value-pair.png\">\n",
    "        </p><p><br></p>\n",
    "        <p>If you can’t find the key in the document, choose Key not found.</p>\n",
    "        <p><img src=\"https://assets.crowd.aws/images/a2i-console/key-is-not-found.png\">\n",
    "        </p><p><br></p>\n",
    "        <p>If the content of a field is empty, choose Value is blank.</p>\n",
    "        <p><img src=\"https://assets.crowd.aws/images/a2i-console/value-is-blank.png\">\n",
    "        </p><p><br></p>\n",
    "        <p><strong>Examples</strong></p>\n",
    "        <p>Key and value are often displayed next or below to each other.\n",
    "        </p><p><br></p>\n",
    "        <p>Key and value displayed in one line.</p>\n",
    "        <p><img src=\"https://assets.crowd.aws/images/a2i-console/sample-key-value-pair-1.png\">\n",
    "        </p><p><br></p>\n",
    "        <p>Key and value displayed in two lines.</p>\n",
    "        <p><img src=\"https://assets.crowd.aws/images/a2i-console/sample-key-value-pair-2.png\">\n",
    "        </p><p><br></p>\n",
    "        <p>If the content of the value has multiple lines, enter all the text without line break. \n",
    "        Include all value text even if it extends beyond the highlight box.</p>\n",
    "        <p><img src=\"https://assets.crowd.aws/images/a2i-console/multiple-lines.png\"></p>\n",
    "    </short-instructions>\n",
    "    <full-instructions header=\"Instructions\"></full-instructions>\n",
    "  </crowd-textract-analyze-document>\n",
    "</crowd-form>\n",
    "\"\"\"\n",
    "\n",
    "def create_task_ui(task_ui_name):\n",
    "    '''\n",
    "    Creates a Human Task UI resource.\n",
    "\n",
    "    Returns:\n",
    "    struct: HumanTaskUiArn\n",
    "    '''\n",
    "    response = sagemaker.create_human_task_ui(\n",
    "        HumanTaskUiName=task_ui_name,\n",
    "        UiTemplate={'Content': template})\n",
    "    return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Task UI name - this value is unique per account and region. You can also provide your own value here.\n",
    "taskUIName = 'ui-textract-dem0'\n",
    "\n",
    "# Create task UI\n",
    "humanTaskUiResponse = create_task_ui(taskUIName)\n",
    "humanTaskUiArn = humanTaskUiResponse['HumanTaskUiArn']\n",
    "print(humanTaskUiArn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating the Flow Definition"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we're going to create a flow definition definition. Flow Definitions allow us to specify:\n",
    "\n",
    "* The conditions under which your human loop will be called.\n",
    "* The workforce that your tasks will be sent to.\n",
    "* The instructions that your workforce will receive. This is called a worker task template.\n",
    "* The configuration of your worker tasks, including the number of workers that receive a task and time limits to complete tasks.\n",
    "* Where your output data will be stored.\n",
    "\n",
    "This demo is going to use the API, but you can optionally create this workflow definition in the console as well. \n",
    "\n",
    "For more details and instructions, see: https://docs.aws.amazon.com/sagemaker/latest/dg/a2i-create-flow-definition.html."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Specify Human Loop Activation Conditions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we are using a built-in integration type for A2I, we can use Human Loop Activation Conditions to provide conditions that trigger a human loop.\n",
    "\n",
    "Here we are specifying conditions for specific keys in our document. If Textract's confidence falls outside of the thresholds set here, the document will be sent to a human for review, with the specific keys that triggered the human loop prompted to the worker. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_flow_definition(flow_definition_name):\n",
    "    '''\n",
    "    Creates a Flow Definition resource\n",
    "\n",
    "    Returns:\n",
    "    struct: FlowDefinitionArn\n",
    "    '''\n",
    "    # Visit https://docs.aws.amazon.com/sagemaker/latest/dg/a2i-human-fallback-conditions-json-schema.html for more information on this schema.\n",
    "    humanLoopActivationConditions = json.dumps(\n",
    "        {\n",
    "            \"Conditions\": [\n",
    "                {\n",
    "                  \"Or\": [\n",
    "                    \n",
    "                    {\n",
    "                        \"ConditionType\": \"ImportantFormKeyConfidenceCheck\",\n",
    "                        \"ConditionParameters\": {\n",
    "                            \"ImportantFormKey\": \"Mail address\",\n",
    "                            \"ImportantFormKeyAliases\": [\"Mail Address:\",\"Mail address:\", \"Mailing Add:\",\"Mailing Addresses\"],\n",
    "                            \"KeyValueBlockConfidenceLessThan\": 100,\n",
    "                            \"WordBlockConfidenceLessThan\": 100\n",
    "                        }\n",
    "                    },\n",
    "                    {\n",
    "                        \"ConditionType\": \"MissingImportantFormKey\",\n",
    "                        \"ConditionParameters\": {\n",
    "                            \"ImportantFormKey\": \"Mail address\",\n",
    "                            \"ImportantFormKeyAliases\": [\"Mail Address:\",\"Mail address:\",\"Mailing Add:\",\"Mailing Addresses\"]\n",
    "                        }\n",
    "                    },\n",
    "                    {\n",
    "                        \"ConditionType\": \"ImportantFormKeyConfidenceCheck\",\n",
    "                        \"ConditionParameters\": {\n",
    "                            \"ImportantFormKey\": \"Phone Number\",\n",
    "                            \"ImportantFormKeyAliases\": [\"Phone number:\", \"Phone No.:\", \"Number:\"],\n",
    "                            \"KeyValueBlockConfidenceLessThan\": 100,\n",
    "                            \"WordBlockConfidenceLessThan\": 100\n",
    "                        }\n",
    "                    },\n",
    "                    {\n",
    "                      \"ConditionType\": \"ImportantFormKeyConfidenceCheck\",\n",
    "                      \"ConditionParameters\": {\n",
    "                        \"ImportantFormKey\": \"*\",\n",
    "                        \"KeyValueBlockConfidenceLessThan\": 100,\n",
    "                        \"WordBlockConfidenceLessThan\": 100\n",
    "                      }\n",
    "                    },\n",
    "                    {\n",
    "                      \"ConditionType\": \"ImportantFormKeyConfidenceCheck\",\n",
    "                      \"ConditionParameters\": {\n",
    "                        \"ImportantFormKey\": \"*\",\n",
    "                        \"KeyValueBlockConfidenceGreaterThan\": 0,\n",
    "                        \"WordBlockConfidenceGreaterThan\": 0\n",
    "                      }\n",
    "                    }\n",
    "            ]\n",
    "        }\n",
    "            ]\n",
    "        }\n",
    "    )\n",
    "\n",
    "    response = sagemaker.create_flow_definition(\n",
    "            FlowDefinitionName= flow_definition_name,\n",
    "            RoleArn= ROLE,\n",
    "            HumanLoopConfig= {\n",
    "                \"WorkteamArn\": WORKTEAM_ARN,\n",
    "                \"HumanTaskUiArn\": humanTaskUiArn,\n",
    "                \"TaskCount\": 1,\n",
    "                \"TaskDescription\": \"Document analysis sample task description\",\n",
    "                \"TaskTitle\": \"Document analysis sample task\"\n",
    "            },\n",
    "            HumanLoopRequestSource={\n",
    "                \"AwsManagedHumanLoopRequestSource\": \"AWS/Textract/AnalyzeDocument/Forms/V1\"\n",
    "            },\n",
    "            HumanLoopActivationConfig={\n",
    "                \"HumanLoopActivationConditionsConfig\": {\n",
    "                    \"HumanLoopActivationConditions\": humanLoopActivationConditions\n",
    "                }\n",
    "            },\n",
    "            OutputConfig={\n",
    "                \"S3OutputPath\" : OUTPUT_PATH\n",
    "            }\n",
    "        )\n",
    "    \n",
    "    return response['FlowDefinitionArn']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Now we are ready to create our Flow Definition!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Flow definition name - this value is unique per account and region. You can also provide your own value here.\n",
    "uniqueId = str(uuid.uuid4())\n",
    "flowDefinitionName = f'fd-textract-{uniqueId}' \n",
    "\n",
    "flowDefinitionArn = create_flow_definition(flowDefinitionName)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def describe_flow_definition(name):\n",
    "    '''\n",
    "    Describes Flow Definition\n",
    "\n",
    "    Returns:\n",
    "    struct: response from DescribeFlowDefinition API invocation\n",
    "    '''\n",
    "    return sagemaker.describe_flow_definition(\n",
    "        FlowDefinitionName=name)\n",
    "\n",
    "# Describe flow definition - status should be active\n",
    "for x in range(60):\n",
    "    describeFlowDefinitionResponse = describe_flow_definition(flowDefinitionName)\n",
    "    print(describeFlowDefinitionResponse['FlowDefinitionStatus'])\n",
    "    if (describeFlowDefinitionResponse['FlowDefinitionStatus'] == 'Active'):\n",
    "        print(\"Flow Definition is active\")\n",
    "        break\n",
    "    time.sleep(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Analyze Document with Textract"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have setup our Flow Definition, all that's left is calling Textract's Analyze Document API, and including our A2I paramters in the HumanLoopConfig."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uniqueId = str(uuid.uuid4())\n",
    "human_loop_unique_id = uniqueId + '1'\n",
    "\n",
    "humanLoopConfig = {\n",
    "    'FlowDefinitionArn':flowDefinitionArn,\n",
    "    'HumanLoopName':human_loop_unique_id, \n",
    "    'DataAttributes': { 'ContentClassifiers': [ 'FreeOfPersonallyIdentifiableInformation' ]}\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def analyze_document_with_a2i(document_name, bucket):\n",
    "    response = textract.analyze_document(\n",
    "        Document={'S3Object': {'Bucket': bucket, 'Name': document_name}},\n",
    "        FeatureTypes=[\"TABLES\", \"FORMS\"], \n",
    "        HumanLoopConfig=humanLoopConfig\n",
    "    )\n",
    "    return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "analyzeDocumentResponse = analyze_document_with_a2i(document, BUCKET)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Human Loops\n",
    "When a document passed to Textract matches the conditions in FlowDefinition, a HumanLoopArn will be present in the response to analyze_document. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If a _HumanLoopArn_ is present in the _HumanLoopActivationOutput_, we know **a Human Loop has been started**!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if 'HumanLoopArn' in analyzeDocumentResponse['HumanLoopActivationOutput']:\n",
    "    # A human loop has been started!\n",
    "    print(f'A human loop has been started with ARN: {analyzeDocumentResponse[\"HumanLoopActivationOutput\"][\"HumanLoopArn\"]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check Status of Human Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_human_loops_in_workflow = a2i_runtime_client.list_human_loops(FlowDefinitionArn=flowDefinitionArn)['HumanLoopSummaries']\n",
    "\n",
    "for human_loop in all_human_loops_in_workflow:\n",
    "    print(f'\\nHuman Loop Name: {human_loop[\"HumanLoopName\"]}')\n",
    "    print(f'Human Loop Status: {human_loop[\"HumanLoopStatus\"]} \\n')\n",
    "    print('\\n')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Wait For Workers to Complete Task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "workteamName = WORKTEAM_ARN[WORKTEAM_ARN.rfind('/') + 1:]\n",
    "print(\"Navigate to the private worker portal and do the tasks. Make sure you've invited yourself to your workteam!\")\n",
    "print('https://' + sagemaker.describe_workteam(WorkteamName=workteamName)['Workteam']['SubDomain'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check Status of Human Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_human_loops_in_workflow = a2i_runtime_client.list_human_loops(FlowDefinitionArn=flowDefinitionArn)['HumanLoopSummaries']\n",
    "\n",
    "completed_loops = []\n",
    "for human_loop in all_human_loops_in_workflow:\n",
    "    print(f'\\nHuman Loop Name: {human_loop[\"HumanLoopName\"]}')\n",
    "    print(f'Human Loop Status: {human_loop[\"HumanLoopStatus\"]} \\n')\n",
    "    print('\\n')\n",
    "    if human_loop['HumanLoopStatus'] == 'Completed':\n",
    "        completed_loops.append(human_loop['HumanLoopName'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### View Task Results  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once work is completed, Amazon A2I stores results in your S3 bucket and sends a Cloudwatch event. Your results should be available in the S3 OUTPUT_PATH when all work is completed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import pprint\n",
    "pp = pprint.PrettyPrinter(indent=2)\n",
    "\n",
    "def retrieve_a2i_results_from_output_s3_uri(bucket, a2i_s3_output_uri):\n",
    "    '''\n",
    "    Gets the json file published by A2I and returns a deserialized object\n",
    "    '''\n",
    "    splitted_string = re.split('s3://' +  bucket + '/', a2i_s3_output_uri)\n",
    "    output_bucket_key = splitted_string[1]\n",
    "\n",
    "    response = s3.get_object(Bucket=bucket, Key=output_bucket_key)\n",
    "    content = response[\"Body\"].read()\n",
    "    return json.loads(content)\n",
    "    \n",
    "\n",
    "for human_loop_name in completed_loops:\n",
    "\n",
    "    describe_human_loop_response = a2i_runtime_client.describe_human_loop(\n",
    "        HumanLoopName=human_loop_name\n",
    "    )\n",
    "    \n",
    "    print(f'\\nHuman Loop Name: {describe_human_loop_response[\"HumanLoopName\"]}')\n",
    "    print(f'Human Loop Status: {describe_human_loop_response[\"HumanLoopStatus\"]}')\n",
    "    print(f'Human Loop Output Location: : {describe_human_loop_response[\"HumanLoopOutput\"][\"OutputS3Uri\"]} \\n')\n",
    "    \n",
    "    # Uncomment below line to print out a2i human answers\n",
    "    output = retrieve_a2i_results_from_output_s3_uri(BUCKET, describe_human_loop_response['HumanLoopOutput']['OutputS3Uri'])\n",
    "#     pp.pprint(output)\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The End!"
   ]
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (Data Science)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-1:081325390199:image/datascience-1.0"
  },
  "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}