src/assembly_engine.py - master - kuzminaa/PDM

src/assembly_engine.py
163 строки9 KB
kuzminaa
refactor: migrate modules to src/, add logging system, add assembly workshop
29 июл 2026, 23:52
29 июл 2026, 23:52a5323aa
# -*- coding: utf-8 -*-
"""
Модуль: assembly_engine.py
Назначение: Умный сборочный стол с валидацией коллизий дубликатов,
            контролем общего объема листов и выявлением "дыр".
"""

import os
import logging
import fitz
from .ocr_processor import EskdSmartOcrProcessor

logger = logging.getLogger(__name__)

class SmartAssemblyEngine:
    def __init__(self):
        self.processor = EskdSmartOcrProcessor()
        # Основная база: { 'Обозначение': { 'expected_total': int/None, 'sheets': { листы }, 'errors': [ошибки] } }
        self.assembled_docs = {}
        self.suspect_pages = []  # Страницы, где вообще не прочиталось Обозначение

    def run_session_analysis(self, folder_path, progress_cb=None):
        self.assembled_docs.clear()
        self.suspect_pages.clear()
        
        files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.lower().endswith('.pdf')]
        total_files = len(files)

        for f_idx, f_path in enumerate(files):
            doc = fitz.open(f_path)
            p_count = len(doc)
            doc.close()

            for p_idx in range(p_count):
                if progress_cb:
                    progress_cb(f_idx + 1, total_files, os.path.basename(f_path), p_idx + 1, p_count)

                des, sheet_num, total_sheets, angle, flipped = self.processor.process_and_extract(f_path, p_idx)

                page_info = {
                    'path': f_path, 'idx': p_idx, 'angle': angle, 'flip': flipped,
                    'detected_sheet': sheet_num, 'detected_total': total_sheets
                }

                if des:
                    if des not in self.assembled_docs:
                        self.assembled_docs[des] = {'expected_total': None, 'sheets': {}, 'collisions': {}, 'status_msg': "OK"}
                    
                    doc_entry = self.assembled_docs[des]
                    
                    # Фиксируем общее количество листов, если оно встретилось на титульном листе
                    if total_sheets and not doc_entry['expected_total']:
                        doc_entry['expected_total'] = total_sheets

                    if sheet_num is None:
                        # Страницы с нечитаемым номером листа — отдельно, не перезаписывают sheets[1..N]
                        if 'unrecognized' not in doc_entry:
                            doc_entry['unrecognized'] = []
                        doc_entry['unrecognized'].append(page_info)
                    else:
                        # ПРОВЕРКА КОЛЛИЗИИ: Если этот номер листа уже занят другой физической страницей
                        if sheet_num in doc_entry['sheets']:
                            if sheet_num not in doc_entry['collisions']:
                                doc_entry['collisions'][sheet_num] = [doc_entry['sheets'][sheet_num]]
                            doc_entry['collisions'][sheet_num].append(page_info)
                        else:
                            doc_entry['sheets'][sheet_num] = page_info
                else:
                    self.suspect_pages.append(page_info)

        # Выполняем пост-анализ целостности документов
        self._validate_integrity()

    def _validate_integrity(self):
        """Проверяет каждый собранный документ на полноту состава и наличие коллизий."""
        for des, doc_entry in self.assembled_docs.items():
            errors = []
            sheets = doc_entry['sheets']
            expected = doc_entry['expected_total']

            if doc_entry.get('collisions'):
                errors.append(f"Коллизия: дублирование листов {list(doc_entry['collisions'].keys())}")

            unrecognized = doc_entry.get('unrecognized', [])
            if unrecognized:
                errors.append(f"Есть {len(unrecognized)} листа с нечитаемым номером страницы")

            if expected:
                # Проверяем наличие всех листов от 1 до expected_total
                missing_sheets = [idx for idx in range(1, expected + 1) if idx not in sheets]
                if missing_sheets:
                    errors.append(f"Пропуск листов: потеряны {missing_sheets} из {expected}")
                
                # Проверяем избыточность (если физически листов больше, чем заявлено в штампе)
                extra_sheets = [idx for idx in sheets if idx > expected]
                if extra_sheets:
                    errors.append(f"Избыточность: обнаружены лишние листы {extra_sheets}")
            else:
                errors.append("Не найден первый лист (графа 'Листов' не определена)")

            if errors:
                doc_entry['status_msg'] = "; ".join(errors)
            else:
                doc_entry['status_msg'] = f"Готов к сборке ({len(sheets)} л.)"

    def resolve_collision(self, designation, sheet_num, keep_page_idx, keep_page_path):
        """Позволяет UI-контроллеру назначить правильную страницу для листа, разрешив коллизию."""
        doc_entry = self.assembled_docs.get(designation)
        if not doc_entry or sheet_num not in doc_entry.get('collisions', {}):
            return
            
        all_candidates = doc_entry['collisions'][sheet_num]
        chosen_page = next((p for p in all_candidates if p['path'] == keep_page_path and p['idx'] == keep_page_idx), None)
        
        if chosen_page:
            # Переносим выбранную страницу в основной состав
            doc_entry['sheets'][sheet_num] = chosen_page
            # Остальных кандидатов переводим в разряд сомнительных (для ручной привязки к другим КД)
            for p in all_candidates:
                if p != chosen_page:
                    self.suspect_pages.append(p)
            
            del doc_entry['collisions'][sheet_num]
            self._validate_integrity()

    def physical_compile(self, designation, output_path):
        """Физическая склейка страниц документа в один PDF с применением углов поворота."""
        doc_entry = self.assembled_docs.get(designation)
        if not doc_entry:
            logger.warning(f"physical_compile: документ '{designation}' не найден")
            return False, "Документ не найден в результатах анализа"
        
        sheets = doc_entry.get('sheets', {})
        if not sheets:
            logger.warning(f"physical_compile: нет страниц для '{designation}'")
            return False, "Нет страниц для сборки"
        
        # Сортируем по номеру листа; -1 и нечисловые — в конец
        sorted_keys = sorted(sheets.keys(), key=lambda k: (k < 0, k if isinstance(k, int) else 0))
        
        try:
            out = fitz.open()
            for sheet_num in sorted_keys:
                page_info = sheets[sheet_num]
                src = fitz.open(page_info['path'])
                page = src[page_info['idx']]
                inserted = out.new_page(width=page.rect.width, height=page.rect.height)
                
                # Применяем поворот OCR-детекта
                angle = page_info.get('angle', 0)
                if angle:
                    inserted.show_pdf_page(inserted.rect, src, page_info['idx'], rotate=int(angle))
                else:
                    inserted.show_pdf_page(inserted.rect, src, page_info['idx'])
                src.close()
            
            out.save(output_path, garbage=4, deflate=True)
            out.close()
            logger.info(f"Сборка завершена: {designation} -> {output_path} ({len(sheets)} страниц)")
            return True, ""
        except Exception as e:
            logger.error(f"Ошибка сборки '{designation}': {e}")
            return False, str(e)