handler.py

import os
import io
import easyocr
import numpy as np
from typing import Dict, List, Union
from PIL import Image

class EndpointHandler:
    def __init__(self, path: str):
        """
        Initialize the OCR inference handler
        
        Args:
            path (str): Directory containing model artifacts
        """
        try:
            # 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 __call__(self, input_data: Union[bytes, Image.Image]) -> Dict[str, List[Dict]]:
        """
        Perform OCR inference
        
        Args:
            input_data (Union[bytes, Image.Image]): Input image to process
        
        Returns:
            Dict containing OCR results
        """
        try:
            # Preprocess the image
            img_array = self._preprocess(input_data)
            
            # Perform OCR detection
            results = self.reader.readtext(img_array)
            
            # Process results
            processed_results = []
            for detection in results:
                bbox, text, confidence = detection
                processed_results.append({
                    'text': text,
                    'confidence': float(confidence),
                    'bbox': [list(map(int, coord)) for coord in bbox]
                })
            
            return {
                'success': True,
                'results': processed_results
            }
        
        except Exception as e:
            return {
                'success': False,
                'error': 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)