src/ocr_worker.py

"""Subprocess worker for PaddleOCR.

Running PaddleOCR in a separate process protects the FastAPI server from
native-library crashes in PaddlePaddle.
"""

import json
import os
import sys
from pathlib import Path


def extract_text_items(ocr_result) -> list[str]:
    """Normalize PaddleOCR output from common v2 and v3 response formats."""
    text_items = []

    if not ocr_result:
        return text_items

    if isinstance(ocr_result, list):
        for item in ocr_result:
            text_items.extend(extract_text_items(item))
        return text_items

    if isinstance(ocr_result, dict):
        for key in ("rec_texts", "texts"):
            values = ocr_result.get(key)
            if isinstance(values, list):
                text_items.extend(str(value) for value in values if str(value).strip())

        for key in ("res", "data"):
            if key in ocr_result:
                text_items.extend(extract_text_items(ocr_result[key]))

        return text_items

    if isinstance(ocr_result, tuple) and len(ocr_result) >= 2:
        possible_text = ocr_result[1]
        if isinstance(possible_text, tuple) and possible_text:
            text_items.append(str(possible_text[0]))
        elif isinstance(possible_text, str):
            text_items.append(possible_text)
        return text_items

    return text_items


def run_ocr(image_path: Path) -> str:
    """Run PaddleOCR on one image and return cleaned text."""
    from paddleocr import PaddleOCR

    try:
        ocr = PaddleOCR(
            lang="en",
            use_doc_orientation_classify=False,
            use_doc_unwarping=False,
            use_textline_orientation=False,
        )
    except ValueError:
        ocr = PaddleOCR(use_angle_cls=True, lang="en")

    if hasattr(ocr, "predict"):
        result = ocr.predict(str(image_path))
    else:
        result = ocr.ocr(str(image_path), cls=True)

    text_items = extract_text_items(result)
    return " ".join(" ".join(text_items).split())


def main() -> int:
    """CLI entry point."""
    if len(sys.argv) != 3:
        print(json.dumps({"error": "Usage: python -m src.ocr_worker IMAGE_PATH CACHE_DIR"}))
        return 2

    image_path = Path(sys.argv[1])
    cache_dir = Path(sys.argv[2])
    cache_dir.mkdir(parents=True, exist_ok=True)

    os.environ.setdefault("PADDLE_PDX_CACHE_HOME", str(cache_dir))
    os.environ.setdefault("PADDLEOCR_DISABLE_AUTO_LOGGING_CONFIG", "1")

    try:
        text = run_ocr(image_path)
    except Exception as error:
        print(json.dumps({"error": str(error)}))
        return 1

    print(json.dumps({"text": text}))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())