Update handler.py

handler.py

@@ -1,49 +1,26 @@
+import os
 import io
 import easyocr
 import numpy as np
 from typing import Dict, List, Union
 from PIL import Image
 
-class OCRInferenceHandler:
-    def __init__(self, model_dir=None):
+class EndpointHandler:
+    def __init__(self, path: str):
         """
         Initialize the OCR inference handler
         
         Args:
-            model_dir (str, optional): Directory containing model artifacts
+            path (str): Directory containing model artifacts
         """
         try:
-            # Initialize EasyOCR reader directly
-            # You can specify languages if needed, e.g., ['en', 'fr']
+            # Detect preferred languages from the model directory name or default to English
+            # You can modify this logic to detect or configure languages
             self.reader = easyocr.Reader(['en'])
         except Exception as e:
             raise RuntimeError(f"Error initializing OCR model: {str(e)}")
     
-    def preprocess(self, input_data: Union[bytes, Image.Image]) -> np.ndarray:
-        """
-        Preprocess the input image
-        
-        Args:
-            input_data (Union[bytes, Image.Image]): Input image in bytes or PIL Image
-        
-        Returns:
-            np.ndarray: Processed image array
-        """
-        # Convert input to PIL Image if it's bytes
-        if isinstance(input_data, bytes):
-            try:
-                image = Image.open(io.BytesIO(input_data))
-            except Exception as e:
-                raise ValueError(f"Invalid image format: {str(e)}")
-        elif isinstance(input_data, Image.Image):
-            image = input_data
-        else:
-            raise TypeError("Input must be bytes or PIL Image")
-        
-        # Convert to numpy array
-        return np.array(image)
-    
-    def predict(self, input_data: Union[bytes, Image.Image]) -> Dict[str, List[Dict]]:
+    def __call__(self, input_data: Union[bytes, Image.Image]) -> Dict[str, List[Dict]]:
         """
         Perform OCR inference
         
@@ -55,7 +32,7 @@ class OCRInferenceHandler:
         """
         try:
             # Preprocess the image
-            img_array = self.preprocess(input_data)
+            img_array = self._preprocess(input_data)
             
             # Perform OCR detection
             results = self.reader.readtext(img_array)
@@ -80,32 +57,27 @@ class OCRInferenceHandler:
                 'success': False,
                 'error': str(e)
             }
-
-# Hugging Face Inference API handler
-def handler(input_data):
-    """
-    Main handler for Hugging Face Inference API
-    
-    Args:
-        input_data: Input image data
     
-    Returns:
-        OCR inference results
-    """
-    # Ensure the model is loaded only once (singleton pattern)
-    if not hasattr(handler, 'inference'):
-        handler.inference = OCRInferenceHandler()
-    
-    # Run inference
-    return handler.inference.predict(input_data)
-
-# Initialization method for model loading
-def init(model_dir=None):
-    """
-    Initialization method for pre-loading the model
-    
-    Args:
-        model_dir (str, optional): Directory containing model artifacts
-    """
-    if not hasattr(handler, 'inference'):
-        handler.inference = OCRInferenceHandler(model_dir)
+    def _preprocess(self, input_data: Union[bytes, Image.Image]) -> np.ndarray:
+        """
+        Preprocess the input image
+        
+        Args:
+            input_data (Union[bytes, Image.Image]): Input image in bytes or PIL Image
+        
+        Returns:
+            np.ndarray: Processed image array
+        """
+        # Convert input to PIL Image if it's bytes
+        if isinstance(input_data, bytes):
+            try:
+                image = Image.open(io.BytesIO(input_data))
+            except Exception as e:
+                raise ValueError(f"Invalid image format: {str(e)}")
+        elif isinstance(input_data, Image.Image):
+            image = input_data
+        else:
+            raise TypeError("Input must be bytes or PIL Image")
+        
+        # Convert to numpy array
+        return np.array(image)