Update api to support base64 files (#246)

* support base64 input files for the ocr api

* add tests

* reformat

Author: zdeveloper

OCR/ocr/api.py

@@ -1,3 +1,5 @@
+import base64
+
 import uvicorn
 import json
 import cv2 as cv
@@ -33,8 +35,8 @@ async def health_check():
     return {"status": "UP"}
 
 
-@app.post("/image_to_text/")
-async def image_to_text(source_image: UploadFile, segmentation_template: UploadFile, labels: str = Form()):
+@app.post("/image_file_to_text/")
+async def image_file_to_text(source_image: UploadFile, segmentation_template: UploadFile, labels: str = Form()):
     source_image_np = np.frombuffer(await source_image.read(), np.uint8)
     source_image_img = cv.imdecode(source_image_np, cv.IMREAD_COLOR)
 
@@ -48,6 +50,23 @@ async def image_to_text(source_image: UploadFile, segmentation_template: UploadF
     return results
 
 
+@app.post("/image_to_text/")
+async def image_to_text(source_image: str = Form(), segmentation_template: str = Form(), labels: str = Form()):
+    source_image_stripped = source_image.replace("data:image/png;base64,", "", 1)
+    source_image_np = np.frombuffer(base64.b64decode(source_image_stripped), np.uint8)
+    source_image_img = cv.imdecode(source_image_np, cv.IMREAD_COLOR)
+
+    segmentation_template_stripped = segmentation_template.replace("data:image/png;base64,", "", 1)
+    segmentation_template_np = np.frombuffer(base64.b64decode(segmentation_template_stripped), np.uint8)
+    segmentation_template_img = cv.imdecode(segmentation_template_np, cv.IMREAD_COLOR)
+
+    loaded_json = json.loads(labels)
+    segments = segmenter.segment(source_image_img, segmentation_template_img, loaded_json)
+    results = ocr.image_to_text(segments)
+
+    return results
+
+
 def start():
     """Launched with `poetry run start` at root level"""
     uvicorn.run(app, host="0.0.0.0", port=8000, reload=False)

OCR/tests/api_test.py

@@ -1,3 +1,4 @@
+import base64
 import os
 import json
 
@@ -20,7 +21,7 @@ def test_health_check(self):
         assert response.status_code == 200
         assert response.json() == {"status": "UP"}
 
-    def test_ocr(self):
+    def test_image_file_to_text(self):
         # load the files
         with (
             open(segmentation_template_path, "rb") as segmentation_template_file,
@@ -35,7 +36,67 @@ def test_ocr(self):
             ]
 
             # call ocr api
-            response = client.post(url="/image_to_text", files=files_to_send, data={"labels": json.dumps(label_data)})
+            response = client.post(
+                url="/image_file_to_text", files=files_to_send, data={"labels": json.dumps(label_data)}
+            )
+
+            assert response.status_code == 200
+
+            # assert output
+            response_json = response.json()
+            assert response_json["nbs_patient_id"][0] == "SIENNA HAMPTON"
+            assert response_json["nbs_cas_id"][0] == "123555"
+
+    def test_image_to_text(self):
+        # load the files
+        with (
+            open(segmentation_template_path, "rb") as segmentation_template_file,
+            open(source_image_path, "rb") as source_image_file,
+            open(labels_path, "rb") as labels,
+        ):
+            label_data = json.load(labels)
+
+            source_image_base64 = base64.b64encode(source_image_file.read()).decode("ascii")
+            segmentation_template_base64 = base64.b64encode(segmentation_template_file.read()).decode("ascii")
+
+            response = client.post(
+                url="/image_to_text",
+                data={
+                    "labels": json.dumps(label_data),
+                    "source_image": str(source_image_base64),
+                    "segmentation_template": str(segmentation_template_base64),
+                },
+                headers={"Content-Type": "application/x-www-form-urlencoded"},
+            )
+
+            assert response.status_code == 200
+
+            # assert output
+            response_json = response.json()
+            assert response_json["nbs_patient_id"][0] == "SIENNA HAMPTON"
+            assert response_json["nbs_cas_id"][0] == "123555"
+
+    def test_image_to_text_with_padding(self):
+        # load the files
+        with (
+            open(segmentation_template_path, "rb") as segmentation_template_file,
+            open(source_image_path, "rb") as source_image_file,
+            open(labels_path, "rb") as labels,
+        ):
+            label_data = json.load(labels)
+
+            source_image_base64 = base64.b64encode(source_image_file.read()).decode("ascii")
+            segmentation_template_base64 = base64.b64encode(segmentation_template_file.read()).decode("ascii")
+
+            response = client.post(
+                url="/image_to_text",
+                data={
+                    "labels": json.dumps(label_data),
+                    "source_image": f"data:image/png;base64,{str(source_image_base64)}",
+                    "segmentation_template": f"data:image/png;base64,{str(segmentation_template_base64)}",
+                },
+                headers={"Content-Type": "application/x-www-form-urlencoded"},
+            )
 
             assert response.status_code == 200