Spaces:
Sleeping
Sleeping
File size: 2,878 Bytes
ba6e49b 1e73641 ba6e49b 1e73641 ba6e49b 1e73641 ba6e49b |
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 |
from abc import ABC, abstractmethod
from typing import Set
# Lazy imports to handle missing dependencies gracefully
try:
import easyocr
EASYOCR_AVAILABLE = True
except ImportError:
EASYOCR_AVAILABLE = False
try:
import pytesseract
PYTESSERACT_AVAILABLE = True
except ImportError:
PYTESSERACT_AVAILABLE = False
try:
import cv2
import numpy as np
CV2_AVAILABLE = True
except ImportError:
CV2_AVAILABLE = False
class OCRStrategy(ABC):
"""Abstract base class for OCR strategies."""
@abstractmethod
def extract_text(self, image_path: str) -> Set[str]:
pass
class EasyOCRStrategy(OCRStrategy):
"""Concrete strategy for EasyOCR."""
def __init__(self):
if not EASYOCR_AVAILABLE:
raise ImportError(
"EasyOCR is not installed. Please install it with: pip install easyocr"
)
# Initialize once to save memory/time
print("Loading EasyOCR Model...")
self.reader = easyocr.Reader(['en'], gpu=False)
def extract_text(self, image_path: str) -> Set[str]:
results = self.reader.readtext(image_path, detail=0)
return {text.lower().strip() for text in results if text.strip()}
class TesseractOCRStrategy(OCRStrategy):
"""Concrete strategy for Tesseract OCR (Free & Fast)."""
def __init__(self):
if not PYTESSERACT_AVAILABLE:
raise ImportError(
"pytesseract is not installed. Please install it with: pip install pytesseract"
)
if not CV2_AVAILABLE:
raise ImportError(
"opencv-python is not installed. Please install it with: pip install opencv-python-headless"
)
def extract_text(self, image_path: str) -> Set[str]:
# Preprocessing for better Tesseract accuracy
img = cv2.imread(image_path)
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
# Apply thresholding
_, thresh = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY)
# Tesseract execution
text = pytesseract.image_to_string(thresh)
# Process results
lines = text.split('\n')
return {line.lower().strip() for line in lines if line.strip()}
class OCRFactory:
"""Factory to get the correct OCR strategy."""
_instances = {}
@staticmethod
def get_strategy(strategy_name: str) -> OCRStrategy:
if strategy_name not in OCRFactory._instances:
if strategy_name == "EasyOCR (Best Accuracy)":
OCRFactory._instances[strategy_name] = EasyOCRStrategy()
elif strategy_name == "Tesseract (Fast & Free)":
OCRFactory._instances[strategy_name] = TesseractOCRStrategy()
else:
raise ValueError(f"Unknown strategy: {strategy_name}")
return OCRFactory._instances[strategy_name] |