Update handler.py

handler.py

@@ -1,31 +1,35 @@
-import os
 import easyocr
+import torch
+from typing import Dict
+from PIL import Image
+import io
+import json
 
 class EndpointHandler:
-    def __init__(self, model_dir):
-        """
-        Initialize the handler and load the model.
-        """
-        self.model_dir = model_dir
-        # Load the EasyOCR reader
-        self.reader = easyocr.Reader(['en'], gpu=False)  # Use GPU=True if GPU is available
+    def __init__(self, model_dir: str):
+        # Path to your model file
+        model_path = f"{model_dir}/easyocr_reader.pkl"
+        self.reader = torch.load(model_path)  # Load the EasyOCR model
 
-    def __call__(self, data):
+    def __call__(self, data: Dict):
         """
-        Handle the inference request.
+        Perform inference on the input image.
         Args:
-            data (dict): Input data for the model.
+            data (Dict): Input dictionary with keys:
+                - `inputs` containing image bytes or a file.
         Returns:
-            dict: Model predictions.
+            str: Extracted text.
         """
-        # Extract the image path or image content from the data
-        image_path = data.get("inputs")
-        if not image_path:
+        # Check if the input is an image
+        if "inputs" not in data:
             return {"error": "No input image provided"}
 
+        image_bytes = data["inputs"]
+        image = Image.open(io.BytesIO(image_bytes))
+
         # Perform OCR
-        try:
-            results = self.reader.readtext(image_path, detail=0)
-            return {"predictions": results}
-        except Exception as e:
-            return {"error": str(e)}
+        results = self.reader.readtext(image)
+
+        # Extract and return the detected text
+        extracted_text = " ".join([text for (_, text, _) in results])
+        return {"text": extracted_text}