|
|
@@ -17,26 +17,28 @@ except ImportError:
|
|
|
class PaddleLayoutDetector(BaseLayoutDetector):
|
|
|
"""PaddleX RT-DETR 布局检测器 (ONNX 版本)"""
|
|
|
|
|
|
- # ⚠️ 修正:使用官方的 RT-DETR-H_layout_17cls 类别定义
|
|
|
- # 映射到 MinerU 的类别体系
|
|
|
+ # 类别映射:PaddleX RT-DETR-H_layout_17cls → MinerU/EnhancedDocPipeline 类别体系
|
|
|
+ # 参考:
|
|
|
+ # - MinerU: mineru/utils/enum_class.py (BlockType, CategoryId)
|
|
|
+ # - Pipeline: universal_doc_parser/core/pipeline_manager_v2.py (EnhancedDocPipeline 类别定义)
|
|
|
CATEGORY_MAP = {
|
|
|
- 0: 'title', # paragraph_title -> title
|
|
|
- 1: 'image_body', # image -> image_body
|
|
|
- 2: 'text', # text -> text
|
|
|
- 3: 'text', # number -> text (合并到text)
|
|
|
- 4: 'text', # abstract -> text
|
|
|
- 5: 'text', # content -> text
|
|
|
- 6: 'image_caption', # figure_title -> image_caption
|
|
|
- 7: 'interline_equation', # formula -> interline_equation
|
|
|
- 8: 'table_body', # table -> table_body
|
|
|
- 9: 'table_caption', # table_title -> table_caption
|
|
|
- 10: 'text', # reference -> text
|
|
|
- 11: 'title', # doc_title -> title
|
|
|
- 12: 'table_footnote', # footnote -> table_footnote
|
|
|
- 13: 'abandon', # header -> abandon (页眉通常不需要)
|
|
|
- 14: 'text', # algorithm -> text
|
|
|
- 15: 'abandon', # footer -> abandon (页脚通常不需要)
|
|
|
- 16: 'abandon' # seal -> abandon (印章通常不需要)
|
|
|
+ 0: 'title', # paragraph_title -> title (TEXT_CATEGORIES)
|
|
|
+ 1: 'image_body', # image -> image_body (IMAGE_BODY_CATEGORIES)
|
|
|
+ 2: 'text', # text -> text (TEXT_CATEGORIES)
|
|
|
+ 3: 'text', # number -> text (TEXT_CATEGORIES)
|
|
|
+ 4: 'text', # abstract -> text (TEXT_CATEGORIES)
|
|
|
+ 5: 'text', # content -> text (TEXT_CATEGORIES)
|
|
|
+ 6: 'image_caption', # figure_title -> image_caption (IMAGE_TEXT_CATEGORIES)
|
|
|
+ 7: 'interline_equation', # formula -> interline_equation (EQUATION_CATEGORIES)
|
|
|
+ 8: 'table_body', # table -> table_body (TABLE_BODY_CATEGORIES)
|
|
|
+ 9: 'table_caption', # table_title -> table_caption (TABLE_TEXT_CATEGORIES)
|
|
|
+ 10: 'ref_text', # reference -> ref_text (TEXT_CATEGORIES)
|
|
|
+ 11: 'title', # doc_title -> title (TEXT_CATEGORIES)
|
|
|
+ 12: 'page_footnote', # footnote -> page_footnote (TEXT_CATEGORIES)
|
|
|
+ 13: 'header', # header -> header (TEXT_CATEGORIES)
|
|
|
+ 14: 'algorithm', # algorithm -> algorithm (CODE_CATEGORIES)
|
|
|
+ 15: 'footer', # footer -> footer (TEXT_CATEGORIES)
|
|
|
+ 16: 'abandon' # seal -> abandon (DISCARD_CATEGORIES)
|
|
|
}
|
|
|
|
|
|
ORIGINAL_CATEGORY_NAMES = {
|
|
|
@@ -434,17 +436,29 @@ class PaddleLayoutDetector(BaseLayoutDetector):
|
|
|
# 为每个类别分配固定颜色
|
|
|
category_colors = {}
|
|
|
|
|
|
- # 预定义一些常用类别的颜色
|
|
|
+ # 预定义类别颜色(与 EnhancedDocPipeline 类别定义保持一致)
|
|
|
predefined_colors = {
|
|
|
- 'text': (0, 255, 0), # 绿色
|
|
|
- 'title': (255, 0, 0), # 红色
|
|
|
- 'table_body': (0, 0, 255), # 蓝色
|
|
|
- 'table_caption': (255, 255, 0), # 青色
|
|
|
- 'table_footnote': (255, 128, 0), # 橙色
|
|
|
- 'image_body': (255, 0, 255), # 洋红
|
|
|
- 'image_caption': (128, 0, 255), # 紫色
|
|
|
- 'interline_equation': (0, 255, 255), # 黄色
|
|
|
- 'abandon': (128, 128, 128), # 灰色
|
|
|
+ # 文本类
|
|
|
+ 'text': (153, 0, 76), # 深红
|
|
|
+ 'title': (102, 102, 255), # 蓝色
|
|
|
+ 'header': (128, 128, 128), # 灰色
|
|
|
+ 'footer': (128, 128, 128), # 灰色
|
|
|
+ 'ref_text': (180, 180, 180), # 浅灰
|
|
|
+ 'page_footnote': (200, 200, 200), # 浅灰
|
|
|
+ # 表格类
|
|
|
+ 'table_body': (204, 204, 0), # 黄色
|
|
|
+ 'table_caption': (255, 255, 102), # 浅黄
|
|
|
+ 'table_footnote': (229, 255, 204), # 浅黄绿
|
|
|
+ # 图片类
|
|
|
+ 'image_body': (153, 255, 51), # 绿色
|
|
|
+ 'image_caption': (102, 178, 255), # 浅蓝
|
|
|
+ # 公式类
|
|
|
+ 'interline_equation': (0, 255, 0), # 亮绿
|
|
|
+ # 代码类
|
|
|
+ 'algorithm': (128, 0, 255), # 紫色
|
|
|
+ 'code': (102, 0, 204), # 紫色
|
|
|
+ # 丢弃类
|
|
|
+ 'abandon': (100, 100, 100), # 深灰
|
|
|
}
|
|
|
|
|
|
# 过滤低置信度结果
|
|
|
@@ -623,7 +637,7 @@ if __name__ == "__main__":
|
|
|
detector.initialize()
|
|
|
|
|
|
# 读取测试图像
|
|
|
- img_path = "/Users/zhch158/workspace/data/流水分析/B用户_扫描流水/PaddleOCR_VL_Results/B用户_扫描流水/B用户_扫描流水_page_001.png"
|
|
|
+ img_path = "/Users/zhch158/workspace/data/流水分析/康强_北京农村商业银行/ppstructurev3_client_results/康强_北京农村商业银行/康强_北京农村商业银行_page_001.png"
|
|
|
print(f"\n📖 Loading image: {img_path}")
|
|
|
img = cv2.imread(img_path)
|
|
|
|