MinerU/zhch/universal_doc_parser/core/pipeline_manager.py

from typing import Dict, List, Any, Optional, Union
from pathlib import Path
import numpy as np
from PIL import Image
import fitz  # PyMuPDF
from loguru import logger

from .model_factory import ModelFactory
from .config_manager import ConfigManager
from models.adapters import BaseAdapter

class FinancialDocPipeline:
    """金融文档处理统一流水线"""
    
    def __init__(self, config_path: str):
        self.config = ConfigManager.load_config(config_path)
        self.scene_name = self.config.get('scene_name', 'unknown')
        
        # 初始化各个组件
        self._init_components()
        
    def _init_components(self):
        """初始化处理组件"""
        try:
            # 1. 预处理器（方向分类、图像矫正等）
            self.preprocessor = ModelFactory.create_preprocessor(
                self.config['preprocessor']
            )
            
            # 2. 版式检测器
            self.layout_detector = ModelFactory.create_layout_detector(
                self.config['layout_detection']
            )
            
            # 3. VL识别器（表格、公式等）
            self.vl_recognizer = ModelFactory.create_vl_recognizer(
                self.config['vl_recognition']
            )
            
            # 4. OCR识别器
            self.ocr_recognizer = ModelFactory.create_ocr_recognizer(
                self.config['ocr_recognition']
            )
            
            logger.info(f"✅ Initialized pipeline for scene: {self.scene_name}")
            
        except Exception as e:
            logger.error(f"❌ Failed to initialize pipeline components: {e}")
            raise
    
    def process_document(self, document_path: str) -> Dict[str, Any]:
        """
        处理文档的主流程
        
        Args:
            document_path: 文档路径
            
        Returns:
            处理结果，包含所有元素的坐标和内容信息
        """
        results = {
            'scene': self.scene_name,
            'document_path': document_path,
            'pages': [],
            'metadata': self._extract_metadata(document_path)
        }
        
        try:
            # 加载文档图像
            images = self._load_document_images(document_path)
            logger.info(f"📄 Loaded {len(images)} pages from document")
            
            for page_idx, image in enumerate(images):
                logger.info(f"🔍 Processing page {page_idx + 1}/{len(images)}")
                page_result = self._process_single_page(image, page_idx)
                results['pages'].append(page_result)
                
        except Exception as e:
            logger.error(f"❌ Failed to process document: {e}")
            raise
            
        return results
    
    def _load_document_images(self, document_path: str) -> List[np.ndarray]:
        """加载文档图像"""
        document_path = Path(document_path)
        
        if not document_path.exists():
            raise FileNotFoundError(f"Document not found: {document_path}")
        
        images = []
        
        if document_path.suffix.lower() == '.pdf':
            # 处理PDF文件
            doc = fitz.open(document_path)
            try:
                for page_num in range(len(doc)):
                    page = doc.load_page(page_num)
                    # 设置合适的DPI
                    dpi = self.config.get('input', {}).get('dpi', 200)
                    mat = fitz.Matrix(dpi/72, dpi/72)
                    pix = page.get_pixmap(matrix=mat)
                    img_data = pix.tobytes("ppm")
                    
                    # 转换为numpy数组
                    from io import BytesIO
                    img = Image.open(BytesIO(img_data))
                    img_array = np.array(img)
                    images.append(img_array)
            finally:
                doc.close()
                
        elif document_path.suffix.lower() in ['.png', '.jpg', '.jpeg', '.bmp', '.tiff']:
            # 处理图像文件
            img = Image.open(document_path)
            if img.mode != 'RGB':
                img = img.convert('RGB')
            img_array = np.array(img)
            images.append(img_array)
            
        else:
            raise ValueError(f"Unsupported file format: {document_path.suffix}")
        
        return images
    
    def _extract_metadata(self, document_path: str) -> Dict[str, Any]:
        """提取文档元数据"""
        document_path = Path(document_path)
        
        metadata = {
            'filename': document_path.name,
            'size': document_path.stat().st_size,
            'format': document_path.suffix.lower()
        }
        
        # 如果是PDF，提取更多元数据
        if document_path.suffix.lower() == '.pdf':
            try:
                doc = fitz.open(document_path)
                metadata.update({
                    'page_count': len(doc),
                    'title': doc.metadata.get('title', ''),
                    'author': doc.metadata.get('author', ''),
                    'subject': doc.metadata.get('subject', ''),
                    'creator': doc.metadata.get('creator', '')
                })
                doc.close()
            except Exception:
                pass
        
        return metadata
    
    def _process_single_page(self, image: np.ndarray, page_idx: int) -> Dict[str, Any]:
        """处理单页文档"""
        # 1. 预处理（方向校正等）
        try:
            preprocessed_image, rotate_angle = self.preprocessor.process(image)
        except Exception as e:
            logger.warning(f"⚠️ Preprocessing failed for page {page_idx}: {e}")
            preprocessed_image = image
        
        # 2. 版式检测
        try:
            layout_results = self.layout_detector.detect(preprocessed_image)
            logger.info(f"📋 Detected {len(layout_results)} layout elements on page {page_idx}")
        except Exception as e:
            logger.error(f"❌ Layout detection failed for page {page_idx}: {e}")
            layout_results = []
        
        # 3. 根据场景类型分别处理不同元素
        page_elements = []
        
        for layout_item in layout_results:
            try:
                element_type = layout_item['category']
                
                if element_type in ['table_body', 'table']:
                    # 表格使用VL模型处理
                    element_result = self._process_table_element(
                        preprocessed_image, layout_item
                    )
                elif element_type in ['text', 'title', 'ocr_text']:
                    # 文本使用OCR处理
                    element_result = self._process_text_element(
                        preprocessed_image, layout_item
                    )
                elif element_type in ['interline_equation', 'inline_equation']:
                    # 公式使用VL模型处理
                    element_result = self._process_formula_element(
                        preprocessed_image, layout_item
                    )
                else:
                    # 其他元素保持原样
                    element_result = layout_item.copy()
                    element_result['type'] = element_type
                    
                page_elements.append(element_result)
                
            except Exception as e:
                logger.warning(f"⚠️ Failed to process element {element_type}: {e}")
                # 添加失败的元素，标记为错误
                error_element = layout_item.copy()
                error_element['type'] = 'error'
                error_element['error'] = str(e)
                page_elements.append(error_element)
        
        return {
            'page_idx': page_idx,
            'elements': page_elements,
            'layout_raw': layout_results,
            'image_shape': preprocessed_image.shape,
            'processed_image': preprocessed_image,
            'angle': rotate_angle
        }
    
    def _process_table_element(self, image: np.ndarray, layout_item: Dict[str, Any]) -> Dict[str, Any]:
        """处理表格元素"""
        try:
            # 裁剪表格区域
            cropped_table = self._crop_region(image, layout_item['bbox'])
            
            # 使用VL模型识别表格
            table_result = self.vl_recognizer.recognize_table(
                cropped_table,
                return_cells_coordinate=True  # 关键：返回单元格坐标
            )
            
            # 转换坐标到原图坐标系
            if 'cells' in table_result:
                for cell in table_result['cells']:
                    cell['absolute_bbox'] = self._convert_to_absolute_coords(
                        cell['bbox'], layout_item['bbox']
                    )
            
            result = {
                'type': 'table',
                'bbox': layout_item['bbox'],
                'confidence': layout_item.get('confidence', 0.0),
                'content': table_result,
                'scene_specific': self._add_scene_specific_info(table_result)
            }
            
            logger.info(f"✅ Table processed with {len(table_result.get('cells', []))} cells")
            return result
            
        except Exception as e:
            logger.error(f"❌ Table processing failed: {e}")
            return {
                'type': 'table',
                'bbox': layout_item['bbox'],
                'content': {'html': '', 'markdown': '', 'cells': []},
                'error': str(e)
            }
    
    def _process_text_element(self, image: np.ndarray, layout_item: Dict[str, Any]) -> Dict[str, Any]:
        """处理文本元素"""
        try:
            # 裁剪文本区域
            cropped_text = self._crop_region(image, layout_item['bbox'])
            
            # 使用OCR识别文本
            text_results = self.ocr_recognizer.recognize_text(cropped_text)
            
            # 合并识别结果
            combined_text = ""
            if text_results:
                text_parts = [item['text'] for item in text_results if item['confidence'] > 0.5]
                combined_text = " ".join(text_parts)
            
            result = {
                'type': layout_item['category'],
                'bbox': layout_item['bbox'],
                'confidence': layout_item.get('confidence', 0.0),
                'content': {
                    'text': combined_text,
                    'ocr_details': text_results
                }
            }
            
            logger.info(f"✅ Text processed: '{combined_text[:50]}...'")
            return result
            
        except Exception as e:
            logger.error(f"❌ Text processing failed: {e}")
            return {
                'type': layout_item['category'],
                'bbox': layout_item['bbox'],
                'content': {'text': '', 'ocr_details': []},
                'error': str(e)
            }
    
    def _process_formula_element(self, image: np.ndarray, layout_item: Dict[str, Any]) -> Dict[str, Any]:
        """处理公式元素"""
        try:
            # 裁剪公式区域
            cropped_formula = self._crop_region(image, layout_item['bbox'])
            
            # 使用VL模型识别公式
            formula_result = self.vl_recognizer.recognize_formula(cropped_formula)
            
            result = {
                'type': 'formula',
                'bbox': layout_item['bbox'],
                'confidence': layout_item.get('confidence', 0.0),
                'content': formula_result
            }
            
            logger.info(f"✅ Formula processed: {formula_result.get('latex', '')[:50]}...")
            return result
            
        except Exception as e:
            logger.error(f"❌ Formula processing failed: {e}")
            return {
                'type': 'formula',
                'bbox': layout_item['bbox'],
                'content': {'latex': '', 'confidence': 0.0},
                'error': str(e)
            }
    
    def _crop_region(self, image: np.ndarray, bbox: List[float]) -> np.ndarray:
        """裁剪图像区域"""
        if len(bbox) < 4:
            return image
            
        x1, y1, x2, y2 = map(int, bbox)
        
        # 边界检查
        h, w = image.shape[:2]
        x1 = max(0, min(x1, w))
        y1 = max(0, min(y1, h))
        x2 = max(x1, min(x2, w))
        y2 = max(y1, min(y2, h))
        
        return image[y1:y2, x1:x2]
    
    def _convert_to_absolute_coords(self, relative_bbox: List[float], region_bbox: List[float]) -> List[float]:
        """将相对坐标转换为绝对坐标"""
        if len(relative_bbox) < 4 or len(region_bbox) < 4:
            return relative_bbox
            
        rx1, ry1, rx2, ry2 = relative_bbox
        bx1, by1, bx2, by2 = region_bbox
        
        # 计算绝对坐标
        abs_x1 = bx1 + rx1
        abs_y1 = by1 + ry1
        abs_x2 = bx1 + rx2
        abs_y2 = by1 + ry2
        
        return [abs_x1, abs_y1, abs_x2, abs_y2]
    
    def _add_scene_specific_info(self, content: Dict[str, Any]) -> Dict[str, Any]:
        """根据场景添加特定信息"""
        if self.scene_name == 'bank_statement':
            return self._process_bank_statement_table(content)
        elif self.scene_name == 'financial_report':
            return self._process_financial_report_table(content)
        return {}
    
    def _process_bank_statement_table(self, content: Dict[str, Any]) -> Dict[str, Any]:
        """处理银行流水表格特定逻辑"""
        scene_info = {
            'table_type': 'bank_statement',
            'expected_columns': ['日期', '摘要', '收入', '支出', '余额'],
            'validation_rules': {
                'amount_format': True,
                'date_format': True,
                'balance_consistency': True
            }
        }
        
        # 进行银行流水特定的验证和处理
        if 'html' in content and content['html']:
            # 这里可以添加银行流水特定的HTML后处理逻辑
            pass
            
        return scene_info
    
    def _process_financial_report_table(self, content: Dict[str, Any]) -> Dict[str, Any]:
        """处理财务报表特定逻辑"""
        scene_info = {
            'table_type': 'financial_report',
            'complex_headers': True,
            'merged_cells': True,
            'validation_rules': {
                'accounting_format': True,
                'sum_validation': True
            }
        }
        
        # 进行财务报表特定的验证和处理
        if 'html' in content and content['html']:
            # 这里可以添加财务报表特定的HTML后处理逻辑
            pass
            
        return scene_info
    
    def cleanup(self):
        """清理资源"""
        try:
            if hasattr(self, 'preprocessor'):
                self.preprocessor.cleanup()
            if hasattr(self, 'layout_detector'):
                self.layout_detector.cleanup()
            if hasattr(self, 'vl_recognizer'):
                self.vl_recognizer.cleanup()
            if hasattr(self, 'ocr_recognizer'):
                self.ocr_recognizer.cleanup()
                
            logger.info("✅ Pipeline cleanup completed")
            
        except Exception as e:
            logger.warning(f"⚠️ Cleanup failed: {e}")
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.cleanup()