70 lines
2.2 KiB
Python
70 lines
2.2 KiB
Python
|
"""Background OCR worker for BadNote.
|
||
|
|
|
||
|
|
Polls the file-based queue and processes jobs using the OcrEngine.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import logging
|
||
|
|
import os
|
||
|
|
|
||
|
|
from ..config import settings
|
||
|
|
from .engine import OcrEngine
|
||
|
|
from . import queue as job_queue
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
async def run_worker(poll_interval: int = 5) -> None:
|
||
|
|
"""Poll the queue and process OCR jobs.
|
||
|
|
|
||
|
|
Loads images from the job's image_path and runs them through EasyOCR.
|
||
|
|
"""
|
||
|
|
# Ensure queue dirs exist
|
||
|
|
for subdir in ("pending", "processing", "done", "failed"):
|
||
|
|
os.makedirs(os.path.join(settings.queue_path, subdir), exist_ok=True)
|
||
|
|
|
||
|
|
engine = OcrEngine()
|
||
|
|
logger.info("OCR worker started (poll_interval=%ds)", poll_interval)
|
||
|
|
|
||
|
|
while True:
|
||
|
|
job = job_queue.dequeue()
|
||
|
|
if job is not None:
|
||
|
|
job_id = job["id"]
|
||
|
|
logger.info("Processing OCR job %s", job_id)
|
||
|
|
try:
|
||
|
|
# Read image from the path specified in the job
|
||
|
|
image_path = job.get("image_path", "")
|
||
|
|
if image_path and os.path.exists(image_path):
|
||
|
|
result = await engine.recognize_file(image_path)
|
||
|
|
else:
|
||
|
|
# Fall back to image_bytes if provided inline
|
||
|
|
image_bytes = job.get("image_bytes", b"")
|
||
|
|
if isinstance(image_bytes, str):
|
||
|
|
import base64
|
||
|
|
image_bytes = base64.b64decode(image_bytes)
|
||
|
|
result = await engine.recognize(image_bytes)
|
||
|
|
|
||
|
|
job_queue.complete(job_id, result)
|
||
|
|
logger.info("OCR job %s completed: %d chars", job_id, len(result))
|
||
|
|
except Exception as exc:
|
||
|
|
logger.error("OCR job %s failed: %s", job_id, exc)
|
||
|
|
job_queue.fail(job_id, str(exc))
|
||
|
|
else:
|
||
|
|
await asyncio.sleep(poll_interval)
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
"""Entry point for `python -m badnote_server.ocr.worker`."""
|
||
|
|
logging.basicConfig(
|
||
|
|
level=logging.INFO,
|
||
|
|
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
asyncio.run(run_worker())
|
||
|
|
except KeyboardInterrupt:
|
||
|
|
logger.info("OCR worker stopped")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|