core/ocr_engine.py - master - chalykhii_OG/OCR_PDF

core/ocr_engine.py
77 строк2 KB
chalykhii_OG
first_commit
07 июл 2026, 18:53
07 июл 2026, 18:531cae060
import os
import warnings
from contextlib import contextmanager, redirect_stderr, redirect_stdout
from dataclasses import dataclass
from pathlib import Path

from config import OCR_CACHE_DIR, OUTPUT_DIR, TOOLS_DIR


@dataclass(frozen=True)
class OcrResult:
    image_path: Path
    lines: list[str]

    @property
    def text(self) -> str:
        return "\n".join(self.lines)


class OcrEngine:
    def __init__(self, lang: str = "ru") -> None:
        OCR_CACHE_DIR.mkdir(parents=True, exist_ok=True)
        os.environ.setdefault("PADDLE_PDX_CACHE_HOME", str(OCR_CACHE_DIR))
        os.environ.setdefault("PADDLE_PDX_ENABLE_MKLDNN_BYDEFAULT", "False")
        os.environ.setdefault("FLAGS_use_mkldnn", "0")
        os.environ.setdefault("GLOG_minloglevel", "2")
        os.environ.setdefault("PADDLE_CPP_LOG_LEVEL", "ERROR")
        _prepare_paddle_environment()

        with _quiet_paddle_output():
            from paddleocr import PaddleOCR

            self._ocr = PaddleOCR(
                lang=lang,
                use_doc_orientation_classify=False,
                use_doc_unwarping=False,
                use_textline_orientation=False,
            )

    def recognize(self, image_path: Path) -> OcrResult:
        with _quiet_paddle_output():
            raw_result = self._ocr.predict(str(image_path))
        lines = _extract_lines(raw_result)
        return OcrResult(image_path=image_path, lines=lines)


def _extract_lines(raw_result: object) -> list[str]:
    lines: list[str] = []

    if isinstance(raw_result, list):
        for page in raw_result:
            if isinstance(page, dict):
                lines.extend(str(text) for text in page.get("rec_texts", []) if text)

    return lines


def _prepare_paddle_environment() -> None:
    TOOLS_DIR.mkdir(exist_ok=True)
    ccache_stub = TOOLS_DIR / "ccache"
    ccache_stub.touch(exist_ok=True)

    path_parts = os.environ.get("PATH", "").split(os.pathsep)
    tools_path = str(TOOLS_DIR)
    if tools_path not in path_parts:
        os.environ["PATH"] = tools_path + os.pathsep + os.environ.get("PATH", "")


@contextmanager
def _quiet_paddle_output():
    OUTPUT_DIR.mkdir(exist_ok=True)
    log_path = OUTPUT_DIR / "ocr.log"
    with log_path.open("a", encoding="utf-8") as log_file:
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", message="No ccache found.*")
            with redirect_stdout(log_file), redirect_stderr(log_file):
                yield