|
|
@@ -21,6 +21,7 @@ from typing import Dict, List, Any, Optional
|
|
|
from pathlib import Path
|
|
|
import numpy as np
|
|
|
from loguru import logger
|
|
|
+import traceback
|
|
|
|
|
|
# 添加 ocr_platform 根目录到 Python 路径(用于导入 ocr_utils)
|
|
|
ocr_platform_root = Path(__file__).parents[3] # core -> universal_doc_parser -> ocr_tools -> ocr_platform -> repository.git
|
|
|
@@ -36,16 +37,17 @@ if str(module_root) not in sys.path:
|
|
|
try:
|
|
|
from .model_factory import ModelFactory
|
|
|
from .config_manager import ConfigManager
|
|
|
- from .coordinate_utils import CoordinateUtils
|
|
|
+ from ocr_utils.coordinate_utils import CoordinateUtils
|
|
|
+ from ocr_utils import PDFUtils
|
|
|
+ from .table_coordinate_utils import TableCoordinateUtils
|
|
|
from .layout_utils import LayoutUtils, SpanMatcher
|
|
|
- from .pdf_utils import PDFUtils
|
|
|
from .element_processors import ElementProcessors
|
|
|
except ImportError:
|
|
|
from model_factory import ModelFactory
|
|
|
from config_manager import ConfigManager
|
|
|
- from coordinate_utils import CoordinateUtils
|
|
|
+ from ocr_utils import CoordinateUtils, PDFUtils
|
|
|
+ from table_coordinate_utils import TableCoordinateUtils
|
|
|
from layout_utils import LayoutUtils, SpanMatcher
|
|
|
- from pdf_utils import PDFUtils
|
|
|
from element_processors import ElementProcessors
|
|
|
|
|
|
# 从 ocr_tools.ocr_merger 导入 merger 组件
|
|
|
@@ -126,6 +128,13 @@ class EnhancedDocPipeline:
|
|
|
self.config['layout_detection']
|
|
|
)
|
|
|
|
|
|
+ # 如果是智能路由器且使用ocr_eval策略,需要设置OCR识别器
|
|
|
+ if hasattr(self.layout_detector, 'set_ocr_recognizer'):
|
|
|
+ # OCR识别器会在后面初始化,这里先标记
|
|
|
+ self._smart_router_needs_ocr = True
|
|
|
+ else:
|
|
|
+ self._smart_router_needs_ocr = False
|
|
|
+
|
|
|
# 3. VL识别器(表格、公式)
|
|
|
if self.config.get('vl_recognition', {}) != {}:
|
|
|
self.vl_recognizer = ModelFactory.create_vl_recognizer(
|
|
|
@@ -138,6 +147,11 @@ class EnhancedDocPipeline:
|
|
|
self.ocr_recognizer = ModelFactory.create_ocr_recognizer(
|
|
|
self.config['ocr_recognition']
|
|
|
)
|
|
|
+
|
|
|
+ # 如果使用智能路由器,设置OCR识别器
|
|
|
+ if self._smart_router_needs_ocr and hasattr(self.layout_detector, 'set_ocr_recognizer'):
|
|
|
+ self.layout_detector.set_ocr_recognizer(self.ocr_recognizer)
|
|
|
+ logger.info("✅ OCR recognizer set for smart router")
|
|
|
|
|
|
# 5. 表格分类器(可选)
|
|
|
self.table_classifier = None
|
|
|
@@ -357,40 +371,65 @@ class EnhancedDocPipeline:
|
|
|
|
|
|
# 2. Layout检测
|
|
|
try:
|
|
|
- layout_results = self.layout_detector.detect(detection_image)
|
|
|
- logger.info(f"📋 Page {page_idx}: detected {len(layout_results)} elements (before dedup)")
|
|
|
+ # 如果使用智能路由器且策略是ocr_eval,需要先获取OCR spans(只检测文本框,不识别文字)
|
|
|
+ ocr_spans_for_layout = None
|
|
|
+ if hasattr(self.layout_detector, 'strategy') and self.layout_detector.strategy == 'ocr_eval':
|
|
|
+ # 先运行OCR检测获取文本框(用于layout评估,不需要识别文字)
|
|
|
+ if self.ocr_recognizer:
|
|
|
+ try:
|
|
|
+ # 使用 detect_text_boxes 只检测文本框,不识别文字(性能优化)
|
|
|
+ if hasattr(self.ocr_recognizer, 'detect_text_boxes'):
|
|
|
+ ocr_spans_for_layout = self.ocr_recognizer.detect_text_boxes(detection_image)
|
|
|
+ else:
|
|
|
+ # 如果没有 detect_text_boxes 方法,使用完整识别
|
|
|
+ ocr_spans_for_layout = self.ocr_recognizer.recognize_text(detection_image)
|
|
|
+
|
|
|
+ ocr_spans_for_layout = SpanMatcher.remove_duplicate_spans(ocr_spans_for_layout)
|
|
|
+ ocr_spans_for_layout = self._sort_spans_by_position(ocr_spans_for_layout)
|
|
|
+ logger.info(f"📝 Pre-OCR text box detection for layout evaluation: {len(ocr_spans_for_layout)} boxes")
|
|
|
+ except Exception as e:
|
|
|
+ logger.warning(f"⚠️ Pre-OCR text box detection for layout evaluation failed: {e}")
|
|
|
+
|
|
|
+ # 如果是 SmartLayoutRouter 且启用调试模式,设置调试信息
|
|
|
+ if hasattr(self.layout_detector, 'set_ocr_recognizer'):
|
|
|
+ # 这是 SmartLayoutRouter
|
|
|
+ if self.debug_mode and output_dir:
|
|
|
+ self.layout_detector.debug_mode = self.debug_mode # type: ignore
|
|
|
+ self.layout_detector.output_dir = output_dir # type: ignore
|
|
|
+
|
|
|
+ # 调用layout检测(传递OCR spans如果可用)
|
|
|
+ if ocr_spans_for_layout is not None and hasattr(self.layout_detector, 'detect'):
|
|
|
+ # SmartLayoutRouter 的 detect 方法支持 ocr_spans 和 page_name 参数
|
|
|
+ try:
|
|
|
+ layout_results = self.layout_detector.detect(
|
|
|
+ detection_image,
|
|
|
+ ocr_spans=ocr_spans_for_layout,
|
|
|
+ page_name=page_name
|
|
|
+ )
|
|
|
+ except TypeError:
|
|
|
+ # 如果方法不支持这些参数,使用默认调用
|
|
|
+ try:
|
|
|
+ layout_results = self.layout_detector.detect(detection_image, ocr_spans=ocr_spans_for_layout)
|
|
|
+ except TypeError:
|
|
|
+ layout_results = self.layout_detector.detect(detection_image)
|
|
|
+ else:
|
|
|
+ # 普通 detector 或没有 ocr_spans
|
|
|
+ if hasattr(self.layout_detector, 'detect') and hasattr(self.layout_detector, 'set_ocr_recognizer'):
|
|
|
+ # SmartLayoutRouter,尝试传递 page_name
|
|
|
+ try:
|
|
|
+ layout_results = self.layout_detector.detect(detection_image, page_name=page_name)
|
|
|
+ except TypeError:
|
|
|
+ layout_results = self.layout_detector.detect(detection_image)
|
|
|
+ else:
|
|
|
+ layout_results = self.layout_detector.detect(detection_image)
|
|
|
+
|
|
|
+ logger.info(f"📋 Page {page_idx}: detected {len(layout_results)} elements (post-processed)")
|
|
|
except Exception as e:
|
|
|
logger.error(f"❌ Layout detection failed: {e}")
|
|
|
layout_results = []
|
|
|
|
|
|
- # 2.5 处理重叠框
|
|
|
- if layout_results:
|
|
|
- original_count = len(layout_results)
|
|
|
- layout_results = LayoutUtils.remove_overlapping_boxes(layout_results)
|
|
|
- removed_count = original_count - len(layout_results)
|
|
|
- if removed_count > 0:
|
|
|
- logger.info(f"🔄 Page {page_idx}: removed {removed_count} overlapping boxes")
|
|
|
-
|
|
|
- # 2.6 将大面积文本块转换为表格(后处理)
|
|
|
- if layout_results:
|
|
|
- convert_large_text = self.config.get('layout', {}).get('convert_large_text_to_table', False)
|
|
|
- if convert_large_text:
|
|
|
- # 获取图像尺寸 (height, width)
|
|
|
- img_shape = detection_image.shape
|
|
|
- if len(img_shape) >= 2:
|
|
|
- image_shape = (int(img_shape[0]), int(img_shape[1]))
|
|
|
- else:
|
|
|
- image_shape = (0, 0)
|
|
|
-
|
|
|
- layout_results = LayoutUtils.convert_large_text_to_table(
|
|
|
- layout_results,
|
|
|
- image_shape,
|
|
|
- min_area_ratio=self.config.get('layout', {}).get('min_text_area_ratio', 0.25),
|
|
|
- min_width_ratio=self.config.get('layout', {}).get('min_text_width_ratio', 0.4),
|
|
|
- min_height_ratio=self.config.get('layout', {}).get('min_text_height_ratio', 0.3)
|
|
|
- )
|
|
|
-
|
|
|
- page_result['layout_raw'] = layout_results
|
|
|
+ # 2.5 后处理布局结果
|
|
|
+ # 注意:所有 detector 的 detect() 方法已经自动执行了后处理
|
|
|
|
|
|
# 3. 整页 OCR 获取所有 text spans(关键改进)
|
|
|
all_text_spans = []
|
|
|
@@ -485,12 +524,12 @@ class EnhancedDocPipeline:
|
|
|
# 8. 坐标转换回原始图片坐标系
|
|
|
if rotate_angle != 0:
|
|
|
sorted_elements = [
|
|
|
- CoordinateUtils.transform_coords_to_original(
|
|
|
+ TableCoordinateUtils.transform_coords_to_original(
|
|
|
element, rotate_angle, detection_image.shape, original_image.shape
|
|
|
) for element in sorted_elements
|
|
|
]
|
|
|
sorted_discarded = [
|
|
|
- CoordinateUtils.transform_coords_to_original(
|
|
|
+ TableCoordinateUtils.transform_coords_to_original(
|
|
|
element, rotate_angle, detection_image.shape, original_image.shape
|
|
|
) for element in sorted_discarded
|
|
|
]
|
|
|
@@ -872,8 +911,8 @@ class EnhancedDocPipeline:
|
|
|
|
|
|
processed_elements.append(element)
|
|
|
except Exception as e:
|
|
|
- logger.warning(f"⚠️ Table processing failed: {e}")
|
|
|
- processed_elements.append(ElementProcessors.create_error_element(item, str(e)))
|
|
|
+ logger.error(f"⚠️ Table processing failed: {e}, traceback: {traceback.format_exc()}")
|
|
|
+ raise Exception(f"Table processing failed: {e}, traceback: {traceback.format_exc()}")
|
|
|
|
|
|
# 处理公式元素
|
|
|
for item in classified_elements['equation']:
|