src/ocr.py

"""Image OCR + translation helper."""
from __future__ import annotations

import logging
from typing import Tuple

import numpy as np
from PIL import Image

from src.translator import translate

logger = logging.getLogger(__name__)

# Lazy EasyOCR loader
_ocr_reader = None


def _get_reader():
    global _ocr_reader
    if _ocr_reader is None:
        import easyocr
        logger.info("Initializing EasyOCR (ja, en)...")
        # Use CPU to avoid competing for GPU memory with the translation model.
        _ocr_reader = easyocr.Reader(["ja", "en"], gpu=False)
    return _ocr_reader


def ocr_image(image: Image.Image) -> Tuple[str, str]:
    """Run OCR on a PIL image and translate the result.

    Returns (detected_text, english_translation).
    """
    if image is None:
        return "", ""

    try:
        arr = np.array(image.convert("RGB"))
        reader = _get_reader()
        results = reader.readtext(arr)
        lines = [r[1] for r in results]
        detected = "\n".join(lines).strip()
        if not detected:
            return "(no text detected)", ""
        translated, _ = translate(detected)
        return detected, translated
    except Exception as e:
        logger.exception("OCR failed")
        return f"OCR error: {e}", ""