scan4llm/src/ocr/paddle_engine.py

from __future__ import annotations

import importlib.resources
from dataclasses import dataclass, field
from typing import Any

import cv2
import yaml
from paddlex import create_pipeline


@dataclass
class ParsedBlock:
    label: str
    content: str
    bbox: tuple[int, int, int, int]
    confidence: float = 0.9


@dataclass
class OcrPageResult:
    blocks: list[ParsedBlock] = field(default_factory=list)
    raw_json: dict[str, Any] = field(default_factory=dict)
    width: int = 0
    height: int = 0


def _load_config() -> dict[str, Any]:
    ref = importlib.resources.files("paddlex") / "configs" / "pipelines" / "PP-StructureV3.yaml"
    with importlib.resources.as_file(ref) as path, open(path) as f:
        config: dict[str, Any] = yaml.safe_load(f)

    config["SubPipelines"]["GeneralOCR"]["SubModules"]["TextRecognition"]["model_name"] = (
        "eslav_PP-OCRv5_mobile_rec"
    )
    table_ocr = config["SubPipelines"]["TableRecognition"]["SubPipelines"]["GeneralOCR"]
    table_ocr["SubModules"]["TextRecognition"]["model_name"] = "eslav_PP-OCRv5_mobile_rec"
    return config


class OcrEngine:
    def __init__(self) -> None:
        self._config: dict[str, Any] = _load_config()
        self._pipeline: Any = None

    def process(self, image_path: str) -> OcrPageResult:
        image = cv2.imread(image_path)
        if image is None:
            raise FileNotFoundError(f"Не удалось загрузить изображение: {image_path}")

        h, w = image.shape[:2]
        if self._pipeline is None:
            self._pipeline = create_pipeline(
                config=self._config,
                use_doc_orientation_classify=False,
                use_doc_unwarping=False,
            )
        raw_results = list(self._pipeline.predict(image))

        blocks, raw_json = self._parse_results(raw_results)
        return OcrPageResult(blocks=blocks, raw_json=raw_json, width=w, height=h)

    @staticmethod
    def _parse_results(raw_results: list[Any]) -> tuple[list[ParsedBlock], dict[str, Any]]:
        if not raw_results:
            return [], {}

        raw_json: dict[str, Any] = raw_results[0].json["res"]
        blocks: list[ParsedBlock] = []
        for item in raw_json.get("parsing_res_list", []):
            text = item.get("block_content", "")
            label = item.get("block_label", "text")
            bbox = item.get("block_bbox", [0, 0, 0, 0])
            if not text.strip():
                continue
            blocks.append(ParsedBlock(
                label=label, content=text,
                bbox=(int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])),
            ))
        return blocks, raw_json


_engine: OcrEngine | None = None


def ocr_image(image_path: str) -> OcrPageResult:
    global _engine
    if _engine is None:
        _engine = OcrEngine()
    return _engine.process(image_path)