File size: 5,296 Bytes
d520909 |
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 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 |
"""
OCR Engine Factory
Provides convenient functions to create and manage OCR engines.
Handles fallback logic and singleton management.
"""
from typing import Optional, Dict
from loguru import logger
from .base import OCREngine, OCRConfig
from .paddle_ocr import PaddleOCREngine, HAS_PADDLEOCR
from .tesseract_ocr import TesseractOCREngine, HAS_TESSERACT
# Singleton instances for reuse
_ocr_engines: Dict[str, OCREngine] = {}
def create_ocr_engine(
engine_type: str = "auto",
config: Optional[OCRConfig] = None,
initialize: bool = True,
) -> OCREngine:
"""
Create an OCR engine instance.
Args:
engine_type: Engine type: "paddle", "paddleocr", "tesseract", or "auto"
config: OCR configuration
initialize: Whether to initialize the engine immediately
Returns:
OCREngine instance
Raises:
RuntimeError: If no OCR engine is available
"""
if config is None:
config = OCRConfig()
# Normalize engine type aliases
if engine_type == "paddleocr":
engine_type = "paddle"
# Auto-select best available engine
if engine_type == "auto":
if HAS_PADDLEOCR:
engine_type = "paddle"
logger.info("Auto-selected PaddleOCR engine")
elif HAS_TESSERACT:
engine_type = "tesseract"
logger.info("Auto-selected Tesseract engine")
else:
raise RuntimeError(
"No OCR engine available. Install one of: "
"pip install paddleocr paddlepaddle-gpu OR "
"pip install pytesseract (+ apt-get install tesseract-ocr)"
)
# Create engine
if engine_type == "paddle":
if not HAS_PADDLEOCR:
raise RuntimeError(
"PaddleOCR not installed. Install with: "
"pip install paddleocr paddlepaddle-gpu"
)
engine = PaddleOCREngine(config)
elif engine_type == "tesseract":
if not HAS_TESSERACT:
raise RuntimeError(
"Tesseract not installed. Install with: "
"pip install pytesseract (+ apt-get install tesseract-ocr)"
)
engine = TesseractOCREngine(config)
else:
raise ValueError(f"Unknown engine type: {engine_type}")
# Initialize if requested
if initialize:
engine.initialize()
return engine
def get_ocr_engine(
engine_type: str = "auto",
config: Optional[OCRConfig] = None,
) -> OCREngine:
"""
Get or create an OCR engine singleton.
Reuses existing engine instances for efficiency.
Args:
engine_type: Engine type: "paddle", "paddleocr", "tesseract", or "auto"
config: OCR configuration (only used for new instances)
Returns:
OCREngine instance
"""
global _ocr_engines
# Normalize engine type aliases
if engine_type == "paddleocr":
engine_type = "paddle"
# Resolve auto to specific type
if engine_type == "auto":
if HAS_PADDLEOCR:
engine_type = "paddle"
elif HAS_TESSERACT:
engine_type = "tesseract"
else:
raise RuntimeError("No OCR engine available")
# Check for existing instance
if engine_type in _ocr_engines:
return _ocr_engines[engine_type]
# Create new instance
engine = create_ocr_engine(engine_type, config, initialize=True)
_ocr_engines[engine_type] = engine
return engine
def get_available_engines() -> Dict[str, bool]:
"""
Get availability status of OCR engines.
Returns:
Dict mapping engine name to availability
"""
return {
"paddle": HAS_PADDLEOCR,
"tesseract": HAS_TESSERACT,
}
def clear_engines():
"""Clear all cached OCR engine instances."""
global _ocr_engines
_ocr_engines.clear()
logger.debug("Cleared OCR engine cache")
class OCREngineManager:
"""
Context manager for OCR engine lifecycle.
Example:
with OCREngineManager("paddle") as engine:
result = engine.recognize(image)
"""
def __init__(
self,
engine_type: str = "auto",
config: Optional[OCRConfig] = None,
use_singleton: bool = True,
):
"""
Initialize OCR engine manager.
Args:
engine_type: Engine type
config: OCR configuration
use_singleton: Whether to use singleton instance
"""
self.engine_type = engine_type
self.config = config
self.use_singleton = use_singleton
self._engine: Optional[OCREngine] = None
self._owned = False
def __enter__(self) -> OCREngine:
"""Enter context and return engine."""
if self.use_singleton:
self._engine = get_ocr_engine(self.engine_type, self.config)
self._owned = False
else:
self._engine = create_ocr_engine(self.engine_type, self.config)
self._owned = True
return self._engine
def __exit__(self, exc_type, exc_val, exc_tb):
"""Exit context."""
# Don't clean up singletons
if self._owned and self._engine:
# Could add cleanup here if needed
pass
self._engine = None
return False
|