Spaces:
Runtime error
Runtime error
File size: 4,211 Bytes
22997db | 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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | # api/inference.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image, ImageOps
import numpy as np
from io import BytesIO
import os
from src.model import OptimizedOCR
class OCRPredictor:
def __init__(self):
self.device = torch.device("cpu") # Force CPU for Koyeb
self.model = None
self.idx_to_label = self._create_label_map()
def _create_label_map(self):
"""Creates the label mapping dictionary"""
label_map = {}
# Add digits (0-9)
for i in range(10):
label_map[i] = str(i)
# Add uppercase letters (A-Z)
for i in range(26):
label_map[i + 10] = chr(65 + i)
return label_map
async def _load_model(self):
"""Lazy loading of the model"""
if self.model is None:
try:
self.model = OptimizedOCR(num_classes=36).to(self.device)
model_path = 'best_model.pth'
if not os.path.exists(model_path):
raise FileNotFoundError(f"Model file not found at {model_path}")
# Load model with map_location to ensure CPU loading
self.model.load_state_dict(
torch.load(model_path, map_location=self.device)
)
self.model.eval()
except Exception as e:
raise Exception(f"Error loading model: {str(e)}")
def _enhance_image(self, image: Image.Image) -> Image.Image:
"""Enhance image quality for better OCR"""
# Convert to grayscale
image = ImageOps.grayscale(image)
# Increase contrast
image = ImageOps.autocontrast(image, cutoff=10)
# Resize to 27x27 using Lanczos resampling
image = image.resize((27, 27), Image.LANCZOS)
return image
async def _preprocess_image(self, image_data: bytes) -> torch.Tensor:
"""Preprocesses the input image for model prediction"""
try:
# Open image using PIL
image = Image.open(BytesIO(image_data))
# Enhance image
image = self._enhance_image(image)
# Convert to numpy array
image_array = np.array(image)
# Normalize and binarize
image_array = image_array.astype(np.float32) / 255.0
threshold = np.mean(image_array)
binary_image = (image_array > threshold).astype(np.float32)
# Ensure dark background (invert if needed)
if np.mean(binary_image) > 0.5:
binary_image = 1.0 - binary_image
# Add batch and channel dimensions
image_tensor = torch.FloatTensor(binary_image).unsqueeze(0).unsqueeze(0)
return image_tensor
except Exception as e:
raise Exception(f"Error preprocessing image: {str(e)}")
async def predict_character(self, file) -> dict:
"""Predicts the character in the uploaded image"""
try:
# Load model if not loaded
await self._load_model()
# Read image data
image_data = await file.read()
# Preprocess image
image_tensor = await self._preprocess_image(image_data)
image_tensor = image_tensor.to(self.device)
# Make prediction
with torch.no_grad():
output = self.model(image_tensor)
probabilities = torch.softmax(output, dim=1)
predicted_idx = output.argmax(1).item()
confidence = float(probabilities[0][predicted_idx])
predicted_char = self.idx_to_label[predicted_idx]
return {
"character": predicted_char,
"confidence": confidence,
"status": "success"
}
except Exception as e:
return {
"status": "error",
"error": str(e)
} |