File size: 2,654 Bytes
efe181a 4dd3e86 fa0e3f8 38b8548 74e00d4 efe181a 38b8548 efe181a 38b8548 efe181a fa0e3f8 38b8548 fa0e3f8 38b8548 efe181a 38b8548 dec7250 38b8548 efe181a 38b8548 dec7250 38b8548 efe181a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
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) |