scan4llm/src/postprocess/corrector.py

from __future__ import annotations

import logging
from typing import Protocol

from src.ocr.paddle_engine import ParsedBlock

logger = logging.getLogger(__name__)


class Corrector(Protocol):
    def __call__(self, block: ParsedBlock) -> ParsedBlock: ...


def apply_corrector(
    blocks: list[ParsedBlock],
    corrector: Corrector,
    name: str,
    skip_formulas: bool = True,
) -> None:
    for i, block in enumerate(blocks):
        if skip_formulas and block.label == "formula":
            continue
        logger.info("%s: блок %d/%d...", name, i + 1, len(blocks))
        corrector(block)
    logger.info("%s: завершено", name)