ocr_parser.py - main - Paff/declarant

ocr_parser.py
256 строк9 KB
1pash1985-gif
Декларант - веб-версия для Timeweb Cloud
15 июл 2026, 09:36
15 июл 2026, 09:36a6b08b8
"""
OCR парсер для инвойсов на основе EasyOCR.
Используется как fallback, когда обычный PDF-парсинг не извлекает текст.
"""
import os
import sys
import tempfile
from typing import Callable

import numpy as np
from PIL import Image


# Глобальный экземпляр InvoiceOCR для переиспользования
_ocr_instance: "InvoiceOCR | None" = None


def is_ocr_available() -> bool:
    """Проверить, доступен ли OCR."""
    try:
        import easyocr
        return True
    except ImportError:
        return False


def get_ocr_version() -> str:
    """Получить версию установленного EasyOCR."""
    if not is_ocr_available():
        return "не установлен"
    try:
        import easyocr
        return getattr(easyocr, "__version__", "неизвестно")
    except Exception:
        return "неизвестно"


def get_downloaded_models() -> list[str]:
    """Получить список уже скачанных моделей EasyOCR."""
    model_dir = os.path.expanduser("~/.EasyOCR/model")
    if not os.path.isdir(model_dir):
        return []
    return sorted(
        name for name in os.listdir(model_dir)
        if os.path.isfile(os.path.join(model_dir, name)) and name.endswith((".pth", ".onnx", ".pt"))
    )


def ensure_ocr_models(
    languages: tuple[str, ...] = ("ru", "en"),
    status_callback: Callable[[str], None] | None = None,
    progress_callback: Callable[[int], None] | None = None,
):
    """Принудительно инициализировать EasyOCR для указанных языков и скачать модели."""
    if not is_ocr_available():
        raise RuntimeError("EasyOCR не установлен")

    original_stdout = sys.stdout
    original_stderr = sys.stderr

    class TeeStream:
        def __init__(self, stream):
            self.stream = stream

        def write(self, msg: str):
            self.stream.write(msg)
            self._process(msg)

        def flush(self):
            self.stream.flush()

        def _process(self, msg: str):
            if not status_callback:
                return
            stripped = msg.strip()
            if not stripped:
                return
            if any(key in stripped for key in ("Downloading", "downloading", "%", "progress", "Loading", "loading")):
                status_callback(stripped)
            if progress_callback and "%" in stripped:
                import re
                match = re.search(r"(\d+(?:\.\d+)?)\s*%", stripped)
                if match:
                    try:
                        pct = int(float(match.group(1)))
                        progress_callback(max(0, min(100, pct)))
                    except ValueError:
                        pass

    if status_callback:
        status_callback(f"Подготовка моделей: {', '.join(languages)}...")

    sys.stdout = TeeStream(original_stdout)
    sys.stderr = TeeStream(original_stderr)
    try:
        import easyocr
        reader = easyocr.Reader(list(languages), gpu=False, verbose=False)
        # Прочитать пустое изображение, чтобы форсировать полную загрузку моделей в память
        reader.readtext(np.zeros((10, 10, 3), dtype=np.uint8), detail=0)
    finally:
        sys.stdout = original_stdout
        sys.stderr = original_stderr

    if status_callback:
        status_callback("OCR готов")
    if progress_callback:
        progress_callback(100)


def get_ocr_instance() -> "InvoiceOCR":
    """Получить глобальный экземпляр InvoiceOCR."""
    global _ocr_instance
    if _ocr_instance is None:
        _ocr_instance = InvoiceOCR(languages=("ru", "en"))
    return _ocr_instance


def preload_ocr_models(
    status_callback: Callable[[str], None] | None = None,
    progress_callback: Callable[[int], None] | None = None,
):
    """Предзагрузить модели EasyOCR. Вызывается один раз.

    Args:
        status_callback: вызывается с текстовым статусом.
        progress_callback: вызывается с числом 0-100 (процент загрузки).
    """
    if not is_ocr_available():
        return

    if status_callback:
        status_callback("Инициализация OCR...")

    original_stdout = sys.stdout
    original_stderr = sys.stderr

    class TeeStream:
        """Перехватывает поток и передаёт строки в статус."""

        def __init__(self, stream):
            self.stream = stream

        def write(self, msg: str):
            self.stream.write(msg)
            self._process(msg)

        def flush(self):
            self.stream.flush()

        def _process(self, msg: str):
            if not status_callback:
                return
            stripped = msg.strip()
            if not stripped:
                return
            # Обновить статус если в сообщении есть ключевые слова загрузки
            if any(key in stripped for key in ("Downloading", "downloading", "%", "progress", "Loading", "loading")):
                status_callback(stripped)
            # Попытаться извлечь процент вида 12.3% или 45%
            if progress_callback and "%" in stripped:
                import re
                match = re.search(r"(\d+(?:\.\d+)?)\s*%", stripped)
                if match:
                    try:
                        pct = int(float(match.group(1)))
                        progress_callback(max(0, min(100, pct)))
                    except ValueError:
                        pass

    sys.stdout = TeeStream(original_stdout)
    sys.stderr = TeeStream(original_stderr)
    try:
        ocr = get_ocr_instance()
        ocr._get_reader()
    finally:
        sys.stdout = original_stdout
        sys.stderr = original_stderr

    if status_callback:
        status_callback("OCR готов")
    if progress_callback:
        progress_callback(100)


class InvoiceOCR:
    """OCR распознавание текста из PDF и изображений инвойсов."""

    def __init__(self, languages: tuple[str, ...] = ("ru", "en")):
        self.languages = languages
        self._reader = None

    def _get_reader(self):
        """Ленивая инициализация EasyOCR reader."""
        if self._reader is None:
            import easyocr
            self._reader = easyocr.Reader(
                list(self.languages),
                gpu=False,
                verbose=False
            )
        return self._reader

    def extract_text_from_image(self, image_path: str) -> str:
        """Распознать текст с изображения."""
        reader = self._get_reader()
        result = reader.readtext(image_path, detail=0)
        return "\n".join(result)

    def extract_text_from_pdf(self, pdf_path: str, dpi: int = 200) -> str:
        """Распознать текст из PDF файла."""
        if not os.path.exists(pdf_path):
            raise FileNotFoundError(f"PDF файл не найден: {pdf_path}")

        # Использовать PyMuPDF для конвертации страниц в изображения
        try:
            import fitz  # PyMuPDF
        except ImportError:
            raise ImportError("Для OCR PDF требуется PyMuPDF (fitz). Установите: pip install pymupdf")

        doc = fitz.open(pdf_path)
        images = []
        for page in doc:
            zoom = dpi / 72
            mat = fitz.Matrix(zoom, zoom)
            pix = page.get_pixmap(matrix=mat)
            # Конвертировать pixmap в PIL Image
            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            images.append(np.array(img))
        doc.close()

        if not images:
            return ""

        reader = self._get_reader()
        full_text = []

        for page_idx, img_array in enumerate(images, 1):
            result = reader.readtext(img_array, detail=0)
            page_text = "\n".join(result)
            if page_text.strip():
                full_text.append(f"--- Page {page_idx} ---\n{page_text}")

        return "\n\n".join(full_text)


def extract_text_with_ocr(file_path: str) -> str:
    """Универсальная функция для извлечения текста через OCR."""
    ocr = get_ocr_instance()
    ext = os.path.splitext(file_path)[1].lower()

    if ext == ".pdf":
        return ocr.extract_text_from_pdf(file_path)
    elif ext in (".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif"):
        return ocr.extract_text_from_image(file_path)
    else:
        raise ValueError(f"Неподдерживаемый формат файла для OCR: {ext}")