78 lines
2.6 KiB
Python
78 lines
2.6 KiB
Python
 
|
"""OCR engine for BadNote using EasyOCR.
|
||
|
|
|
||
|
|
Lightweight handwriting-capable OCR using EasyOCR with CPU-only inference.
|
||
|
|
Suitable for Zen2 APU 25W / 16GB RAM (~200MB memory once loaded).
|
||
|
|
"""
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import logging
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
# NOTE: `easyocr` (and its torch dependency) is heavy and optional. It is
|
||
|
|
# imported lazily inside the engine so that importing the FastAPI app — and
|
||
|
|
# running its test suite — does not require the OCR dependencies. Install them
|
||
|
|
# with `pip install -r requirements-ocr.txt` when running the OCR worker.
|
||
|
|
|
||
|
|
|
||
|
|
class OcrEngine:
|
||
|
|
"""EasyOCR-based text recognition engine.
|
||
|
|
|
||
|
|
Lazy-loads the reader on first use to avoid startup overhead.
|
||
|
|
Supports Chinese (simplified) + English. Runs on CPU only.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
self._reader = None
|
||
|
|
|
||
|
|
def _ensure_reader(self):
|
||
|
|
"""Lazy-initialize the EasyOCR reader."""
|
||
|
|
if self._reader is None:
|
||
|
|
import easyocr # imported lazily; see module docstring note
|
||
|
|
|
||
|
|
logger.info("Loading EasyOCR reader (ch_sim + en, CPU)...")
|
||
|
|
self._reader = easyocr.Reader(['ch_sim', 'en'], gpu=False)
|
||
|
|
logger.info("EasyOCR reader loaded")
|
||
|
|
|
||
|
|
async def recognize(self, image_bytes: bytes) -> str:
|
||
|
|
"""Recognize text from image bytes.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
image_bytes: Raw image file bytes (PNG, JPEG, etc.)
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Recognized text as a single string, or empty string on failure.
|
||
|
|
"""
|
||
|
|
if not image_bytes:
|
||
|
|
return ""
|
||
|
|
|
||
|
|
loop = asyncio.get_event_loop()
|
||
|
|
|
||
|
|
try:
|
||
|
|
self._ensure_reader()
|
||
|
|
results = await loop.run_in_executor(
|
||
|
|
None, self._reader.readtext, image_bytes
|
||
|
|
)
|
||
|
|
# results is a list of (bbox, text, confidence) tuples
|
||
|
|
text_parts = [text for _, text, _ in results if text.strip()]
|
||
|
|
return ' '.join(text_parts)
|
||
|
|
except Exception as exc:
|
||
|
|
logger.error("OCR recognition failed: %s", exc)
|
||
|
|
return ""
|
||
|
|
|
||
|
|
async def recognize_file(self, image_path: str) -> str:
|
||
|
|
"""Recognize text from an image file path.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
image_path: Path to the image file on disk.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Recognized text as a single string, or empty string on failure.
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
with open(image_path, 'rb') as f:
|
||
|
|
image_bytes = f.read()
|
||
|
|
return await self.recognize(image_bytes)
|
||
|
|
except Exception as exc:
|
||
|
|
logger.error("Failed to read image file %s: %s", image_path, exc)
|
||
|
|
return ""
|