Bladeren bron

refactor: update category mapping and color definitions in PaddleLayoutDetector for improved clarity and consistency with EnhancedDocPipeline

zhch158_admin 1 week geleden
bovenliggende
commit
e14a1607c7
1 gewijzigde bestanden met toevoegingen van 44 en 30 verwijderingen
  1. 44 30
      zhch/universal_doc_parser/models/adapters/paddle_layout_detector.py

+ 44 - 30
zhch/universal_doc_parser/models/adapters/paddle_layout_detector.py

@@ -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)