|
|
@@ -0,0 +1,1068 @@
|
|
|
+"""
|
|
|
+增强版文档处理流水线 v2
|
|
|
+按照优化后的流程实现:
|
|
|
+1. PDF分类 → 扫描件/数字原生PDF
|
|
|
+2. 页面方向识别(仅扫描件)
|
|
|
+3. Layout检测
|
|
|
+4. 并行处理:
|
|
|
+ - 文本区域:OCR检测+识别 / PDF字符提取
|
|
|
+ - 表格区域:OCR检测(坐标) + VLM结构识别 → 坐标匹配
|
|
|
+5. 合并结果 → 跨页表格合并
|
|
|
+"""
|
|
|
+import os
|
|
|
+import sys
|
|
|
+from typing import Dict, List, Any, Optional, Union, Tuple
|
|
|
+from pathlib import Path
|
|
|
+from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
|
+import numpy as np
|
|
|
+from PIL import Image
|
|
|
+import cv2
|
|
|
+import fitz # PyMuPDF
|
|
|
+from loguru import logger
|
|
|
+from io import BytesIO
|
|
|
+
|
|
|
+# 添加项目路径
|
|
|
+project_root = Path(__file__).parents[3]
|
|
|
+if str(project_root) not in sys.path:
|
|
|
+ sys.path.insert(0, str(project_root))
|
|
|
+
|
|
|
+module_root = Path(__file__).parents[1]
|
|
|
+if str(module_root) not in sys.path:
|
|
|
+ sys.path.insert(0, str(module_root))
|
|
|
+
|
|
|
+try:
|
|
|
+ from .model_factory import ModelFactory
|
|
|
+ from .config_manager import ConfigManager
|
|
|
+except ImportError as e:
|
|
|
+ from model_factory import ModelFactory
|
|
|
+ from config_manager import ConfigManager
|
|
|
+
|
|
|
+# 导入 MinerU 组件
|
|
|
+try:
|
|
|
+ from mineru.utils.pdf_classify import classify as pdf_classify
|
|
|
+ from mineru.utils.pdf_image_tools import load_images_from_pdf
|
|
|
+ from mineru.utils.enum_class import ImageType
|
|
|
+ MINERU_AVAILABLE = True
|
|
|
+except ImportError as e:
|
|
|
+ logger.warning(f"MinerU components not available: {e}")
|
|
|
+ MINERU_AVAILABLE = False
|
|
|
+
|
|
|
+# 导入 merger 组件用于坐标匹配
|
|
|
+try:
|
|
|
+ from merger import TableCellMatcher
|
|
|
+ from merger import TextMatcher
|
|
|
+ from merger import BBoxExtractor
|
|
|
+ from merger import DataProcessor
|
|
|
+ MERGER_AVAILABLE = True
|
|
|
+except ImportError as e:
|
|
|
+ logger.warning(f"Merger components not available: {e}")
|
|
|
+ MERGER_AVAILABLE = False
|
|
|
+
|
|
|
+
|
|
|
+class EnhancedDocPipeline:
|
|
|
+ """增强版文档处理流水线"""
|
|
|
+
|
|
|
+ def __init__(self, config_path: str):
|
|
|
+ """
|
|
|
+ 初始化流水线
|
|
|
+
|
|
|
+ Args:
|
|
|
+ config_path: 配置文件路径
|
|
|
+ """
|
|
|
+ self.config = ConfigManager.load_config(config_path)
|
|
|
+ self.scene_name = self.config.get('scene_name', 'unknown')
|
|
|
+ self.debug_mode = self.config.get('output', {}).get('debug_mode', False)
|
|
|
+
|
|
|
+ # 初始化组件
|
|
|
+ self._init_components()
|
|
|
+
|
|
|
+ # 初始化 merger 组件
|
|
|
+ if MERGER_AVAILABLE:
|
|
|
+ self.text_matcher = TextMatcher()
|
|
|
+ self.table_cell_matcher = TableCellMatcher(
|
|
|
+ text_matcher=self.text_matcher,
|
|
|
+ x_tolerance=3,
|
|
|
+ y_tolerance=10
|
|
|
+ )
|
|
|
+ else:
|
|
|
+ self.text_matcher = None
|
|
|
+ self.table_cell_matcher = None
|
|
|
+ logger.warning("⚠️ Merger components not available, cell coordinate matching disabled")
|
|
|
+
|
|
|
+ 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"✅ Pipeline initialized for scene: {self.scene_name}")
|
|
|
+
|
|
|
+ except Exception as e:
|
|
|
+ logger.error(f"❌ Failed to initialize pipeline: {e}")
|
|
|
+ raise
|
|
|
+
|
|
|
+ def process_document(self, document_path: str) -> Dict[str, Any]:
|
|
|
+ """
|
|
|
+ 处理文档主流程
|
|
|
+
|
|
|
+ Args:
|
|
|
+ document_path: 文档路径(PDF、图片或目录)
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 处理结果字典
|
|
|
+ """
|
|
|
+ document_path = Path(document_path)
|
|
|
+
|
|
|
+ results = {
|
|
|
+ 'scene': self.scene_name,
|
|
|
+ 'document_path': str(document_path),
|
|
|
+ 'pages': [],
|
|
|
+ 'metadata': {},
|
|
|
+ 'debug_info': {} if self.debug_mode else None
|
|
|
+ }
|
|
|
+
|
|
|
+ try:
|
|
|
+ # 1. 加载文档并分类
|
|
|
+ images, pdf_type, pdf_doc = self._load_and_classify_document(document_path)
|
|
|
+ results['metadata']['pdf_type'] = pdf_type
|
|
|
+ results['metadata']['page_count'] = len(images)
|
|
|
+
|
|
|
+ logger.info(f"📄 Loaded {len(images)} pages, type: {pdf_type}")
|
|
|
+
|
|
|
+ # 2. 处理每一页
|
|
|
+ for page_idx, image_dict in enumerate(images):
|
|
|
+ logger.info(f"🔍 Processing page {page_idx + 1}/{len(images)}")
|
|
|
+
|
|
|
+ page_result = self._process_single_page(
|
|
|
+ image_dict=image_dict,
|
|
|
+ page_idx=page_idx,
|
|
|
+ pdf_type=pdf_type,
|
|
|
+ pdf_doc=pdf_doc
|
|
|
+ )
|
|
|
+ results['pages'].append(page_result)
|
|
|
+
|
|
|
+ # 3. 跨页表格合并
|
|
|
+ results = self._merge_cross_page_tables(results)
|
|
|
+
|
|
|
+ # 4. 关闭 PDF 文档
|
|
|
+ if pdf_doc is not None:
|
|
|
+ pdf_doc.close()
|
|
|
+
|
|
|
+ logger.info(f"✅ Document processing completed")
|
|
|
+ return results
|
|
|
+
|
|
|
+ except Exception as e:
|
|
|
+ logger.error(f"❌ Document processing failed: {e}")
|
|
|
+ raise
|
|
|
+
|
|
|
+ def _load_and_classify_document(
|
|
|
+ self,
|
|
|
+ document_path: Path
|
|
|
+ ) -> Tuple[List[Dict], str, Optional[Any]]:
|
|
|
+ """
|
|
|
+ 加载文档并分类
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ (images_list, pdf_type, pdf_doc)
|
|
|
+ - images_list: 图像列表,每个元素包含 {'img_pil': PIL.Image, 'scale': float}
|
|
|
+ - pdf_type: 'ocr' 或 'txt'
|
|
|
+ - pdf_doc: PDF文档对象(如果是PDF)
|
|
|
+ """
|
|
|
+ pdf_doc = None
|
|
|
+ pdf_type = 'ocr' # 默认使用OCR模式
|
|
|
+ images = []
|
|
|
+
|
|
|
+ if document_path.is_dir():
|
|
|
+ # 处理目录:遍历所有图片
|
|
|
+ image_extensions = {'.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif'}
|
|
|
+ image_files = sorted([
|
|
|
+ f for f in document_path.iterdir()
|
|
|
+ if f.suffix.lower() in image_extensions
|
|
|
+ ])
|
|
|
+
|
|
|
+ for img_file in image_files:
|
|
|
+ img = Image.open(img_file)
|
|
|
+ if img.mode != 'RGB':
|
|
|
+ img = img.convert('RGB')
|
|
|
+ images.append({
|
|
|
+ 'img_pil': img,
|
|
|
+ 'scale': 1.0,
|
|
|
+ 'source_path': str(img_file)
|
|
|
+ })
|
|
|
+
|
|
|
+ pdf_type = 'ocr' # 图片目录始终使用OCR模式
|
|
|
+
|
|
|
+ elif document_path.suffix.lower() == '.pdf':
|
|
|
+ # 处理PDF文件
|
|
|
+ with open(document_path, 'rb') as f:
|
|
|
+ pdf_bytes = f.read()
|
|
|
+
|
|
|
+ # PDF分类
|
|
|
+ if MINERU_AVAILABLE:
|
|
|
+ pdf_type = pdf_classify(pdf_bytes)
|
|
|
+ logger.info(f"📋 PDF classified as: {pdf_type}")
|
|
|
+
|
|
|
+ # 加载图像
|
|
|
+ dpi = self.config.get('input', {}).get('dpi', 200)
|
|
|
+ images_list, pdf_doc = load_images_from_pdf(
|
|
|
+ pdf_bytes,
|
|
|
+ dpi=dpi,
|
|
|
+ image_type=ImageType.PIL
|
|
|
+ )
|
|
|
+
|
|
|
+ for img_dict in images_list:
|
|
|
+ images.append({
|
|
|
+ 'img_pil': img_dict['img_pil'],
|
|
|
+ 'scale': img_dict.get('scale', dpi / 72),
|
|
|
+ 'source_path': str(document_path)
|
|
|
+ })
|
|
|
+
|
|
|
+ elif document_path.suffix.lower() in ['.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif']:
|
|
|
+ # 处理单个图片
|
|
|
+ img = Image.open(document_path)
|
|
|
+ if img.mode != 'RGB':
|
|
|
+ img = img.convert('RGB')
|
|
|
+ images.append({
|
|
|
+ 'img_pil': img,
|
|
|
+ 'scale': 1.0,
|
|
|
+ 'source_path': str(document_path)
|
|
|
+ })
|
|
|
+ pdf_type = 'ocr'
|
|
|
+
|
|
|
+ else:
|
|
|
+ raise ValueError(f"Unsupported file format: {document_path.suffix}")
|
|
|
+
|
|
|
+ return images, pdf_type, pdf_doc
|
|
|
+
|
|
|
+ def _process_single_page(
|
|
|
+ self,
|
|
|
+ image_dict: Dict[str, Any],
|
|
|
+ page_idx: int,
|
|
|
+ pdf_type: str,
|
|
|
+ pdf_doc: Optional[Any] = None
|
|
|
+ ) -> Dict[str, Any]:
|
|
|
+ """
|
|
|
+ 处理单页文档
|
|
|
+
|
|
|
+ 流程:
|
|
|
+ 1. 扫描件:页面方向识别 → Layout检测
|
|
|
+ 2. 数字PDF:直接Layout检测
|
|
|
+ 3. 并行处理文本区域和表格区域
|
|
|
+ 4. 合并结果
|
|
|
+ 5. **坐标转换回原始图片坐标系**
|
|
|
+
|
|
|
+ 注意:
|
|
|
+ - 输出的图片是原始未旋转的图片
|
|
|
+ - 输出的坐标是基于原始未旋转图片的坐标系
|
|
|
+ """
|
|
|
+ pil_image = image_dict['img_pil']
|
|
|
+ scale = image_dict.get('scale', 1.0)
|
|
|
+ original_image = np.array(pil_image) # 保存原始图片用于输出
|
|
|
+
|
|
|
+ page_result = {
|
|
|
+ 'page_idx': page_idx,
|
|
|
+ 'elements': [],
|
|
|
+ 'image_shape': original_image.shape, # 原始图片尺寸
|
|
|
+ 'original_image': original_image, # 输出用原始图片(未旋转)
|
|
|
+ 'angle': 0,
|
|
|
+ 'scale': scale,
|
|
|
+ 'pdf_type': pdf_type
|
|
|
+ }
|
|
|
+
|
|
|
+ # 用于检测的图片(可能被旋转以提高检测准确率)
|
|
|
+ detection_image = original_image.copy()
|
|
|
+ rotate_angle = 0
|
|
|
+
|
|
|
+ # 1. 页面方向识别(仅扫描件)
|
|
|
+ if pdf_type == 'ocr':
|
|
|
+ try:
|
|
|
+ detection_image, rotate_angle = self.preprocessor.process(original_image)
|
|
|
+ page_result['angle'] = rotate_angle
|
|
|
+
|
|
|
+ if rotate_angle != 0:
|
|
|
+ logger.info(f"📐 Page {page_idx}: rotated {rotate_angle}° for detection")
|
|
|
+ except Exception as e:
|
|
|
+ logger.warning(f"⚠️ Orientation detection failed: {e}")
|
|
|
+
|
|
|
+ # 2. Layout检测(在旋转后的图片上进行以提高准确率)
|
|
|
+ try:
|
|
|
+ layout_results = self.layout_detector.detect(detection_image)
|
|
|
+ logger.info(f"📋 Page {page_idx}: detected {len(layout_results)} elements")
|
|
|
+ except Exception as e:
|
|
|
+ logger.error(f"❌ Layout detection failed: {e}")
|
|
|
+ layout_results = []
|
|
|
+
|
|
|
+ page_result['layout_raw'] = layout_results
|
|
|
+
|
|
|
+ # 3. 分类元素
|
|
|
+ text_elements = []
|
|
|
+ table_elements = []
|
|
|
+ other_elements = []
|
|
|
+
|
|
|
+ for item in layout_results:
|
|
|
+ category = item.get('category', '')
|
|
|
+ if category in ['text', 'title', 'ocr_text', 'ref_text', 'header', 'footer',
|
|
|
+ 'image_caption', 'image_footnote', 'table_caption', 'table_footnote']:
|
|
|
+ text_elements.append(item)
|
|
|
+ elif category in ['table', 'table_body']:
|
|
|
+ table_elements.append(item)
|
|
|
+ else:
|
|
|
+ other_elements.append(item)
|
|
|
+
|
|
|
+ # 4. 并行处理文本和表格(在旋转后的图片上进行识别)
|
|
|
+ processed_elements = []
|
|
|
+
|
|
|
+ # 4.1 处理文本区域
|
|
|
+ for text_item in text_elements:
|
|
|
+ try:
|
|
|
+ element = self._process_text_element(
|
|
|
+ detection_image, text_item, pdf_type, pdf_doc, page_idx, scale
|
|
|
+ )
|
|
|
+ # **关键**: 将坐标转换回原始图片坐标系
|
|
|
+ if rotate_angle != 0:
|
|
|
+ element = self._transform_coords_to_original(
|
|
|
+ element, rotate_angle, detection_image.shape, original_image.shape
|
|
|
+ )
|
|
|
+ processed_elements.append(element)
|
|
|
+ except Exception as e:
|
|
|
+ logger.warning(f"⚠️ Text processing failed: {e}")
|
|
|
+ processed_elements.append(self._create_error_element(text_item, str(e)))
|
|
|
+
|
|
|
+ # 4.2 处理表格区域(OCR检测 + VLM结构识别 + 坐标匹配)
|
|
|
+ for table_item in table_elements:
|
|
|
+ try:
|
|
|
+ element = self._process_table_element(
|
|
|
+ detection_image, table_item, scale
|
|
|
+ )
|
|
|
+ # **关键**: 将坐标转换回原始图片坐标系
|
|
|
+ if rotate_angle != 0:
|
|
|
+ element = self._transform_coords_to_original(
|
|
|
+ element, rotate_angle, detection_image.shape, original_image.shape
|
|
|
+ )
|
|
|
+ processed_elements.append(element)
|
|
|
+ except Exception as e:
|
|
|
+ logger.warning(f"⚠️ Table processing failed: {e}")
|
|
|
+ processed_elements.append(self._create_error_element(table_item, str(e)))
|
|
|
+
|
|
|
+ # 4.3 处理其他元素
|
|
|
+ for other_item in other_elements:
|
|
|
+ try:
|
|
|
+ element = self._process_other_element(detection_image, other_item)
|
|
|
+ # **关键**: 将坐标转换回原始图片坐标系
|
|
|
+ if rotate_angle != 0:
|
|
|
+ element = self._transform_coords_to_original(
|
|
|
+ element, rotate_angle, detection_image.shape, original_image.shape
|
|
|
+ )
|
|
|
+ processed_elements.append(element)
|
|
|
+ except Exception as e:
|
|
|
+ logger.warning(f"⚠️ Element processing failed: {e}")
|
|
|
+ processed_elements.append(self._create_error_element(other_item, str(e)))
|
|
|
+
|
|
|
+ page_result['elements'] = processed_elements
|
|
|
+ return page_result
|
|
|
+
|
|
|
+ def _process_text_element(
|
|
|
+ self,
|
|
|
+ image: np.ndarray,
|
|
|
+ layout_item: Dict[str, Any],
|
|
|
+ pdf_type: str,
|
|
|
+ pdf_doc: Optional[Any],
|
|
|
+ page_idx: int,
|
|
|
+ scale: float
|
|
|
+ ) -> Dict[str, Any]:
|
|
|
+ """
|
|
|
+ 处理文本元素
|
|
|
+
|
|
|
+ - 扫描件:OCR检测+识别
|
|
|
+ - 数字PDF:尝试PDF字符提取,失败则OCR补充
|
|
|
+ """
|
|
|
+ bbox = layout_item.get('bbox', [0, 0, 0, 0])
|
|
|
+ cropped_image = self._crop_region(image, bbox)
|
|
|
+
|
|
|
+ text_content = ""
|
|
|
+ ocr_details = []
|
|
|
+ extraction_method = "ocr"
|
|
|
+
|
|
|
+ # 数字原生PDF:尝试直接提取文本
|
|
|
+ if pdf_type == 'txt' and pdf_doc is not None:
|
|
|
+ try:
|
|
|
+ # 从PDF直接提取文本(坐标更精确)
|
|
|
+ text_content, extraction_success = self._extract_text_from_pdf(
|
|
|
+ pdf_doc, page_idx, bbox, scale
|
|
|
+ )
|
|
|
+ if extraction_success and text_content.strip():
|
|
|
+ extraction_method = "pdf_extract"
|
|
|
+ logger.debug(f"📝 Text extracted from PDF: '{text_content[:30]}...'")
|
|
|
+ except Exception as e:
|
|
|
+ logger.debug(f"PDF text extraction failed: {e}")
|
|
|
+
|
|
|
+ # OCR识别(扫描件或PDF提取失败)
|
|
|
+ if extraction_method == "ocr" or not text_content.strip():
|
|
|
+ try:
|
|
|
+ ocr_results = self.ocr_recognizer.recognize_text(cropped_image)
|
|
|
+ if ocr_results:
|
|
|
+ text_parts = [
|
|
|
+ item['text'] for item in ocr_results
|
|
|
+ if item.get('confidence', 0) > 0.5
|
|
|
+ ]
|
|
|
+ text_content = " ".join(text_parts)
|
|
|
+ ocr_details = ocr_results
|
|
|
+ extraction_method = "ocr"
|
|
|
+ except Exception as e:
|
|
|
+ logger.warning(f"OCR recognition failed: {e}")
|
|
|
+
|
|
|
+ return {
|
|
|
+ 'type': layout_item.get('category', 'text'),
|
|
|
+ 'bbox': bbox,
|
|
|
+ 'confidence': layout_item.get('confidence', 0.0),
|
|
|
+ 'content': {
|
|
|
+ 'text': text_content,
|
|
|
+ 'ocr_details': ocr_details,
|
|
|
+ 'extraction_method': extraction_method
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ def _process_table_element(
|
|
|
+ self,
|
|
|
+ image: np.ndarray,
|
|
|
+ layout_item: Dict[str, Any],
|
|
|
+ scale: float
|
|
|
+ ) -> Dict[str, Any]:
|
|
|
+ """
|
|
|
+ 处理表格元素
|
|
|
+
|
|
|
+ 流程:
|
|
|
+ 1. 表格方向检测
|
|
|
+ 2. OCR检测获取文本框坐标(旋转后的图片)
|
|
|
+ 3. VLM识别获取表格结构HTML
|
|
|
+ 4. 匹配OCR坐标与VLM结构
|
|
|
+ 5. 将坐标逆向转换回原图坐标(参考 merger.DataProcessor)
|
|
|
+
|
|
|
+ 关键点:
|
|
|
+ - OCR 在旋转后的图片上进行,坐标是旋转后的
|
|
|
+ - 匹配完成后,需要将坐标逆向转换回原图坐标
|
|
|
+ """
|
|
|
+ bbox = layout_item.get('bbox', [0, 0, 0, 0])
|
|
|
+ cropped_table = self._crop_region(image, bbox)
|
|
|
+
|
|
|
+ # 获取裁剪后表格图片的尺寸(原始尺寸,用于坐标逆转换)
|
|
|
+ orig_table_h, orig_table_w = cropped_table.shape[:2]
|
|
|
+ orig_table_size = (orig_table_w, orig_table_h) # (width, height)
|
|
|
+
|
|
|
+ # 1. 表格方向检测(使用预处理器)
|
|
|
+ table_angle = 0
|
|
|
+ try:
|
|
|
+ rotated_table, table_angle = self.preprocessor.process(cropped_table)
|
|
|
+ if table_angle != 0:
|
|
|
+ logger.info(f"📐 Table rotated {table_angle}°")
|
|
|
+ cropped_table = rotated_table
|
|
|
+ except Exception as e:
|
|
|
+ logger.debug(f"Table orientation detection skipped: {e}")
|
|
|
+ table_angle = 0
|
|
|
+
|
|
|
+ # 2. OCR检测获取文本框坐标(在旋转后的图片上)
|
|
|
+ ocr_boxes = []
|
|
|
+ try:
|
|
|
+ ocr_results = self.ocr_recognizer.recognize_text(cropped_table)
|
|
|
+ if ocr_results:
|
|
|
+ for idx, item in enumerate(ocr_results):
|
|
|
+ ocr_poly = item.get('bbox', []) # OCR返回的是4点多边形
|
|
|
+ if ocr_poly:
|
|
|
+ # 转换为 TableCellMatcher 期望的格式
|
|
|
+ # 注意:这里不加表格偏移量,保持相对于裁剪表格的坐标
|
|
|
+ formatted_box = self._convert_ocr_to_matcher_format(
|
|
|
+ ocr_poly,
|
|
|
+ item.get('text', ''),
|
|
|
+ item.get('confidence', 0.0),
|
|
|
+ idx,
|
|
|
+ table_bbox=None # 不加偏移量
|
|
|
+ )
|
|
|
+ if formatted_box:
|
|
|
+ ocr_boxes.append(formatted_box)
|
|
|
+ logger.info(f"📊 OCR detected {len(ocr_boxes)} text boxes in table")
|
|
|
+ except Exception as e:
|
|
|
+ logger.warning(f"Table OCR detection failed: {e}")
|
|
|
+
|
|
|
+ # 3. VLM识别获取表格结构HTML
|
|
|
+ table_html = ""
|
|
|
+ try:
|
|
|
+ vl_result = self.vl_recognizer.recognize_table(
|
|
|
+ cropped_table,
|
|
|
+ return_cells_coordinate=True
|
|
|
+ )
|
|
|
+ table_html = vl_result.get('html', '')
|
|
|
+ logger.info(f"📊 VLM recognized table structure")
|
|
|
+ except Exception as e:
|
|
|
+ logger.warning(f"VLM table recognition failed: {e}")
|
|
|
+
|
|
|
+ # 4. 匹配OCR坐标与VLM结构
|
|
|
+ cells = []
|
|
|
+ enhanced_html = table_html
|
|
|
+ skew_angle = 0.0
|
|
|
+
|
|
|
+ if table_html and ocr_boxes and self.table_cell_matcher:
|
|
|
+ try:
|
|
|
+ # OCR 坐标是相对于旋转后的表格图片的
|
|
|
+ # TableCellMatcher 返回的 cells 坐标也是旋转后的
|
|
|
+ enhanced_html, cells, _, skew_angle = self.table_cell_matcher.enhance_table_html_with_bbox(
|
|
|
+ html=table_html,
|
|
|
+ paddle_text_boxes=ocr_boxes,
|
|
|
+ start_pointer=0,
|
|
|
+ table_bbox=None # 不使用 table_bbox 筛选
|
|
|
+ )
|
|
|
+ logger.info(f"📊 Matched {len(cells)} cells with coordinates (skew: {skew_angle:.2f}°)")
|
|
|
+ except Exception as e:
|
|
|
+ logger.warning(f"Cell coordinate matching failed: {e}")
|
|
|
+
|
|
|
+ # 5. 坐标转换:将旋转后的坐标转换回原图坐标
|
|
|
+ if table_angle != 0 and MERGER_AVAILABLE:
|
|
|
+ # 使用 BBoxExtractor 进行逆向旋转
|
|
|
+ cells, enhanced_html = self._inverse_rotate_table_coords(
|
|
|
+ cells=cells,
|
|
|
+ html=enhanced_html,
|
|
|
+ rotation_angle=table_angle,
|
|
|
+ orig_table_size=orig_table_size,
|
|
|
+ table_bbox=bbox
|
|
|
+ )
|
|
|
+ logger.info(f"📐 Coordinates transformed back to original image")
|
|
|
+ else:
|
|
|
+ # 没有旋转,只需要加上表格偏移量转换为整页坐标
|
|
|
+ cells = self._add_table_offset_to_cells(cells, bbox)
|
|
|
+ enhanced_html = self._add_table_offset_to_html(enhanced_html, bbox)
|
|
|
+
|
|
|
+ return {
|
|
|
+ 'type': 'table',
|
|
|
+ 'bbox': bbox,
|
|
|
+ 'confidence': layout_item.get('confidence', 0.0),
|
|
|
+ 'content': {
|
|
|
+ 'html': enhanced_html,
|
|
|
+ 'original_html': table_html,
|
|
|
+ 'cells': cells,
|
|
|
+ 'ocr_boxes': ocr_boxes,
|
|
|
+ 'table_angle': table_angle,
|
|
|
+ 'skew_angle': skew_angle
|
|
|
+ },
|
|
|
+ }
|
|
|
+
|
|
|
+ def _process_other_element(
|
|
|
+ self,
|
|
|
+ image: np.ndarray,
|
|
|
+ layout_item: Dict[str, Any]
|
|
|
+ ) -> Dict[str, Any]:
|
|
|
+ """处理其他类型元素(公式、图片等)"""
|
|
|
+ bbox = layout_item.get('bbox', [0, 0, 0, 0])
|
|
|
+ category = layout_item.get('category', '')
|
|
|
+ cropped_region = self._crop_region(image, bbox)
|
|
|
+
|
|
|
+ content = {}
|
|
|
+
|
|
|
+ # 公式识别
|
|
|
+ if category in ['interline_equation', 'inline_equation', 'equation']:
|
|
|
+ try:
|
|
|
+ formula_result = self.vl_recognizer.recognize_formula(cropped_region)
|
|
|
+ content = {
|
|
|
+ 'latex': formula_result.get('latex', ''),
|
|
|
+ 'confidence': formula_result.get('confidence', 0.0)
|
|
|
+ }
|
|
|
+ except Exception as e:
|
|
|
+ logger.warning(f"Formula recognition failed: {e}")
|
|
|
+ content = {'latex': '', 'confidence': 0.0}
|
|
|
+
|
|
|
+ # 图片保持原样
|
|
|
+ elif category in ['image', 'image_body', 'figure']:
|
|
|
+ content = {
|
|
|
+ 'type': 'image',
|
|
|
+ 'description': ''
|
|
|
+ }
|
|
|
+
|
|
|
+ return {
|
|
|
+ 'type': category,
|
|
|
+ 'bbox': bbox,
|
|
|
+ 'confidence': layout_item.get('confidence', 0.0),
|
|
|
+ 'content': content
|
|
|
+ }
|
|
|
+
|
|
|
+ def _extract_text_from_pdf(
|
|
|
+ self,
|
|
|
+ pdf_doc: Any,
|
|
|
+ page_idx: int,
|
|
|
+ bbox: List[float],
|
|
|
+ scale: float
|
|
|
+ ) -> Tuple[str, bool]:
|
|
|
+ """
|
|
|
+ 从PDF直接提取文本
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ (text, success)
|
|
|
+ """
|
|
|
+ try:
|
|
|
+ page = pdf_doc[page_idx]
|
|
|
+
|
|
|
+ # 将图像坐标转换为PDF坐标
|
|
|
+ pdf_bbox = fitz.Rect(
|
|
|
+ bbox[0] / scale,
|
|
|
+ bbox[1] / scale,
|
|
|
+ bbox[2] / scale,
|
|
|
+ bbox[3] / scale
|
|
|
+ )
|
|
|
+
|
|
|
+ # 提取区域内的文本
|
|
|
+ text = page.get_text("text", clip=pdf_bbox)
|
|
|
+
|
|
|
+ return text.strip(), bool(text.strip())
|
|
|
+
|
|
|
+ except Exception as e:
|
|
|
+ logger.debug(f"PDF text extraction error: {e}")
|
|
|
+ return "", False
|
|
|
+
|
|
|
+ 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[:4])
|
|
|
+ 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,
|
|
|
+ region_bbox: List[float]
|
|
|
+ ) -> List:
|
|
|
+ """将相对坐标转换为绝对坐标"""
|
|
|
+ if not relative_bbox or len(region_bbox) < 4:
|
|
|
+ return relative_bbox
|
|
|
+
|
|
|
+ bx1, by1 = region_bbox[0], region_bbox[1]
|
|
|
+
|
|
|
+ # 处理4点坐标格式 [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
|
|
|
+ if isinstance(relative_bbox[0], (list, tuple)):
|
|
|
+ return [
|
|
|
+ [p[0] + bx1, p[1] + by1] for p in relative_bbox
|
|
|
+ ]
|
|
|
+
|
|
|
+ # 处理4值坐标格式 [x1, y1, x2, y2]
|
|
|
+ if len(relative_bbox) >= 4:
|
|
|
+ return [
|
|
|
+ relative_bbox[0] + bx1,
|
|
|
+ relative_bbox[1] + by1,
|
|
|
+ relative_bbox[2] + bx1,
|
|
|
+ relative_bbox[3] + by1
|
|
|
+ ]
|
|
|
+
|
|
|
+ return relative_bbox
|
|
|
+
|
|
|
+ def _convert_ocr_to_matcher_format(
|
|
|
+ self,
|
|
|
+ ocr_poly: List,
|
|
|
+ text: str,
|
|
|
+ confidence: float,
|
|
|
+ idx: int,
|
|
|
+ table_bbox: Optional[List[float]] = None
|
|
|
+ ) -> Optional[Dict[str, Any]]:
|
|
|
+ """
|
|
|
+ 将 OCR 结果转换为 TableCellMatcher 期望的格式
|
|
|
+
|
|
|
+ OCR 返回格式:
|
|
|
+ bbox: [[x1,y1], [x2,y1], [x2,y2], [x1,y2]] # 4点多边形
|
|
|
+ text: str
|
|
|
+ confidence: float
|
|
|
+
|
|
|
+ TableCellMatcher 期望格式:
|
|
|
+ text: str
|
|
|
+ bbox: [x_min, y_min, x_max, y_max] # 4值矩形
|
|
|
+ poly: [[x1,y1], [x2,y1], [x2,y2], [x1,y2]] # 4点多边形
|
|
|
+ score: float
|
|
|
+ paddle_bbox_index: int
|
|
|
+ used: bool
|
|
|
+
|
|
|
+ Args:
|
|
|
+ ocr_poly: OCR返回的4点多边形坐标
|
|
|
+ text: 识别文本
|
|
|
+ confidence: 置信度
|
|
|
+ idx: 索引
|
|
|
+ table_bbox: 表格的绝对坐标,用于转换相对坐标
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 转换后的字典,或 None(如果无效)
|
|
|
+ """
|
|
|
+ if not ocr_poly or not text:
|
|
|
+ return None
|
|
|
+
|
|
|
+ # 处理不同的坐标格式
|
|
|
+ poly = []
|
|
|
+
|
|
|
+ # 格式1: [[x1,y1], [x2,y2], [x3,y3], [x4,y4]] - 4点多边形
|
|
|
+ if isinstance(ocr_poly[0], (list, tuple)) and len(ocr_poly) == 4:
|
|
|
+ poly = [[float(p[0]), float(p[1])] for p in ocr_poly]
|
|
|
+
|
|
|
+ # 格式2: [x1, y1, x2, y1, x2, y2, x1, y2] - 8值展平格式
|
|
|
+ elif len(ocr_poly) == 8 and isinstance(ocr_poly[0], (int, float)):
|
|
|
+ poly = [
|
|
|
+ [float(ocr_poly[0]), float(ocr_poly[1])],
|
|
|
+ [float(ocr_poly[2]), float(ocr_poly[3])],
|
|
|
+ [float(ocr_poly[4]), float(ocr_poly[5])],
|
|
|
+ [float(ocr_poly[6]), float(ocr_poly[7])]
|
|
|
+ ]
|
|
|
+
|
|
|
+ # 格式3: [x1, y1, x2, y2] - 4值矩形格式
|
|
|
+ elif len(ocr_poly) == 4 and isinstance(ocr_poly[0], (int, float)):
|
|
|
+ x1, y1, x2, y2 = ocr_poly
|
|
|
+ poly = [
|
|
|
+ [float(x1), float(y1)],
|
|
|
+ [float(x2), float(y1)],
|
|
|
+ [float(x2), float(y2)],
|
|
|
+ [float(x1), float(y2)]
|
|
|
+ ]
|
|
|
+ else:
|
|
|
+ logger.warning(f"Unknown OCR bbox format: {ocr_poly}")
|
|
|
+ return None
|
|
|
+
|
|
|
+ # 转换为绝对坐标(相对于整页图片)
|
|
|
+ if table_bbox and len(table_bbox) >= 2:
|
|
|
+ offset_x, offset_y = table_bbox[0], table_bbox[1]
|
|
|
+ poly = [[p[0] + offset_x, p[1] + offset_y] for p in poly]
|
|
|
+
|
|
|
+ # 从多边形计算 bbox [x_min, y_min, x_max, y_max]
|
|
|
+ xs = [p[0] for p in poly]
|
|
|
+ ys = [p[1] for p in poly]
|
|
|
+ bbox = [min(xs), min(ys), max(xs), max(ys)]
|
|
|
+
|
|
|
+ return {
|
|
|
+ 'text': text,
|
|
|
+ 'bbox': bbox,
|
|
|
+ 'poly': poly,
|
|
|
+ 'score': confidence,
|
|
|
+ 'paddle_bbox_index': idx,
|
|
|
+ 'used': False
|
|
|
+ }
|
|
|
+
|
|
|
+ def _inverse_rotate_table_coords(
|
|
|
+ self,
|
|
|
+ cells: List[Dict],
|
|
|
+ html: str,
|
|
|
+ rotation_angle: float,
|
|
|
+ orig_table_size: Tuple[int, int],
|
|
|
+ table_bbox: List[float]
|
|
|
+ ) -> Tuple[List[Dict], str]:
|
|
|
+ """
|
|
|
+ 将旋转后的坐标逆向转换回原图坐标
|
|
|
+
|
|
|
+ 参考 merger.DataProcessor 的实现
|
|
|
+
|
|
|
+ Args:
|
|
|
+ cells: 单元格列表(坐标是旋转后的)
|
|
|
+ html: HTML字符串(data-bbox是旋转后的)
|
|
|
+ rotation_angle: 旋转角度
|
|
|
+ orig_table_size: 原始表格尺寸 (width, height)
|
|
|
+ table_bbox: 表格在整页图片中的位置 [x1, y1, x2, y2]
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ (转换后的cells, 转换后的html)
|
|
|
+ """
|
|
|
+ import re
|
|
|
+ import json
|
|
|
+
|
|
|
+ table_offset_x, table_offset_y = table_bbox[0], table_bbox[1]
|
|
|
+
|
|
|
+ # 转换 cells 中的 bbox
|
|
|
+ converted_cells = []
|
|
|
+ for cell in cells:
|
|
|
+ cell_copy = cell.copy()
|
|
|
+ cell_bbox = cell.get('bbox', [])
|
|
|
+ if cell_bbox and len(cell_bbox) == 4:
|
|
|
+ # 先逆向旋转,再加上表格偏移量
|
|
|
+ orig_bbox = BBoxExtractor.inverse_rotate_box_coordinates(
|
|
|
+ cell_bbox, rotation_angle, orig_table_size
|
|
|
+ )
|
|
|
+ # 加上表格偏移量转换为整页坐标
|
|
|
+ cell_copy['bbox'] = [
|
|
|
+ orig_bbox[0] + table_offset_x,
|
|
|
+ orig_bbox[1] + table_offset_y,
|
|
|
+ orig_bbox[2] + table_offset_x,
|
|
|
+ orig_bbox[3] + table_offset_y
|
|
|
+ ]
|
|
|
+ converted_cells.append(cell_copy)
|
|
|
+
|
|
|
+ # 转换 HTML 中的 data-bbox
|
|
|
+ def replace_bbox(match):
|
|
|
+ bbox_str = match.group(1)
|
|
|
+ try:
|
|
|
+ bbox = json.loads(bbox_str)
|
|
|
+ if len(bbox) == 4:
|
|
|
+ # 先逆向旋转
|
|
|
+ orig_bbox = BBoxExtractor.inverse_rotate_box_coordinates(
|
|
|
+ bbox, rotation_angle, orig_table_size
|
|
|
+ )
|
|
|
+ # 加上表格偏移量
|
|
|
+ new_bbox = [
|
|
|
+ orig_bbox[0] + table_offset_x,
|
|
|
+ orig_bbox[1] + table_offset_y,
|
|
|
+ orig_bbox[2] + table_offset_x,
|
|
|
+ orig_bbox[3] + table_offset_y
|
|
|
+ ]
|
|
|
+ return f'data-bbox="[{new_bbox[0]},{new_bbox[1]},{new_bbox[2]},{new_bbox[3]}]"'
|
|
|
+ except:
|
|
|
+ pass
|
|
|
+ return match.group(0)
|
|
|
+
|
|
|
+ converted_html = re.sub(
|
|
|
+ r'data-bbox="\[([^\]]+)\]"',
|
|
|
+ replace_bbox,
|
|
|
+ html
|
|
|
+ )
|
|
|
+
|
|
|
+ return converted_cells, converted_html
|
|
|
+
|
|
|
+ def _add_table_offset_to_cells(
|
|
|
+ self,
|
|
|
+ cells: List[Dict],
|
|
|
+ table_bbox: List[float]
|
|
|
+ ) -> List[Dict]:
|
|
|
+ """
|
|
|
+ 为单元格坐标添加表格偏移量(无旋转情况)
|
|
|
+
|
|
|
+ Args:
|
|
|
+ cells: 单元格列表
|
|
|
+ table_bbox: 表格位置 [x1, y1, x2, y2]
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 转换后的 cells
|
|
|
+ """
|
|
|
+ offset_x, offset_y = table_bbox[0], table_bbox[1]
|
|
|
+
|
|
|
+ converted_cells = []
|
|
|
+ for cell in cells:
|
|
|
+ cell_copy = cell.copy()
|
|
|
+ cell_bbox = cell.get('bbox', [])
|
|
|
+ if cell_bbox and len(cell_bbox) == 4:
|
|
|
+ cell_copy['bbox'] = [
|
|
|
+ cell_bbox[0] + offset_x,
|
|
|
+ cell_bbox[1] + offset_y,
|
|
|
+ cell_bbox[2] + offset_x,
|
|
|
+ cell_bbox[3] + offset_y
|
|
|
+ ]
|
|
|
+ converted_cells.append(cell_copy)
|
|
|
+
|
|
|
+ return converted_cells
|
|
|
+
|
|
|
+ def _add_table_offset_to_html(
|
|
|
+ self,
|
|
|
+ html: str,
|
|
|
+ table_bbox: List[float]
|
|
|
+ ) -> str:
|
|
|
+ """
|
|
|
+ 为HTML中的data-bbox添加表格偏移量(无旋转情况)
|
|
|
+
|
|
|
+ Args:
|
|
|
+ html: HTML字符串
|
|
|
+ table_bbox: 表格位置 [x1, y1, x2, y2]
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 转换后的 HTML
|
|
|
+ """
|
|
|
+ import re
|
|
|
+ import json
|
|
|
+
|
|
|
+ offset_x, offset_y = table_bbox[0], table_bbox[1]
|
|
|
+
|
|
|
+ def replace_bbox(match):
|
|
|
+ bbox_str = match.group(1)
|
|
|
+ try:
|
|
|
+ bbox = json.loads(f"[{bbox_str}]")
|
|
|
+ if len(bbox) == 4:
|
|
|
+ new_bbox = [
|
|
|
+ bbox[0] + offset_x,
|
|
|
+ bbox[1] + offset_y,
|
|
|
+ bbox[2] + offset_x,
|
|
|
+ bbox[3] + offset_y
|
|
|
+ ]
|
|
|
+ return f'data-bbox="[{new_bbox[0]},{new_bbox[1]},{new_bbox[2]},{new_bbox[3]}]"'
|
|
|
+ except:
|
|
|
+ pass
|
|
|
+ return match.group(0)
|
|
|
+
|
|
|
+ converted_html = re.sub(
|
|
|
+ r'data-bbox="\[([^\]]+)\]"',
|
|
|
+ replace_bbox,
|
|
|
+ html
|
|
|
+ )
|
|
|
+
|
|
|
+ return converted_html
|
|
|
+
|
|
|
+ def _merge_cross_page_tables(self, results: Dict[str, Any]) -> Dict[str, Any]:
|
|
|
+ """合并跨页表格"""
|
|
|
+ # TODO: 实现跨页表格合并逻辑
|
|
|
+ # 可以参考 MinerU 的 cross_page_table_merge 实现
|
|
|
+ return results
|
|
|
+
|
|
|
+ def _transform_coords_to_original(
|
|
|
+ self,
|
|
|
+ element: Dict[str, Any],
|
|
|
+ rotate_angle: int,
|
|
|
+ rotated_shape: Tuple,
|
|
|
+ original_shape: Tuple
|
|
|
+ ) -> Dict[str, Any]:
|
|
|
+ """
|
|
|
+ 将坐标从旋转后的图片坐标系转换回原始图片坐标系
|
|
|
+
|
|
|
+ 使用 BBoxExtractor.inverse_rotate_box_coordinates 进行转换
|
|
|
+
|
|
|
+ Args:
|
|
|
+ element: 元素字典(包含bbox等坐标信息)
|
|
|
+ rotate_angle: 旋转角度(0, 90, 180, 270)
|
|
|
+ rotated_shape: 旋转后图片的shape (h, w, c)
|
|
|
+ original_shape: 原始图片的shape (h, w, c)
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 坐标转换后的元素字典(深拷贝)
|
|
|
+ """
|
|
|
+ import copy
|
|
|
+ element = copy.deepcopy(element)
|
|
|
+
|
|
|
+ if rotate_angle == 0 or not MERGER_AVAILABLE:
|
|
|
+ return element
|
|
|
+
|
|
|
+ # 原始图片尺寸 (width, height) - BBoxExtractor 使用的格式
|
|
|
+ orig_h, orig_w = original_shape[:2]
|
|
|
+ orig_image_size = (orig_w, orig_h)
|
|
|
+
|
|
|
+ # 转换主bbox
|
|
|
+ if 'bbox' in element and element['bbox']:
|
|
|
+ element['bbox'] = BBoxExtractor.inverse_rotate_box_coordinates(
|
|
|
+ element['bbox'], rotate_angle, orig_image_size
|
|
|
+ )
|
|
|
+
|
|
|
+ # 转换表格相关坐标
|
|
|
+ if element.get('type') == 'table' and 'content' in element:
|
|
|
+ content = element['content']
|
|
|
+
|
|
|
+ # 转换 OCR boxes
|
|
|
+ if 'ocr_boxes' in content and content['ocr_boxes']:
|
|
|
+ for box in content['ocr_boxes']:
|
|
|
+ if 'bbox' in box and box['bbox']:
|
|
|
+ box['bbox'] = BBoxExtractor.inverse_rotate_box_coordinates(
|
|
|
+ box['bbox'], rotate_angle, orig_image_size
|
|
|
+ )
|
|
|
+
|
|
|
+ # 转换 cells
|
|
|
+ if 'cells' in content and content['cells']:
|
|
|
+ for cell in content['cells']:
|
|
|
+ if 'bbox' in cell and cell['bbox']:
|
|
|
+ cell['bbox'] = BBoxExtractor.inverse_rotate_box_coordinates(
|
|
|
+ cell.get('bbox', []), rotate_angle, orig_image_size
|
|
|
+ )
|
|
|
+
|
|
|
+ # 转换 HTML 中的 data-bbox 属性
|
|
|
+ if 'html' in content and content['html']:
|
|
|
+ content['html'] = self._transform_html_data_bbox(
|
|
|
+ content['html'], rotate_angle, orig_image_size
|
|
|
+ )
|
|
|
+
|
|
|
+ # 转换文本OCR details
|
|
|
+ if 'content' in element and 'ocr_details' in element.get('content', {}):
|
|
|
+ ocr_details = element['content'].get('ocr_details', [])
|
|
|
+ if ocr_details:
|
|
|
+ for detail in ocr_details:
|
|
|
+ if 'bbox' in detail and detail['bbox']:
|
|
|
+ detail['bbox'] = BBoxExtractor.inverse_rotate_box_coordinates(
|
|
|
+ detail.get('bbox', []), rotate_angle, orig_image_size
|
|
|
+ )
|
|
|
+
|
|
|
+ return element
|
|
|
+
|
|
|
+ def _transform_html_data_bbox(
|
|
|
+ self,
|
|
|
+ html: str,
|
|
|
+ rotate_angle: int,
|
|
|
+ orig_image_size: Tuple[int, int]
|
|
|
+ ) -> str:
|
|
|
+ """
|
|
|
+ 转换 HTML 中所有 data-bbox 属性的坐标
|
|
|
+
|
|
|
+ Args:
|
|
|
+ html: 包含 data-bbox 属性的 HTML 字符串
|
|
|
+ rotate_angle: 旋转角度
|
|
|
+ orig_image_size: 原始图片尺寸 (width, height)
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 转换后的 HTML 字符串
|
|
|
+ """
|
|
|
+ import re
|
|
|
+ import json
|
|
|
+
|
|
|
+ def replace_bbox(match):
|
|
|
+ try:
|
|
|
+ bbox_str = match.group(1)
|
|
|
+ bbox = json.loads(bbox_str)
|
|
|
+ if bbox and len(bbox) == 4:
|
|
|
+ transformed = BBoxExtractor.inverse_rotate_box_coordinates(
|
|
|
+ bbox, rotate_angle, orig_image_size
|
|
|
+ )
|
|
|
+ return f'data-bbox="{json.dumps(transformed)}"'
|
|
|
+ except (json.JSONDecodeError, ValueError):
|
|
|
+ pass
|
|
|
+ return match.group(0) # 返回原始匹配
|
|
|
+
|
|
|
+ # 匹配 data-bbox="[...]" 格式
|
|
|
+ pattern = r'data-bbox="(\[[^\]]+\])"'
|
|
|
+ return re.sub(pattern, replace_bbox, html)
|
|
|
+
|
|
|
+ def _create_error_element(
|
|
|
+ self,
|
|
|
+ layout_item: Dict[str, Any],
|
|
|
+ error_msg: str
|
|
|
+ ) -> Dict[str, Any]:
|
|
|
+ """创建错误元素"""
|
|
|
+ return {
|
|
|
+ 'type': 'error',
|
|
|
+ 'bbox': layout_item.get('bbox', [0, 0, 0, 0]),
|
|
|
+ 'confidence': 0.0,
|
|
|
+ 'content': {'error': error_msg},
|
|
|
+ 'original_category': layout_item.get('category', 'unknown')
|
|
|
+ }
|
|
|
+
|
|
|
+ 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()
|
|
|
+
|