Update incorrectly uploaded script

Author: leschan

lambda/a2i_create_humanloop.py

@@ -1,72 +1,62 @@
 import json
 import os
 import boto3
-ssm = boto3.client('ssm')
+import uuid
 sns = boto3.client('sns')
-rekognition = boto3.client('rekognition')
+a2i = boto3.client('sagemaker-a2i-runtime')
+ssm = boto3.client('ssm')
 dynamodb = boto3.client('dynamodb')
-def get_parameters():
-    response = ssm.get_parameters_by_path(
-        Path=os.environ['parameter_store_path'],
-        Recursive=True
-    )
-    parameter_store = {}
-    for parameter in response['Parameters']:
-        parameter_name = (parameter['Name'].split('/'))[-1]
-        parameter_store[parameter_name] = parameter['Value']
-    return parameter_store
-def store_detection_results(dynamodb_table, detectlabel_request_id, detectlabel_date, s3_bucket_name, s3_object_key, s3_object_eTag, project_version, detected_label, confidence_level, minimum_confidence_level, A2I):
-    response = dynamodb.put_item(
-        TableName=dynamodb_table,
-        Item={
-            'DetectLabelRequestId': {'S':detectlabel_request_id},
-            'DetectLabelDate': {'S':detectlabel_date},
-            'S3Bucket': {'S':s3_bucket_name},
-            'S3ObjectKey': {'S':s3_object_key},
-            'S3ObjectEtag': {'S':s3_object_eTag},
-            'ProjectVersionArn': {'S':project_version},
-            'DectectedLabel': {'S':detected_label},
-            'DetectedConfidenceLevel': {'N':str(confidence_level)},
-            'MinimumConfidenceLevel': {'N':str(minimum_confidence_level)},
-            'A2IEnabled': {'BOOL':A2I}
-        }
-    )
+def get_parameter(parameter_name):
+    response = ssm.get_parameter(Name=parameter_name)
+    parameter_value = json.loads(response['Parameter']['Value'])
+    return parameter_value
 def publish_message(sns_subject, sns_message, topic_arn):
     response = sns.publish(
         TopicArn=topic_arn,
         Message=sns_message,
         Subject=sns_subject
     )
+def append_a2i_request(dynamodb_table, detectlabel_request_id, human_loop_name, humanloop_request_id):
+    response = dynamodb.update_item(
+        TableName=dynamodb_table,
+        Key={
+            'DetectLabelRequestId': {'S':detectlabel_request_id}
+        },
+        ExpressionAttributeNames={
+            '#HLN':'HumanLoopName',
+            '#HLRI':'HumanLoopRequestId'
+        },
+        ExpressionAttributeValues={
+            ':n': {'S':human_loop_name},
+            ':i': {'S':humanloop_request_id}
+        },
+        UpdateExpression='SET #HLN=:n, #HLRI=:i'
+    )
 def handler(event, context):
-    parameter_store = get_parameters()
-    a2i = True if parameter_store['Enable-A2I-Workflow'].lower().capitalize() == 'True' else False
-    sys_vars = json.loads(parameter_store['For-System-Use-Only'])
-    project_version = sys_vars['rekognition_project_version_arn']
+    parameter_name = os.environ['parameter_store_path'] + 'For-System-Use-Only'
+    sys_vars = get_parameter(parameter_name)
     dynamodb_table = sys_vars['dynamodb_table']
-    minimum_confidence_level=parameter_store['Minimum-Label-Detection-Confidence']
-    s3_bucket_name=event['s3event']['s3']['bucket']['name']
-    s3_object_key=event['s3event']['s3']['object']['key']
-    s3_object_eTag=event['s3event']['s3']['object']['eTag']
-    s3_event_time=event['s3event']['eventTime']
-    response = rekognition.detect_custom_labels(
-        ProjectVersionArn=project_version,
-        Image={
-            'S3Object': {
-                'Bucket': s3_bucket_name,
-                'Name': s3_object_key
-            }
-        },
-        MaxResults=1,
-        MinConfidence=0
+    detectlabel_request_id=event['message']['ResponseMetadata']['RequestId']
+    s3_bucket_name=event['message']['s3event']['s3']['bucket']['name']
+    s3_object_key=event['message']['s3event']['s3']['object']['key']
+    detected_label=event['message']['CustomLabels'][0]['Name']
+    confidence_level=event['message']['CustomLabels'][0]['Confidence']
+    human_loop_name=str(uuid.uuid4())
+    response = a2i.start_human_loop(
+        HumanLoopName = human_loop_name,
+        FlowDefinitionArn = sys_vars['flow_definition_arn'],
+        HumanLoopInput = {
+            'InputContent': json.dumps({
+                'initialValue': confidence_level,
+                'detectLabelRequestId': detectlabel_request_id,
+                'taskObject': 's3://'+s3_bucket_name+'/'+s3_object_key
+            })
+        }
     )
-    confidence_level = response['CustomLabels'][0]['Confidence']
-    detected_label = response['CustomLabels'][0]['Name']
-    detectlabel_request_id = response['ResponseMetadata']['RequestId']
-    if confidence_level < float(parameter_store['Minimum-Label-Detection-Confidence']):
-        response['A2I'] = True if a2i else False
-    store_detection_results(dynamodb_table, detectlabel_request_id, s3_event_time, s3_bucket_name, s3_object_key, s3_object_eTag, project_version, detected_label, confidence_level, minimum_confidence_level, a2i)
-    publish_message('Rekgnition Custom Labels Detection Invoked', json.dumps(response), sys_vars['sns-topic'])
-    response['s3event'] = event['s3event']
+    humanloop_request_id=response['ResponseMetadata']['RequestId']
+    append_a2i_request(dynamodb_table, detectlabel_request_id, human_loop_name, humanloop_request_id)
+    response['s3event'] = event['message']['s3event']
+    publish_message('A2I Human Loop Initiated', json.dumps(response), sys_vars['sns-topic'])
     return {
         'message': response
     }