Accept image with blob format

app.py

@@ -1,3 +1,5 @@
+import base64
+import io
 import sys
 import torch
 import onnx
@@ -8,12 +10,7 @@ from PIL import Image
 
 
 class DocumentParserModel:
-    def __init__(
-        self,
-        model_path,
-        img_size,
-        charset
-    ):
+    def __init__(self, model_path, img_size, charset):
         self.charset = charset
         self.tokenizer_base = Tokenizer(self.charset)
         self.transform = self.create_transform_pipeline(img_size)
@@ -32,10 +29,15 @@ class DocumentParserModel:
         onnx.checker.check_model(onnx_model)
         return rt.InferenceSession(model_path)
 
-    # TODO: test with image blob
+    def load_image_from_base64(self, base64_string):
+        img_data = base64.b64decode(base64_string)
+        image_buffer = io.BytesIO(img_data)
+        image = Image.open(image_buffer)
+        return image
+
     def predict_text(self, image_path):
         try:
-            with Image.open(image_path) as img_org:
+            with self.load_image_from_base64(image_path) as img_org:
                 x = self.transform(img_org.convert("RGB")).unsqueeze(0)
                 ort_inputs = {self.ort_session.get_inputs()[0].name: x.cpu().numpy()}
                 logits = self.ort_session.run(None, ort_inputs)[0]