|
@@ -62,6 +62,39 @@ except ImportError as e:
|
|
|
class EnhancedDocPipeline:
|
|
class EnhancedDocPipeline:
|
|
|
"""增强版文档处理流水线"""
|
|
"""增强版文档处理流水线"""
|
|
|
|
|
|
|
|
|
|
+ # 元素类别分组(参考 MinerU CategoryId 定义)
|
|
|
|
|
+ # CategoryId: 0=Title, 1=Text, 2=Abandon, 3=ImageBody, 4=ImageCaption,
|
|
|
|
|
+ # 5=TableBody, 6=TableCaption, 7=TableFootnote, 8=InterlineEquation_Layout,
|
|
|
|
|
+ # 9=InterlineEquationNumber_Layout, 13=InlineEquation, 14=InterlineEquation_YOLO,
|
|
|
|
|
+ # 15=OcrText, 16=LowScoreText, 101=ImageFootnote
|
|
|
|
|
+
|
|
|
|
|
+ # 文本类元素:需要OCR识别或PDF文本提取
|
|
|
|
|
+ TEXT_CATEGORIES = [
|
|
|
|
|
+ 'text', 'title', 'ocr_text', 'low_score_text', # 基础文本
|
|
|
|
|
+ 'header', 'footer', 'page_number', # 页面元素
|
|
|
|
|
+ 'ref_text', 'aside_text', 'page_footnote', # 辅助文本
|
|
|
|
|
+ ]
|
|
|
|
|
+
|
|
|
|
|
+ # 表格相关元素
|
|
|
|
|
+ TABLE_BODY_CATEGORIES = ['table', 'table_body'] # 表格主体:需要OCR+VLM处理
|
|
|
|
|
+ TABLE_TEXT_CATEGORIES = ['table_caption', 'table_footnote'] # 表格附属文本
|
|
|
|
|
+
|
|
|
|
|
+ # 图片相关元素
|
|
|
|
|
+ IMAGE_BODY_CATEGORIES = ['image', 'image_body', 'figure'] # 图片主体
|
|
|
|
|
+ IMAGE_TEXT_CATEGORIES = ['image_caption', 'image_footnote'] # 图片附属文本
|
|
|
|
|
+
|
|
|
|
|
+ # 公式类元素
|
|
|
|
|
+ EQUATION_CATEGORIES = [
|
|
|
|
|
+ 'interline_equation', 'inline_equation', 'equation',
|
|
|
|
|
+ 'interline_equation_yolo', 'interline_equation_number'
|
|
|
|
|
+ ]
|
|
|
|
|
+
|
|
|
|
|
+ # 丢弃类元素(水印、装饰等)
|
|
|
|
|
+ DISCARD_CATEGORIES = ['abandon', 'discarded']
|
|
|
|
|
+
|
|
|
|
|
+ # 代码类元素
|
|
|
|
|
+ CODE_CATEGORIES = ['code', 'code_body', 'code_caption', 'algorithm']
|
|
|
|
|
+
|
|
|
def __init__(self, config_path: str):
|
|
def __init__(self, config_path: str):
|
|
|
"""
|
|
"""
|
|
|
初始化流水线
|
|
初始化流水线
|
|
@@ -311,71 +344,154 @@ class EnhancedDocPipeline:
|
|
|
|
|
|
|
|
page_result['layout_raw'] = layout_results
|
|
page_result['layout_raw'] = layout_results
|
|
|
|
|
|
|
|
- # 3. 分类元素
|
|
|
|
|
- text_elements = []
|
|
|
|
|
- table_elements = []
|
|
|
|
|
- other_elements = []
|
|
|
|
|
|
|
+ # 3. 分类元素(参考 MinerU 的分类策略)
|
|
|
|
|
+ text_elements = [] # 纯文本元素
|
|
|
|
|
+ table_body_elements = [] # 表格主体
|
|
|
|
|
+ table_text_elements = [] # 表格标题/脚注
|
|
|
|
|
+ image_body_elements = [] # 图片主体
|
|
|
|
|
+ image_text_elements = [] # 图片标题/脚注
|
|
|
|
|
+ equation_elements = [] # 公式
|
|
|
|
|
+ code_elements = [] # 代码
|
|
|
|
|
+ discard_elements = [] # 丢弃元素
|
|
|
|
|
|
|
|
for item in layout_results:
|
|
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']:
|
|
|
|
|
|
|
+ category = item.get('category', '').lower()
|
|
|
|
|
+
|
|
|
|
|
+ if category in self.TEXT_CATEGORIES:
|
|
|
text_elements.append(item)
|
|
text_elements.append(item)
|
|
|
- elif category in ['table', 'table_body']:
|
|
|
|
|
- table_elements.append(item)
|
|
|
|
|
|
|
+ elif category in self.TABLE_BODY_CATEGORIES:
|
|
|
|
|
+ table_body_elements.append(item)
|
|
|
|
|
+ elif category in self.TABLE_TEXT_CATEGORIES:
|
|
|
|
|
+ table_text_elements.append(item)
|
|
|
|
|
+ elif category in self.IMAGE_BODY_CATEGORIES:
|
|
|
|
|
+ image_body_elements.append(item)
|
|
|
|
|
+ elif category in self.IMAGE_TEXT_CATEGORIES:
|
|
|
|
|
+ image_text_elements.append(item)
|
|
|
|
|
+ elif category in self.EQUATION_CATEGORIES:
|
|
|
|
|
+ equation_elements.append(item)
|
|
|
|
|
+ elif category in self.CODE_CATEGORIES:
|
|
|
|
|
+ code_elements.append(item)
|
|
|
|
|
+ elif category in self.DISCARD_CATEGORIES:
|
|
|
|
|
+ discard_elements.append(item)
|
|
|
else:
|
|
else:
|
|
|
- other_elements.append(item)
|
|
|
|
|
|
|
+ # 未知类别归入文本处理
|
|
|
|
|
+ logger.debug(f"Unknown category '{category}', treating as text")
|
|
|
|
|
+ text_elements.append(item)
|
|
|
|
|
|
|
|
- # 4. 并行处理文本和表格(在旋转后的图片上进行识别)
|
|
|
|
|
|
|
+ logger.info(f"📊 Page {page_idx} elements: "
|
|
|
|
|
+ f"text={len(text_elements)}, "
|
|
|
|
|
+ f"table={len(table_body_elements)}, "
|
|
|
|
|
+ f"table_text={len(table_text_elements)}, "
|
|
|
|
|
+ f"image={len(image_body_elements)}, "
|
|
|
|
|
+ f"image_text={len(image_text_elements)}, "
|
|
|
|
|
+ f"equation={len(equation_elements)}, "
|
|
|
|
|
+ f"code={len(code_elements)}, "
|
|
|
|
|
+ f"discard={len(discard_elements)}")
|
|
|
|
|
+
|
|
|
|
|
+ # 4. 并行处理各类元素(在旋转后的图片上进行识别)
|
|
|
processed_elements = []
|
|
processed_elements = []
|
|
|
|
|
+ discarded_elements_result = [] # 单独存储丢弃元素
|
|
|
|
|
+
|
|
|
|
|
+ # 辅助函数:处理元素并转换坐标
|
|
|
|
|
+ def process_and_transform(item, processor_func, *args):
|
|
|
|
|
+ element = processor_func(detection_image, item, *args)
|
|
|
|
|
+ if rotate_angle != 0:
|
|
|
|
|
+ element = self._transform_coords_to_original(
|
|
|
|
|
+ element, rotate_angle, detection_image.shape, original_image.shape
|
|
|
|
|
+ )
|
|
|
|
|
+ return element
|
|
|
|
|
|
|
|
- # 4.1 处理文本区域
|
|
|
|
|
|
|
+ # 4.1 处理纯文本区域
|
|
|
for text_item in text_elements:
|
|
for text_item in text_elements:
|
|
|
try:
|
|
try:
|
|
|
- element = self._process_text_element(
|
|
|
|
|
- detection_image, text_item, pdf_type, pdf_doc, page_idx, scale
|
|
|
|
|
|
|
+ element = process_and_transform(
|
|
|
|
|
+ text_item, self._process_text_element, 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)
|
|
processed_elements.append(element)
|
|
|
except Exception as e:
|
|
except Exception as e:
|
|
|
logger.warning(f"⚠️ Text processing failed: {e}")
|
|
logger.warning(f"⚠️ Text processing failed: {e}")
|
|
|
processed_elements.append(self._create_error_element(text_item, str(e)))
|
|
processed_elements.append(self._create_error_element(text_item, str(e)))
|
|
|
|
|
|
|
|
- # 4.2 处理表格区域(OCR检测 + VLM结构识别 + 坐标匹配)
|
|
|
|
|
- for table_item in table_elements:
|
|
|
|
|
|
|
+ # 4.2 处理表格标题/脚注(作为文本处理)
|
|
|
|
|
+ for table_text_item in table_text_elements:
|
|
|
try:
|
|
try:
|
|
|
- element = self._process_table_element(
|
|
|
|
|
- detection_image, table_item, scale
|
|
|
|
|
|
|
+ element = process_and_transform(
|
|
|
|
|
+ table_text_item, self._process_text_element, pdf_type, pdf_doc, page_idx, scale
|
|
|
|
|
+ )
|
|
|
|
|
+ processed_elements.append(element)
|
|
|
|
|
+ except Exception as e:
|
|
|
|
|
+ logger.warning(f"⚠️ Table text processing failed: {e}")
|
|
|
|
|
+ processed_elements.append(self._create_error_element(table_text_item, str(e)))
|
|
|
|
|
+
|
|
|
|
|
+ # 4.3 处理图片标题/脚注(作为文本处理)
|
|
|
|
|
+ for image_text_item in image_text_elements:
|
|
|
|
|
+ try:
|
|
|
|
|
+ element = process_and_transform(
|
|
|
|
|
+ image_text_item, self._process_text_element, pdf_type, pdf_doc, page_idx, scale
|
|
|
|
|
+ )
|
|
|
|
|
+ processed_elements.append(element)
|
|
|
|
|
+ except Exception as e:
|
|
|
|
|
+ logger.warning(f"⚠️ Image text processing failed: {e}")
|
|
|
|
|
+ processed_elements.append(self._create_error_element(image_text_item, str(e)))
|
|
|
|
|
+
|
|
|
|
|
+ # 4.4 处理表格主体(OCR检测 + VLM结构识别 + 坐标匹配)
|
|
|
|
|
+ for table_item in table_body_elements:
|
|
|
|
|
+ try:
|
|
|
|
|
+ element = process_and_transform(
|
|
|
|
|
+ table_item, self._process_table_element, scale
|
|
|
)
|
|
)
|
|
|
- # **关键**: 将坐标转换回原始图片坐标系
|
|
|
|
|
- if rotate_angle != 0:
|
|
|
|
|
- element = self._transform_coords_to_original(
|
|
|
|
|
- element, rotate_angle, detection_image.shape, original_image.shape
|
|
|
|
|
- )
|
|
|
|
|
processed_elements.append(element)
|
|
processed_elements.append(element)
|
|
|
except Exception as e:
|
|
except Exception as e:
|
|
|
logger.warning(f"⚠️ Table processing failed: {e}")
|
|
logger.warning(f"⚠️ Table processing failed: {e}")
|
|
|
processed_elements.append(self._create_error_element(table_item, str(e)))
|
|
processed_elements.append(self._create_error_element(table_item, str(e)))
|
|
|
|
|
|
|
|
- # 4.3 处理其他元素
|
|
|
|
|
- for other_item in other_elements:
|
|
|
|
|
|
|
+ # 4.5 处理公式元素
|
|
|
|
|
+ for equation_item in equation_elements:
|
|
|
try:
|
|
try:
|
|
|
- element = self._process_other_element(detection_image, other_item)
|
|
|
|
|
- # **关键**: 将坐标转换回原始图片坐标系
|
|
|
|
|
|
|
+ element = process_and_transform(
|
|
|
|
|
+ equation_item, self._process_equation_element
|
|
|
|
|
+ )
|
|
|
|
|
+ processed_elements.append(element)
|
|
|
|
|
+ except Exception as e:
|
|
|
|
|
+ logger.warning(f"⚠️ Equation processing failed: {e}")
|
|
|
|
|
+ processed_elements.append(self._create_error_element(equation_item, str(e)))
|
|
|
|
|
+
|
|
|
|
|
+ # 4.6 处理图片主体
|
|
|
|
|
+ for image_item in image_body_elements:
|
|
|
|
|
+ try:
|
|
|
|
|
+ element = process_and_transform(
|
|
|
|
|
+ image_item, self._process_image_element
|
|
|
|
|
+ )
|
|
|
|
|
+ processed_elements.append(element)
|
|
|
|
|
+ except Exception as e:
|
|
|
|
|
+ logger.warning(f"⚠️ Image processing failed: {e}")
|
|
|
|
|
+ processed_elements.append(self._create_error_element(image_item, str(e)))
|
|
|
|
|
+
|
|
|
|
|
+ # 4.7 处理代码元素(作为文本处理)
|
|
|
|
|
+ for code_item in code_elements:
|
|
|
|
|
+ try:
|
|
|
|
|
+ element = process_and_transform(
|
|
|
|
|
+ code_item, self._process_code_element, pdf_type, pdf_doc, page_idx, scale
|
|
|
|
|
+ )
|
|
|
|
|
+ processed_elements.append(element)
|
|
|
|
|
+ except Exception as e:
|
|
|
|
|
+ logger.warning(f"⚠️ Code processing failed: {e}")
|
|
|
|
|
+ processed_elements.append(self._create_error_element(code_item, str(e)))
|
|
|
|
|
+
|
|
|
|
|
+ # 4.8 处理丢弃元素(仅记录,不做OCR)
|
|
|
|
|
+ for discard_item in discard_elements:
|
|
|
|
|
+ try:
|
|
|
|
|
+ element = self._process_discard_element(discard_item)
|
|
|
if rotate_angle != 0:
|
|
if rotate_angle != 0:
|
|
|
element = self._transform_coords_to_original(
|
|
element = self._transform_coords_to_original(
|
|
|
element, rotate_angle, detection_image.shape, original_image.shape
|
|
element, rotate_angle, detection_image.shape, original_image.shape
|
|
|
)
|
|
)
|
|
|
- processed_elements.append(element)
|
|
|
|
|
|
|
+ discarded_elements_result.append(element)
|
|
|
except Exception as e:
|
|
except Exception as e:
|
|
|
- logger.warning(f"⚠️ Element processing failed: {e}")
|
|
|
|
|
- processed_elements.append(self._create_error_element(other_item, str(e)))
|
|
|
|
|
|
|
+ logger.debug(f"Discard element processing failed: {e}")
|
|
|
|
|
|
|
|
page_result['elements'] = processed_elements
|
|
page_result['elements'] = processed_elements
|
|
|
|
|
+ page_result['discarded_blocks'] = discarded_elements_result # MinerU 兼容字段
|
|
|
return page_result
|
|
return page_result
|
|
|
|
|
|
|
|
def _process_text_element(
|
|
def _process_text_element(
|
|
@@ -541,11 +657,20 @@ class EnhancedDocPipeline:
|
|
|
orig_table_size=orig_table_size,
|
|
orig_table_size=orig_table_size,
|
|
|
table_bbox=bbox
|
|
table_bbox=bbox
|
|
|
)
|
|
)
|
|
|
|
|
+ # ocr_boxes 也需要逆向旋转并加上表格偏移量
|
|
|
|
|
+ ocr_boxes = self._inverse_rotate_ocr_boxes(
|
|
|
|
|
+ ocr_boxes=ocr_boxes,
|
|
|
|
|
+ rotation_angle=table_angle,
|
|
|
|
|
+ orig_table_size=orig_table_size,
|
|
|
|
|
+ table_bbox=bbox
|
|
|
|
|
+ )
|
|
|
logger.info(f"📐 Coordinates transformed back to original image")
|
|
logger.info(f"📐 Coordinates transformed back to original image")
|
|
|
else:
|
|
else:
|
|
|
# 没有旋转,只需要加上表格偏移量转换为整页坐标
|
|
# 没有旋转,只需要加上表格偏移量转换为整页坐标
|
|
|
cells = self._add_table_offset_to_cells(cells, bbox)
|
|
cells = self._add_table_offset_to_cells(cells, bbox)
|
|
|
enhanced_html = self._add_table_offset_to_html(enhanced_html, bbox)
|
|
enhanced_html = self._add_table_offset_to_html(enhanced_html, bbox)
|
|
|
|
|
+ # ocr_boxes 也需要加上表格偏移量
|
|
|
|
|
+ ocr_boxes = self._add_table_offset_to_ocr_boxes(ocr_boxes, bbox)
|
|
|
|
|
|
|
|
return {
|
|
return {
|
|
|
'type': 'table',
|
|
'type': 'table',
|
|
@@ -561,36 +686,26 @@ class EnhancedDocPipeline:
|
|
|
},
|
|
},
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
- def _process_other_element(
|
|
|
|
|
|
|
+ def _process_equation_element(
|
|
|
self,
|
|
self,
|
|
|
image: np.ndarray,
|
|
image: np.ndarray,
|
|
|
layout_item: Dict[str, Any]
|
|
layout_item: Dict[str, Any]
|
|
|
) -> Dict[str, Any]:
|
|
) -> Dict[str, Any]:
|
|
|
- """处理其他类型元素(公式、图片等)"""
|
|
|
|
|
|
|
+ """处理公式元素"""
|
|
|
bbox = layout_item.get('bbox', [0, 0, 0, 0])
|
|
bbox = layout_item.get('bbox', [0, 0, 0, 0])
|
|
|
category = layout_item.get('category', '')
|
|
category = layout_item.get('category', '')
|
|
|
cropped_region = self._crop_region(image, bbox)
|
|
cropped_region = self._crop_region(image, bbox)
|
|
|
|
|
|
|
|
- content = {}
|
|
|
|
|
|
|
+ content = {'latex': '', 'confidence': 0.0}
|
|
|
|
|
|
|
|
- # 公式识别
|
|
|
|
|
- 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']:
|
|
|
|
|
|
|
+ try:
|
|
|
|
|
+ formula_result = self.vl_recognizer.recognize_formula(cropped_region)
|
|
|
content = {
|
|
content = {
|
|
|
- 'type': 'image',
|
|
|
|
|
- 'description': ''
|
|
|
|
|
|
|
+ 'latex': formula_result.get('latex', ''),
|
|
|
|
|
+ 'confidence': formula_result.get('confidence', 0.0)
|
|
|
}
|
|
}
|
|
|
|
|
+ except Exception as e:
|
|
|
|
|
+ logger.warning(f"Formula recognition failed: {e}")
|
|
|
|
|
|
|
|
return {
|
|
return {
|
|
|
'type': category,
|
|
'type': category,
|
|
@@ -599,6 +714,111 @@ class EnhancedDocPipeline:
|
|
|
'content': content
|
|
'content': content
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
|
|
+ def _process_image_element(
|
|
|
|
|
+ self,
|
|
|
|
|
+ image: np.ndarray,
|
|
|
|
|
+ layout_item: Dict[str, Any]
|
|
|
|
|
+ ) -> Dict[str, Any]:
|
|
|
|
|
+ """
|
|
|
|
|
+ 处理图片元素
|
|
|
|
|
+
|
|
|
|
|
+ 裁剪图片区域,返回图片信息(实际保存由 OutputFormatter 处理)
|
|
|
|
|
+ """
|
|
|
|
|
+ bbox = layout_item.get('bbox', [0, 0, 0, 0])
|
|
|
|
|
+ category = layout_item.get('category', '')
|
|
|
|
|
+
|
|
|
|
|
+ # 裁剪图片区域
|
|
|
|
|
+ cropped_image = self._crop_region(image, bbox)
|
|
|
|
|
+
|
|
|
|
|
+ return {
|
|
|
|
|
+ 'type': category,
|
|
|
|
|
+ 'bbox': bbox,
|
|
|
|
|
+ 'confidence': layout_item.get('confidence', 0.0),
|
|
|
|
|
+ 'content': {
|
|
|
|
|
+ 'type': 'image',
|
|
|
|
|
+ 'description': '',
|
|
|
|
|
+ 'image_data': cropped_image, # 保存裁剪的图片数据,由 OutputFormatter 保存
|
|
|
|
|
+ 'image_path': '' # 由 OutputFormatter 填充实际路径
|
|
|
|
|
+ }
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ def _process_code_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]:
|
|
|
|
|
+ """
|
|
|
|
|
+ 处理代码元素
|
|
|
|
|
+
|
|
|
|
|
+ 代码块作为特殊文本处理,保留格式
|
|
|
|
|
+ """
|
|
|
|
|
+ bbox = layout_item.get('bbox', [0, 0, 0, 0])
|
|
|
|
|
+ category = layout_item.get('category', '')
|
|
|
|
|
+ cropped_image = self._crop_region(image, bbox)
|
|
|
|
|
+
|
|
|
|
|
+ code_content = ""
|
|
|
|
|
+ ocr_details = []
|
|
|
|
|
+
|
|
|
|
|
+ # 优先从PDF提取(保留格式更好)
|
|
|
|
|
+ if pdf_type == 'txt' and pdf_doc is not None:
|
|
|
|
|
+ code_content, success = self._extract_text_from_pdf(
|
|
|
|
|
+ pdf_doc, page_idx, bbox, scale
|
|
|
|
|
+ )
|
|
|
|
|
+ if not success:
|
|
|
|
|
+ code_content = ""
|
|
|
|
|
+
|
|
|
|
|
+ # PDF提取失败或扫描件,使用OCR
|
|
|
|
|
+ if not code_content:
|
|
|
|
|
+ try:
|
|
|
|
|
+ ocr_results = self.ocr_recognizer.recognize_text(cropped_image)
|
|
|
|
|
+ if ocr_results:
|
|
|
|
|
+ # 代码保留换行
|
|
|
|
|
+ lines = []
|
|
|
|
|
+ for item in ocr_results:
|
|
|
|
|
+ lines.append(item.get('text', ''))
|
|
|
|
|
+ ocr_details.append({
|
|
|
|
|
+ 'text': item.get('text', ''),
|
|
|
|
|
+ 'bbox': item.get('bbox', []),
|
|
|
|
|
+ 'confidence': item.get('confidence', 0.0)
|
|
|
|
|
+ })
|
|
|
|
|
+ code_content = '\n'.join(lines)
|
|
|
|
|
+ except Exception as e:
|
|
|
|
|
+ logger.warning(f"Code OCR failed: {e}")
|
|
|
|
|
+
|
|
|
|
|
+ return {
|
|
|
|
|
+ 'type': category,
|
|
|
|
|
+ 'bbox': bbox,
|
|
|
|
|
+ 'confidence': layout_item.get('confidence', 0.0),
|
|
|
|
|
+ 'content': {
|
|
|
|
|
+ 'code': code_content,
|
|
|
|
|
+ 'language': '', # 可以后续添加语言检测
|
|
|
|
|
+ 'ocr_details': ocr_details
|
|
|
|
|
+ }
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ def _process_discard_element(
|
|
|
|
|
+ self,
|
|
|
|
|
+ layout_item: Dict[str, Any]
|
|
|
|
|
+ ) -> Dict[str, Any]:
|
|
|
|
|
+ """
|
|
|
|
|
+ 处理丢弃元素(水印、装饰等)
|
|
|
|
|
+
|
|
|
|
|
+ 仅记录位置,不做OCR识别
|
|
|
|
|
+ """
|
|
|
|
|
+ bbox = layout_item.get('bbox', [0, 0, 0, 0])
|
|
|
|
|
+ category = layout_item.get('category', 'abandon')
|
|
|
|
|
+
|
|
|
|
|
+ return {
|
|
|
|
|
+ 'type': 'discarded',
|
|
|
|
|
+ 'bbox': bbox,
|
|
|
|
|
+ 'confidence': layout_item.get('confidence', 0.0),
|
|
|
|
|
+ 'original_category': category
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
def _extract_text_from_pdf(
|
|
def _extract_text_from_pdf(
|
|
|
self,
|
|
self,
|
|
|
pdf_doc: Any,
|
|
pdf_doc: Any,
|
|
@@ -913,6 +1133,129 @@ class EnhancedDocPipeline:
|
|
|
|
|
|
|
|
return converted_html
|
|
return converted_html
|
|
|
|
|
|
|
|
|
|
+ def _add_table_offset_to_ocr_boxes(
|
|
|
|
|
+ self,
|
|
|
|
|
+ ocr_boxes: List[Dict],
|
|
|
|
|
+ table_bbox: List[float]
|
|
|
|
|
+ ) -> List[Dict]:
|
|
|
|
|
+ """
|
|
|
|
|
+ 为 ocr_boxes 添加表格偏移量,将相对坐标转换为页面绝对坐标
|
|
|
|
|
+
|
|
|
|
|
+ Args:
|
|
|
|
|
+ ocr_boxes: OCR 框列表
|
|
|
|
|
+ table_bbox: 表格在页面中的位置 [x1, y1, x2, y2]
|
|
|
|
|
+
|
|
|
|
|
+ Returns:
|
|
|
|
|
+ 转换后的 ocr_boxes
|
|
|
|
|
+ """
|
|
|
|
|
+ if not ocr_boxes or not table_bbox:
|
|
|
|
|
+ return ocr_boxes
|
|
|
|
|
+
|
|
|
|
|
+ offset_x = table_bbox[0]
|
|
|
|
|
+ offset_y = table_bbox[1]
|
|
|
|
|
+
|
|
|
|
|
+ converted_boxes = []
|
|
|
|
|
+ for box in ocr_boxes:
|
|
|
|
|
+ new_box = box.copy()
|
|
|
|
|
+
|
|
|
|
|
+ # 转换 bbox [x1, y1, x2, y2]
|
|
|
|
|
+ if 'bbox' in new_box and new_box['bbox']:
|
|
|
|
|
+ bbox = new_box['bbox']
|
|
|
|
|
+ if len(bbox) >= 4:
|
|
|
|
|
+ new_box['bbox'] = [
|
|
|
|
|
+ bbox[0] + offset_x,
|
|
|
|
|
+ bbox[1] + offset_y,
|
|
|
|
|
+ bbox[2] + offset_x,
|
|
|
|
|
+ bbox[3] + offset_y
|
|
|
|
|
+ ]
|
|
|
|
|
+
|
|
|
|
|
+ # 转换 poly [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
|
|
|
|
|
+ if 'poly' in new_box and new_box['poly']:
|
|
|
|
|
+ poly = new_box['poly']
|
|
|
|
|
+ new_poly = []
|
|
|
|
|
+ for point in poly:
|
|
|
|
|
+ if isinstance(point, (list, tuple)) and len(point) >= 2:
|
|
|
|
|
+ new_poly.append([point[0] + offset_x, point[1] + offset_y])
|
|
|
|
|
+ else:
|
|
|
|
|
+ new_poly.append(point)
|
|
|
|
|
+ new_box['poly'] = new_poly
|
|
|
|
|
+
|
|
|
|
|
+ converted_boxes.append(new_box)
|
|
|
|
|
+
|
|
|
|
|
+ return converted_boxes
|
|
|
|
|
+
|
|
|
|
|
+ def _inverse_rotate_ocr_boxes(
|
|
|
|
|
+ self,
|
|
|
|
|
+ ocr_boxes: List[Dict],
|
|
|
|
|
+ rotation_angle: float,
|
|
|
|
|
+ orig_table_size: Tuple[int, int],
|
|
|
|
|
+ table_bbox: List[float]
|
|
|
|
|
+ ) -> List[Dict]:
|
|
|
|
|
+ """
|
|
|
|
|
+ 对 ocr_boxes 进行逆向旋转并添加表格偏移量
|
|
|
|
|
+
|
|
|
|
|
+ Args:
|
|
|
|
|
+ ocr_boxes: OCR 框列表
|
|
|
|
|
+ rotation_angle: 表格旋转角度
|
|
|
|
|
+ orig_table_size: 原始表格尺寸 (width, height)
|
|
|
|
|
+ table_bbox: 表格在页面中的位置
|
|
|
|
|
+
|
|
|
|
|
+ Returns:
|
|
|
|
|
+ 转换后的 ocr_boxes
|
|
|
|
|
+ """
|
|
|
|
|
+ if not ocr_boxes:
|
|
|
|
|
+ return ocr_boxes
|
|
|
|
|
+
|
|
|
|
|
+ if not MERGER_AVAILABLE:
|
|
|
|
|
+ # 如果 merger 不可用,只添加偏移量
|
|
|
|
|
+ return self._add_table_offset_to_ocr_boxes(ocr_boxes, table_bbox)
|
|
|
|
|
+
|
|
|
|
|
+ offset_x = table_bbox[0]
|
|
|
|
|
+ offset_y = table_bbox[1]
|
|
|
|
|
+
|
|
|
|
|
+ converted_boxes = []
|
|
|
|
|
+ for box in ocr_boxes:
|
|
|
|
|
+ new_box = box.copy()
|
|
|
|
|
+
|
|
|
|
|
+ # 逆向旋转 bbox
|
|
|
|
|
+ if 'bbox' in new_box and new_box['bbox']:
|
|
|
|
|
+ bbox = new_box['bbox']
|
|
|
|
|
+ if len(bbox) >= 4:
|
|
|
|
|
+ try:
|
|
|
|
|
+ rotated_bbox = BBoxExtractor.inverse_rotate_box_coordinates(
|
|
|
|
|
+ bbox, rotation_angle, orig_table_size
|
|
|
|
|
+ )
|
|
|
|
|
+ # 添加表格偏移量
|
|
|
|
|
+ new_box['bbox'] = [
|
|
|
|
|
+ rotated_bbox[0] + offset_x,
|
|
|
|
|
+ rotated_bbox[1] + offset_y,
|
|
|
|
|
+ rotated_bbox[2] + offset_x,
|
|
|
|
|
+ rotated_bbox[3] + offset_y
|
|
|
|
|
+ ]
|
|
|
|
|
+ except Exception as e:
|
|
|
|
|
+ logger.debug(f"Failed to inverse rotate ocr_box bbox: {e}")
|
|
|
|
|
+
|
|
|
|
|
+ # 逆向旋转 poly
|
|
|
|
|
+ if 'poly' in new_box and new_box['poly']:
|
|
|
|
|
+ poly = new_box['poly']
|
|
|
|
|
+ try:
|
|
|
|
|
+ # 将 poly 转换为 list 格式 [[x1,y1], [x2,y2], ...]
|
|
|
|
|
+ poly_list = [[float(p[0]), float(p[1])] for p in poly]
|
|
|
|
|
+ rotated_poly = BBoxExtractor.inverse_rotate_coordinates(
|
|
|
|
|
+ poly_list, rotation_angle, orig_table_size
|
|
|
|
|
+ )
|
|
|
|
|
+ # 添加表格偏移量
|
|
|
|
|
+ new_poly = []
|
|
|
|
|
+ for point in rotated_poly:
|
|
|
|
|
+ new_poly.append([point[0] + offset_x, point[1] + offset_y])
|
|
|
|
|
+ new_box['poly'] = new_poly
|
|
|
|
|
+ except Exception as e:
|
|
|
|
|
+ logger.debug(f"Failed to inverse rotate ocr_box poly: {e}")
|
|
|
|
|
+
|
|
|
|
|
+ converted_boxes.append(new_box)
|
|
|
|
|
+
|
|
|
|
|
+ return converted_boxes
|
|
|
|
|
+
|
|
|
def _merge_cross_page_tables(self, results: Dict[str, Any]) -> Dict[str, Any]:
|
|
def _merge_cross_page_tables(self, results: Dict[str, Any]) -> Dict[str, Any]:
|
|
|
"""合并跨页表格"""
|
|
"""合并跨页表格"""
|
|
|
# TODO: 实现跨页表格合并逻辑
|
|
# TODO: 实现跨页表格合并逻辑
|
|
@@ -988,12 +1331,31 @@ class EnhancedDocPipeline:
|
|
|
if ocr_details:
|
|
if ocr_details:
|
|
|
for detail in ocr_details:
|
|
for detail in ocr_details:
|
|
|
if 'bbox' in detail and detail['bbox']:
|
|
if 'bbox' in detail and detail['bbox']:
|
|
|
- detail['bbox'] = BBoxExtractor.inverse_rotate_box_coordinates(
|
|
|
|
|
- detail.get('bbox', []), rotate_angle, orig_image_size
|
|
|
|
|
|
|
+ if self._is_poly_format(detail['bbox']):
|
|
|
|
|
+ detail['bbox'] = BBoxExtractor.inverse_rotate_coordinates(
|
|
|
|
|
+ detail['bbox'], rotate_angle, orig_image_size
|
|
|
|
|
+ )
|
|
|
|
|
+ else:
|
|
|
|
|
+ detail['bbox'] = BBoxExtractor.inverse_rotate_box_coordinates(
|
|
|
|
|
+ detail.get('bbox', []), rotate_angle, orig_image_size
|
|
|
)
|
|
)
|
|
|
|
|
|
|
|
return element
|
|
return element
|
|
|
|
|
|
|
|
|
|
+ def _is_poly_format(self, bbox: Any) -> bool:
|
|
|
|
|
+ """
|
|
|
|
|
+ 检测 bbox 是否为四点多边形格式
|
|
|
|
|
+
|
|
|
|
|
+ 四点格式: [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
|
|
|
|
|
+ 矩形格式: [x_min, y_min, x_max, y_max]
|
|
|
|
|
+ """
|
|
|
|
|
+ if not bbox or not isinstance(bbox, list):
|
|
|
|
|
+ return False
|
|
|
|
|
+ if len(bbox) != 4:
|
|
|
|
|
+ return False
|
|
|
|
|
+ # 如果第一个元素是列表,则为四点格式
|
|
|
|
|
+ return isinstance(bbox[0], (list, tuple))
|
|
|
|
|
+
|
|
|
def _transform_html_data_bbox(
|
|
def _transform_html_data_bbox(
|
|
|
self,
|
|
self,
|
|
|
html: str,
|
|
html: str,
|