|
|
@@ -36,14 +36,14 @@ try:
|
|
|
from .model_factory import ModelFactory
|
|
|
from .config_manager import ConfigManager
|
|
|
from .coordinate_utils import CoordinateUtils
|
|
|
- from .layout_utils import LayoutUtils
|
|
|
+ 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 layout_utils import LayoutUtils
|
|
|
+ from layout_utils import LayoutUtils, SpanMatcher
|
|
|
from pdf_utils import PDFUtils
|
|
|
from element_processors import ElementProcessors
|
|
|
|
|
|
@@ -162,12 +162,19 @@ class EnhancedDocPipeline:
|
|
|
|
|
|
# ==================== 主处理流程 ====================
|
|
|
|
|
|
- def process_document(self, document_path: str) -> Dict[str, Any]:
|
|
|
+ def process_document(
|
|
|
+ self,
|
|
|
+ document_path: str,
|
|
|
+ page_range: Optional[str] = None
|
|
|
+ ) -> Dict[str, Any]:
|
|
|
"""
|
|
|
处理文档主流程
|
|
|
|
|
|
Args:
|
|
|
document_path: 文档路径(PDF、图片或目录)
|
|
|
+ page_range: 页面范围字符串,如 "1-5,7,9-12"
|
|
|
+ - PDF:按页码(从1开始)
|
|
|
+ - 图片目录:按文件名排序后的位置(从1开始)
|
|
|
|
|
|
Returns:
|
|
|
处理结果字典
|
|
|
@@ -186,22 +193,27 @@ class EnhancedDocPipeline:
|
|
|
# 1. 加载文档并分类
|
|
|
dpi = self.config.get('input', {}).get('dpi', 200)
|
|
|
images, pdf_type, pdf_doc = PDFUtils.load_and_classify_document(
|
|
|
- doc_path, dpi=dpi
|
|
|
+ doc_path, dpi=dpi, page_range=page_range
|
|
|
)
|
|
|
results['metadata']['pdf_type'] = pdf_type
|
|
|
results['metadata']['page_count'] = len(images)
|
|
|
+ results['metadata']['page_range'] = page_range
|
|
|
|
|
|
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)}")
|
|
|
+ for idx, image_dict in enumerate(images):
|
|
|
+ # 使用原始页码索引(支持页面范围过滤)
|
|
|
+ page_idx = image_dict.get('page_idx', idx)
|
|
|
+ page_name = image_dict.get('page_name', f'page_{page_idx + 1:03d}')
|
|
|
+ logger.info(f"🔍 Processing page {idx + 1}/{len(images)} (original index: {page_idx})")
|
|
|
|
|
|
page_result = self._process_single_page(
|
|
|
image_dict=image_dict,
|
|
|
page_idx=page_idx,
|
|
|
pdf_type=pdf_type,
|
|
|
- pdf_doc=pdf_doc
|
|
|
+ pdf_doc=pdf_doc,
|
|
|
+ page_name=page_name
|
|
|
)
|
|
|
results['pages'].append(page_result)
|
|
|
|
|
|
@@ -224,7 +236,8 @@ class EnhancedDocPipeline:
|
|
|
image_dict: Dict[str, Any],
|
|
|
page_idx: int,
|
|
|
pdf_type: str,
|
|
|
- pdf_doc: Optional[Any] = None
|
|
|
+ pdf_doc: Optional[Any] = None,
|
|
|
+ page_name: Optional[str] = None
|
|
|
) -> Dict[str, Any]:
|
|
|
"""
|
|
|
处理单页文档
|
|
|
@@ -240,8 +253,13 @@ class EnhancedDocPipeline:
|
|
|
scale = image_dict.get('scale', 1.0)
|
|
|
original_image = np.array(pil_image)
|
|
|
|
|
|
+ # 页面名称(用于输出文件命名)
|
|
|
+ if page_name is None:
|
|
|
+ page_name = f"page_{page_idx + 1:03d}"
|
|
|
+
|
|
|
page_result = {
|
|
|
'page_idx': page_idx,
|
|
|
+ 'page_name': page_name, # 用于输出文件命名
|
|
|
'elements': [],
|
|
|
'image_shape': original_image.shape,
|
|
|
'original_image': original_image,
|
|
|
@@ -283,24 +301,41 @@ class EnhancedDocPipeline:
|
|
|
|
|
|
page_result['layout_raw'] = layout_results
|
|
|
|
|
|
- # 3. 分类元素
|
|
|
+ # 3. 整页 OCR 获取所有 text spans(关键改进)
|
|
|
+ all_ocr_spans = []
|
|
|
+ try:
|
|
|
+ all_ocr_spans = self.ocr_recognizer.recognize_text(detection_image)
|
|
|
+ # 去除重复 spans
|
|
|
+ all_ocr_spans = SpanMatcher.remove_duplicate_spans(all_ocr_spans)
|
|
|
+ logger.info(f"📝 Page {page_idx}: OCR detected {len(all_ocr_spans)} text spans")
|
|
|
+ except Exception as e:
|
|
|
+ logger.warning(f"⚠️ Full-page OCR failed: {e}")
|
|
|
+
|
|
|
+ # 4. 将 OCR spans 匹配到 layout blocks
|
|
|
+ matched_spans = SpanMatcher.match_spans_to_blocks(
|
|
|
+ all_ocr_spans, layout_results, overlap_threshold=0.5
|
|
|
+ )
|
|
|
+
|
|
|
+ # 5. 分类元素
|
|
|
classified_elements = self._classify_elements(layout_results, page_idx)
|
|
|
|
|
|
- # 4. 处理各类元素
|
|
|
+ # 6. 处理各类元素(传入匹配的 spans)
|
|
|
processed_elements, discarded_elements = self._process_all_elements(
|
|
|
detection_image=detection_image,
|
|
|
classified_elements=classified_elements,
|
|
|
pdf_type=pdf_type,
|
|
|
pdf_doc=pdf_doc,
|
|
|
page_idx=page_idx,
|
|
|
- scale=scale
|
|
|
+ scale=scale,
|
|
|
+ matched_spans=matched_spans,
|
|
|
+ layout_results=layout_results
|
|
|
)
|
|
|
|
|
|
- # 5. 按阅读顺序排序
|
|
|
+ # 7. 按阅读顺序排序
|
|
|
sorted_elements = LayoutUtils.sort_elements_by_reading_order(processed_elements)
|
|
|
sorted_discarded = LayoutUtils.sort_elements_by_reading_order(discarded_elements)
|
|
|
|
|
|
- # 6. 坐标转换回原始图片坐标系
|
|
|
+ # 8. 坐标转换回原始图片坐标系
|
|
|
if rotate_angle != 0:
|
|
|
sorted_elements = [
|
|
|
CoordinateUtils.transform_coords_to_original(
|
|
|
@@ -388,7 +423,9 @@ class EnhancedDocPipeline:
|
|
|
pdf_type: str,
|
|
|
pdf_doc: Optional[Any],
|
|
|
page_idx: int,
|
|
|
- scale: float
|
|
|
+ scale: float,
|
|
|
+ matched_spans: Optional[Dict[int, List[Dict[str, Any]]]] = None,
|
|
|
+ layout_results: Optional[List[Dict[str, Any]]] = None
|
|
|
) -> tuple:
|
|
|
"""
|
|
|
处理所有分类后的元素
|
|
|
@@ -400,6 +437,8 @@ class EnhancedDocPipeline:
|
|
|
pdf_doc: PDF文档对象
|
|
|
page_idx: 页码
|
|
|
scale: 缩放比例
|
|
|
+ matched_spans: 匹配的 OCR spans {block_idx: [spans]}
|
|
|
+ layout_results: 原始 layout 检测结果(用于索引匹配)
|
|
|
|
|
|
Returns:
|
|
|
(processed_elements, discarded_elements)
|
|
|
@@ -407,11 +446,29 @@ class EnhancedDocPipeline:
|
|
|
processed_elements = []
|
|
|
discarded_elements = []
|
|
|
|
|
|
+ # 构建 layout item 到 block index 的映射
|
|
|
+ item_to_block_idx = {}
|
|
|
+ if layout_results:
|
|
|
+ for idx, item in enumerate(layout_results):
|
|
|
+ # 使用 bbox 作为唯一标识(转为元组方便哈希)
|
|
|
+ bbox_key = tuple(item.get('bbox', [0, 0, 0, 0]))
|
|
|
+ item_to_block_idx[bbox_key] = idx
|
|
|
+
|
|
|
+ def get_matched_spans_for_item(item: Dict[str, Any]) -> List[Dict[str, Any]]:
|
|
|
+ """获取某个 layout item 匹配的 spans"""
|
|
|
+ if not matched_spans or not layout_results:
|
|
|
+ return []
|
|
|
+ bbox_key = tuple[Any, ...](item.get('bbox', [0, 0, 0, 0]))
|
|
|
+ block_idx = item_to_block_idx.get(bbox_key, -1)
|
|
|
+ return matched_spans.get(block_idx, [])
|
|
|
+
|
|
|
# 处理纯文本区域
|
|
|
for item in classified_elements['text']:
|
|
|
try:
|
|
|
+ spans = get_matched_spans_for_item(item)
|
|
|
element = self.element_processors.process_text_element(
|
|
|
- detection_image, item, pdf_type, pdf_doc, page_idx, scale
|
|
|
+ detection_image, item, pdf_type, pdf_doc, page_idx, scale,
|
|
|
+ pre_matched_spans=spans
|
|
|
)
|
|
|
processed_elements.append(element)
|
|
|
except Exception as e:
|
|
|
@@ -421,8 +478,10 @@ class EnhancedDocPipeline:
|
|
|
# 处理表格标题/脚注
|
|
|
for item in classified_elements['table_text']:
|
|
|
try:
|
|
|
+ spans = get_matched_spans_for_item(item)
|
|
|
element = self.element_processors.process_text_element(
|
|
|
- detection_image, item, pdf_type, pdf_doc, page_idx, scale
|
|
|
+ detection_image, item, pdf_type, pdf_doc, page_idx, scale,
|
|
|
+ pre_matched_spans=spans
|
|
|
)
|
|
|
processed_elements.append(element)
|
|
|
except Exception as e:
|
|
|
@@ -432,8 +491,10 @@ class EnhancedDocPipeline:
|
|
|
# 处理图片标题/脚注
|
|
|
for item in classified_elements['image_text']:
|
|
|
try:
|
|
|
+ spans = get_matched_spans_for_item(item)
|
|
|
element = self.element_processors.process_text_element(
|
|
|
- detection_image, item, pdf_type, pdf_doc, page_idx, scale
|
|
|
+ detection_image, item, pdf_type, pdf_doc, page_idx, scale,
|
|
|
+ pre_matched_spans=spans
|
|
|
)
|
|
|
processed_elements.append(element)
|
|
|
except Exception as e:
|
|
|
@@ -443,8 +504,9 @@ class EnhancedDocPipeline:
|
|
|
# 处理表格主体
|
|
|
for item in classified_elements['table_body']:
|
|
|
try:
|
|
|
+ spans = get_matched_spans_for_item(item)
|
|
|
element = self.element_processors.process_table_element(
|
|
|
- detection_image, item, scale
|
|
|
+ detection_image, item, scale, pre_matched_spans=spans
|
|
|
)
|
|
|
processed_elements.append(element)
|
|
|
except Exception as e:
|
|
|
@@ -476,8 +538,10 @@ class EnhancedDocPipeline:
|
|
|
# 处理代码元素
|
|
|
for item in classified_elements['code']:
|
|
|
try:
|
|
|
+ spans = get_matched_spans_for_item(item)
|
|
|
element = self.element_processors.process_code_element(
|
|
|
- detection_image, item, pdf_type, pdf_doc, page_idx, scale
|
|
|
+ detection_image, item, pdf_type, pdf_doc, page_idx, scale,
|
|
|
+ pre_matched_spans=spans
|
|
|
)
|
|
|
processed_elements.append(element)
|
|
|
except Exception as e:
|
|
|
@@ -487,8 +551,9 @@ class EnhancedDocPipeline:
|
|
|
# 处理丢弃元素
|
|
|
for item in classified_elements['discard']:
|
|
|
try:
|
|
|
+ spans = get_matched_spans_for_item(item)
|
|
|
element = self.element_processors.process_discard_element(
|
|
|
- detection_image, item
|
|
|
+ detection_image, item, pre_matched_spans=spans
|
|
|
)
|
|
|
discarded_elements.append(element)
|
|
|
except Exception as e:
|