From 093e6a81fc744f60b07594f52c512d9f8189dfe3 Mon Sep 17 00:00:00 2001
From: Maxence Guindon <maxence.guindon@inspection.gc.ca>
Date: Wed, 14 Feb 2024 19:59:33 +0000
Subject: [PATCH] fixes #51: Standardize headers with azureml-model-deployment
 being mandantory

---
 app.py                                   | 15 ++++++++++-----
 model_utilitary_functions/model_UTILS.py | 20 +++++---------------
 2 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/app.py b/app.py
index 113ebb04..f6fa0672 100644
--- a/app.py
+++ b/app.py
@@ -35,7 +35,11 @@
 NACHET_MODEL = os.getenv("NACHET_MODEL")
 
 # The following tuples will be used to store the endpoints and their respective utilitary functions
-tuple_endpoints = (((endpoint_url, endpoint_api_key, ""),),((sd_endpoint, sd_api_key, utils.image_slicing),(swin_endpoint, swin_api_key, utils.swin_result_parser)))
+tuple_endpoints = (
+    ((endpoint_url, endpoint_api_key, "m-14of15seeds-6seedsmag", None),)
+    ,((sd_endpoint, sd_api_key, "seed-detector-1", utils.image_slicing),
+      (swin_endpoint, swin_api_key, "swinv1-base-dataaugv2-1", utils.swin_result_parser))
+)
 
 CACHE = {
     "seeds": None,
@@ -156,7 +160,8 @@ async def inference_request():
     """
     Performs inference on an image, and returns the results.
     The image and inference results are uploaded to a folder in the user's container.
-    """   
+    """  
+     
     seconds = time.perf_counter() # transform into logging
     try:
         print("Entering inference request") # Transform into logging
@@ -191,18 +196,18 @@ async def inference_request():
             cache_json_result = None
             for model in pipelines_endpoints.get(pipeline_name):
                 
-                endpoint_url, endpoint_api_key, utilitary_function = model
+                endpoint_url, endpoint_api_key, model_name, utilitary_function = model
 
                 if isinstance(image_bytes, list):
                     result_json = []
                     for img in image_bytes:
-                        req = await utils.request_factory(img, endpoint_url, endpoint_api_key)
+                        req = await utils.request_factory(img, endpoint_url, endpoint_api_key, model_name)
                         response = urllib.request.urlopen(req)
                         result = response.read()
                         result_json.append(json.loads(result.decode("utf-8")))
 
                 elif isinstance(image_bytes, str):
-                    req = await utils.request_factory(image_bytes, endpoint_url, endpoint_api_key)
+                    req = await utils.request_factory(image_bytes, endpoint_url, endpoint_api_key, model_name)
                     response = urllib.request.urlopen(req)
                     result = response.read()
                     result_json = json.loads(result.decode("utf-8"))
diff --git a/model_utilitary_functions/model_UTILS.py b/model_utilitary_functions/model_UTILS.py
index 420f8b4c..f3ee15da 100644
--- a/model_utilitary_functions/model_UTILS.py
+++ b/model_utilitary_functions/model_UTILS.py
@@ -58,21 +58,14 @@ async def swin_result_parser(img_box:dict, results: dict) -> list:
         list: The updated image box with modified labels and scores.
     """
     for i, result in enumerate(results):
-        img_box[0]['boxes'][i]['label'] = result[0].get('label')
-        img_box[0]['boxes'][i]['score'] = result[0].get('score')
+        img_box[0]['boxes'][i]['label'] = [d.get("label") for d in result]
+        img_box[0]['boxes'][i]['score'] = [d.get("score") for d in result]
     
     return img_box
 
-async def seed_detector_header(api_key: str) -> dict:
-    return {
-        "Content-Type": "application/json",
-        "Authorization": ("Bearer " + api_key),
-        "azureml-model-deployment": "seed-detector-1",
-    }
-
 # Eventually the goals would be to have a request factory that would return
 # a request for the specified models such as the following:
-async def request_factory(img_bytes: str | bytes, endpoint_url: str, api_key: str) -> Request:
+async def request_factory(img_bytes: str | bytes, endpoint_url: str, api_key: str, model_name: str) -> Request:
     """
     Args:
         img_bytes (str | bytes): The image data as either a string or bytes.
@@ -82,14 +75,11 @@ async def request_factory(img_bytes: str | bytes, endpoint_url: str, api_key: st
     Returns:
         Request: The request object for calling the AI model.
     """
-
-    model_name = endpoint_url.split("/")[2].split(".")[0]
-
     headers = {
         "Content-Type": "application/json",
         "Authorization": ("Bearer " + api_key),
-    } if model_name != "seed-detector" else await seed_detector_header(api_key)
-
+        "azureml-model-deployment": model_name,
+    }
 
     if isinstance(img_bytes, str): 
         data = {