yimlo / dependencies.py
One2SkipAfew's picture
Upload 15 files
a8c31fd verified
Raw
History Blame Contribute Delete
1.19 kB
import re
import easyocr
import torch
import numpy as np
from PIL import Image, ImageOps
# Initialize EasyOCR Reader globally so it stays in memory once
# This will use GPU (CUDA) if available
reader = easyocr.Reader(['en'], gpu=torch.cuda.is_available())
def resize_image(image: Image.Image, max_dim: int = 1500) -> Image.Image:
"""Resizes image if either dimension exceeds max_dim to speed up CPU OCR."""
w, h = image.size
if w <= max_dim and h <= max_dim:
return image
scale = max_dim / max(w, h)
new_size = (int(w * scale), int(h * scale))
print(f"--- [CPU OPTIMIZATION] Resizing from {w}x{h} to {new_size[0]}x{new_size[1]} ---")
return image.resize(new_size, Image.Resampling.LANCZOS)
def normalize_text(text: str) -> str:
"""Removes special characters, extra spaces, and converts to uppercase for reliable matching."""
if not text: return ""
text = text.upper()
return re.sub(r'[^A-Z0-9]', '', text)
def preprocess_image(image: Image.Image) -> Image.Image:
"""Applies basic grayscale to reduce dimensionality while preserving natural contrast for EasyOCR."""
image = ImageOps.grayscale(image)
return image