code-update-to-handle-queries-feature

Author: srcecde

textract/async-form-table/lambda-create-job/lambda_function.py

@@ -13,6 +13,7 @@
 import boto3
 import logging
 import traceback
+from ast import literal_eval
 from urllib.parse import unquote_plus
 
 logger = logging.getLogger()
@@ -37,17 +38,18 @@ def process_error():
     OUTPUT_S3_PREFIX = os.environ["OUTPUT_S3_PREFIX"]
     SNS_TOPIC_ARN = os.environ["SNS_TOPIC_ARN"]
     SNS_ROLE_ARN = os.environ["SNS_ROLE_ARN"]
+    FEATURES = literal_eval(os.environ["FEATURES"])
     logger.info(
-        f"OUTPUT_BUCKET_NAME: {OUTPUT_BUCKET_NAME}, OUTPUT_S3_PREFIX: {OUTPUT_S3_PREFIX}, SNS_TOPIC_ARN: {SNS_TOPIC_ARN}, SNS_ROLE_ARN: {SNS_ROLE_ARN}"
+        f"OUTPUT_BUCKET_NAME: {OUTPUT_BUCKET_NAME}, OUTPUT_S3_PREFIX: {OUTPUT_S3_PREFIX}, SNS_TOPIC_ARN: {SNS_TOPIC_ARN}, SNS_ROLE_ARN: {SNS_ROLE_ARN}, FEATURES: {FEATURES}"
     )
 except Exception as e:
     error_msg = process_error()
     logger.error(error_msg)
 
 
 def lambda_handler(event, context):
-
     textract = boto3.client("textract")
+    s3 = boto3.client("s3")
     try:
         if "Records" in event:
             logger.info(f"Event: {event}")
@@ -56,18 +58,44 @@ def lambda_handler(event, context):
             filename = unquote_plus(str(file_obj["s3"]["object"]["key"]))
             logger.info(f"Bucket: {bucketname} ::: Key: {filename}")
 
-            response = textract.start_document_analysis(
-                DocumentLocation={"S3Object": {"Bucket": bucketname, "Name": filename}},
-                FeatureTypes=["TABLES", "FORMS", "SIGNATURES"],
-                OutputConfig={
-                    "S3Bucket": OUTPUT_BUCKET_NAME,
-                    "S3Prefix": OUTPUT_S3_PREFIX,
-                },
-                NotificationChannel={
-                    "SNSTopicArn": SNS_TOPIC_ARN,
-                    "RoleArn": SNS_ROLE_ARN,
-                },
-            )
+            if "QUERIES" in FEATURES:
+                file_obj = s3.get_object(
+                    Bucket=bucketname, Key="async-input/queries.json"
+                )
+                queries_data = json.loads(file_obj["Body"].read().decode("utf-8"))
+
+                response = textract.start_document_analysis(
+                    DocumentLocation={
+                        "S3Object": {"Bucket": bucketname, "Name": filename}
+                    },
+                    FeatureTypes=FEATURES,
+                    QueriesConfig=queries_data,
+                    OutputConfig={
+                        "S3Bucket": OUTPUT_BUCKET_NAME,
+                        "S3Prefix": OUTPUT_S3_PREFIX,
+                    },
+                    NotificationChannel={
+                        "SNSTopicArn": SNS_TOPIC_ARN,
+                        "RoleArn": SNS_ROLE_ARN,
+                    },
+                )
+
+            else:
+                response = textract.start_document_analysis(
+                    DocumentLocation={
+                        "S3Object": {"Bucket": bucketname, "Name": filename}
+                    },
+                    FeatureTypes=FEATURES,
+                    OutputConfig={
+                        "S3Bucket": OUTPUT_BUCKET_NAME,
+                        "S3Prefix": OUTPUT_S3_PREFIX,
+                    },
+                    NotificationChannel={
+                        "SNSTopicArn": SNS_TOPIC_ARN,
+                        "RoleArn": SNS_ROLE_ARN,
+                    },
+                )
+
             if response["ResponseMetadata"]["HTTPStatusCode"] == 200:
                 logger.info(f"Job created successfully")
                 return {

textract/async-form-table/lambda-process-response/helper/helper.py

@@ -101,6 +101,15 @@ def save_text_csv(keys, values, job_id, BUCKET_NAME):
     upload_to_s3(csv_buffer, BUCKET_NAME, key)
 
 
+def save_queries_csv(queries, job_id, BUCKET_NAME):
+    key = f"queries/{job_id}/queryAnswer.csv"
+    df = pd.DataFrame.from_dict(queries, orient="index").reset_index()
+    df.drop(["index", "answer_ids"], axis=1, inplace=True)
+    csv_buffer = io.StringIO()
+    df.to_csv(csv_buffer)
+    upload_to_s3(csv_buffer, BUCKET_NAME, key)
+
+
 def map_word_id(response):
     word_map = {}
     for block in response["Blocks"]:
@@ -112,7 +121,13 @@ def map_word_id(response):
 
 
 def process_response(
-    BUCKET_NAME, job_id, get_table=True, get_kv=True, get_text=True, get_signatures=True
+    BUCKET_NAME,
+    job_id,
+    get_table=True,
+    get_kv=True,
+    get_text=True,
+    get_signatures=True,
+    get_queries=True,
 ):
     textract = boto3.client("textract")
 
@@ -149,8 +164,9 @@ def process_response(
         get_kv=get_kv,
         get_text=get_text,
         get_signatures=get_signatures,
+        get_queries=get_queries,
     )
-    table, final_map, text, sign = parse.process_response()
+    table, final_map, text, sign, queries = parse.process_response()
 
     if get_kv:
         keys = list(map(itemgetter(0), final_map))
@@ -165,4 +181,6 @@ def process_response(
         save_text_csv(text_key, text_value, job_id, BUCKET_NAME)
     if get_signatures:
         save_sign_csv(sign, job_id, BUCKET_NAME)
+    if get_queries:
+        save_queries_csv(queries, job_id, BUCKET_NAME)
     logger.info("Parsing completed")

textract/async-form-table/lambda-process-response/helper/parser.py

@@ -16,7 +16,7 @@
 
 
 class Parse:
-    def __init__(self, page, get_table, get_kv, get_text, get_signatures):
+    def __init__(self, page, get_table, get_kv, get_text, get_signatures, get_queries):
         self.response = page
         self.word_map = {}
         self.table_page_map = {}
@@ -28,6 +28,7 @@ def __init__(self, page, get_table, get_kv, get_text, get_signatures):
         self.get_kv = get_kv
         self.get_text = get_text
         self.get_signatures = get_signatures
+        self.get_queries = get_queries
 
     def extract_text(self, extract_by="LINE"):
         for block in self.response:
@@ -57,7 +58,6 @@ def extract_table_info(self):
         response_block_len = len(self.response) - 1
 
         for n, block in enumerate(self.response):
-
             if block["BlockType"] == "TABLE":
                 key = f"table_{uuid.uuid4().hex}_page_{block['Page']}"
                 temp_table = []
@@ -94,7 +94,6 @@ def extract_table_info(self):
 
     def get_key_map(self):
         for block in self.response:
-
             if block["BlockType"] == "KEY_VALUE_SET" and "KEY" in block["EntityTypes"]:
                 for relation in block["Relationships"]:
                     if relation["Type"] == "VALUE":
@@ -136,8 +135,66 @@ def get_signature_info(self):
                 temp_counter += 1
         return (page, signature, confidence)
 
+    def get_queries_info(self):
+        temp_id = []
+        f_response = {}
+
+        for e, block in enumerate(self.response):
+            if block["BlockType"] == "QUERY":
+                if "Relationships" not in block:
+                    rp = {
+                        "query": block.get("Query").get("Text"),
+                        "alias": block.get("Query").get("Alias"),
+                        "answer_ids": None,
+                        "answer": None,
+                        "confidence": None,
+                        "page": None,
+                    }
+                else:
+                    child_ids = [
+                        ids
+                        for rel in block.get("Relationships")
+                        for ids in rel["Ids"]
+                        if rel.get("Type") == "ANSWER"
+                    ]
+                    rp = {
+                        "query": block.get("Query").get("Text"),
+                        "alias": block.get("Query").get("Alias"),
+                        "answer_ids": child_ids,
+                        "answer": None,
+                        "confidence": None,
+                        "page": None,
+                    }
+
+                f_response[block.get("Id")] = rp
+                temp_id.append({block.get("Id"): rp})
+
+            if block["BlockType"] == "QUERY_RESULT":
+                q_id = list(temp_id[-1].keys())[0]
+                q_val = temp_id[-1].get(q_id)
+
+                if q_val.get("answer_ids"):
+                    if block.get("Id") in q_val.get("answer_ids"):
+                        q_ans = block.get("Text")
+                        confidence_s = block.get("Confidence")
+                        q_val["confidence"] = confidence_s
+                        if q_val.get("answer"):
+                            q_val["answer"] = f"{q_val.get('answer')} {q_ans}"
+                        else:
+                            q_val["answer"] = q_ans
+                        q_val["page"] = block.get("Page")
+                f_response[q_id] = q_val
+                temp_id = []
+        return f_response
+
     def process_response(self):
-        final_map, table_info, text = None, None, None
+        final_map, table_info, text, sign_info, queries_info = (
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
 
         logging.info("Mapping Id with word")
         self.map_word_id()
@@ -160,4 +217,8 @@ def process_response(self):
             logging.info("Extracting signature information")
             sign_info = self.get_signature_info()
 
-        return table_info, final_map, text, sign_info
+        if self.get_queries:
+            logging.info("Extracting queries information")
+            queries_info = self.get_queries_info()
+
+        return table_info, final_map, text, sign_info, queries_info

textract/async-form-table/lambda-process-response/lambda_function.py

@@ -17,7 +17,6 @@
 
 
 def lambda_handler(event, context):
-
     try:
         BUCKET_NAME = os.environ["BUCKET_NAME"]
         logger.info(f"Destination bucket: {BUCKET_NAME}")
@@ -39,6 +38,7 @@ def lambda_handler(event, context):
                 get_kv=True,
                 get_text=True,
                 get_signatures=True,
+                get_queries=True,
             )
 
             return {