src/ocr_processor.py - master - kuzminaa/PDM

src/ocr_processor.py
83 строки3 KB
kuzminaa
refactor: migrate modules to src/, add logging system, add assembly workshop
29 июл 2026, 23:52
29 июл 2026, 23:52a5323aa
# -*- coding: utf-8 -*-
"""
Модуль: ocr_processor.py
Назначение: Зональный OCR штампов ЕСКД с извлечением Обозначения, 
            текущего листа и общего количества листов.
Платформа: ОС Debian 13 (python3-tesserocr, python3-pymupdf).
"""

import re
import fitz  # PyMuPDF
import tesserocr

ESKD_PATTERN = re.compile(r'[А-Я]{4}\.\d{6}\.\d{3}(?:-\d{2})?')
# Регулярные выражения для текущего листа и общего количества листов
SHEET_PATTERN = re.compile(r'(?:Лист|Л\.)\s*(\d+)', re.IGNORECASE)
TOTAL_SHEETS_PATTERN = re.compile(r'(?:Листов|Л-ов|Лтв)\s*(\d+)', re.IGNORECASE)

class EskdSmartOcrProcessor:
    def __init__(self, lang="rus"):
        self.lang = lang

    def process_and_extract(self, pdf_path, page_index):
        """Многопроходной конвейер рендеринга и распознавания с компенсацией скосов."""
        doc = fitz.open(pdf_path)
        page = doc[page_index]
        rect = page.rect
        
        # Перебираем варианты ориентации (0° и 180°) и углы скоса
        for flip in [False, True]:
            for angle in [0.0, -2.0, 2.0, -4.0, 4.0]:
                des, sheet, total = self._ocr_with_rotation(page, rect, angle, flip_180=flip)
                if des:
                    doc.close()
                    return des, sheet, total, angle, flip

        doc.close()
        return None, None, None, 0.0, False

    def _ocr_with_rotation(self, page, rect, angle, flip_180=False):
        w, h = rect.width, rect.height
        matrix = fitz.Matrix(2.0, 2.0)
        if angle != 0.0:
            matrix.prerotate(angle)
        if flip_180:
            matrix.prerotate(180)
            stamp_region = fitz.Rect(0, 0, w, 140)  # Увеличили зону для захвата графы "Листов"
        else:
            stamp_region = fitz.Rect(0, h - 140, w, h)

        try:
            pix = page.get_pixmap(matrix=matrix, clip=stamp_region)
            img_bytes = pix.tobytes("png") if hasattr(pix, "tobytes") else pix.getImageData(format="png")

            with tesserocr.PyTessBaseAPI(lang=self.lang) as api:
                api.SetImageBytes(img_bytes, pix.width, pix.height, 3, pix.width * 3)
                text = api.GetUTF8Text()

            return self._parse_text(text)
        except Exception:
            return None, None, None

    def _parse_text(self, text):
        if not text: 
            return None, None, None
            
        # Убираем пробелы для склеивания разорванных OCR децимальных номеров
        clean_text = " ".join(text.split())
        no_spaces_text = clean_text.replace(" ", "")
        
        des_match = ESKD_PATTERN.search(no_spaces_text)
        if not des_match:
            return None, None, None
            
        des = des_match.group(0)
        
        # Ищем текущий лист и общее количество листов
        sheet_match = SHEET_PATTERN.search(clean_text)
        sheet_num = int(sheet_match.group(1)) if sheet_match else None
        
        total_match = TOTAL_SHEETS_PATTERN.search(clean_text)
        total_sheets = int(total_match.group(1)) if total_match else None
        
        return des, sheet_num, total_sheets