feat: enhance bounding box processing by adding support for CHART and REF_TEXT block types

mineru/backend/pipeline/model_json_to_middle_json.py

@@ -300,7 +300,7 @@ def _apply_post_ocr(pdf_info_list, lang=None):
             span['score'] = 0.0
 
 
-def _apply_default_title_levels(pdf_info_list):
+def _post_block_process(pdf_info_list):
     for page_info in pdf_info_list:
         for block_key in ["preproc_blocks", "para_blocks"]:
             for block in page_info.get(block_key, []):
@@ -311,6 +311,8 @@ def _apply_default_title_levels(pdf_info_list):
                 elif block_type == BlockType.PARAGRAPH_TITLE:
                     block["type"] = BlockType.TITLE
                     block["level"] = 2
+                elif block_type == BlockType.VERTICAL_TEXT:
+                    block["type"] = BlockType.TEXT
 
 
 def finalize_middle_json(pdf_info_list, lang=None, ocr_enable=False):
@@ -328,7 +330,7 @@ def finalize_middle_json(pdf_info_list, lang=None, ocr_enable=False):
             llm_aided_title(pdf_info_list, title_aided_config)
             logger.info(f'llm aided title time: {round(time.time() - llm_aided_title_start_time, 2)}')
 
-    _apply_default_title_levels(pdf_info_list)
+    _post_block_process(pdf_info_list)
 
     if os.getenv('MINERU_DONOT_CLEAN_MEM') is None and len(pdf_info_list) >= 10:
         clean_memory(get_device())

mineru/utils/draw_bbox.py

@@ -314,7 +314,7 @@ def draw_span_bbox(pdf_info, pdf_bytes, out_path, filename):
             page_inline_equation_list.append(span['bbox'])
         elif span['type'] == ContentType.INTERLINE_EQUATION:
             page_interline_equation_list.append(span['bbox'])
-        elif span['type'] == ContentType.IMAGE:
+        elif span['type'] in [ContentType.IMAGE, ContentType.CHART, ContentType.SEAL]:
             page_image_list.append(span['bbox'])
         elif span['type'] == ContentType.TABLE:
             page_table_list.append(span['bbox'])
@@ -330,10 +330,9 @@ def draw_span_bbox(pdf_info, pdf_bytes, out_path, filename):
 
         # 构造dropped_list
         for block in page['discarded_blocks']:
-            if block['type'] == BlockType.DISCARDED:
-                for line in block['lines']:
-                    for span in line['spans']:
-                        page_dropped_list.append(span['bbox'])
+            for line in block['lines']:
+                for span in line['spans']:
+                    page_dropped_list.append(span['bbox'])
         dropped_list.append(page_dropped_list)
         # 构造其余useful_list
         # for block in page['para_blocks']:  # span直接用分段合并前的结果就可以
@@ -344,11 +343,14 @@ def draw_span_bbox(pdf_info, pdf_bytes, out_path, filename):
                 BlockType.INTERLINE_EQUATION,
                 BlockType.LIST,
                 BlockType.INDEX,
+                BlockType.REF_TEXT,
+                BlockType.ABSTRACT,
+                BlockType.SEAL,
             ]:
                 for line in block['lines']:
                     for span in line['spans']:
                         get_span_info(span)
-            elif block['type'] in [BlockType.IMAGE, BlockType.TABLE]:
+            elif block['type'] in [BlockType.IMAGE, BlockType.TABLE, BlockType.CHART, BlockType.CODE]:
                 for sub_block in block['blocks']:
                     for line in sub_block['lines']:
                         for span in line['spans']:

mineru/utils/enum_class.py

@@ -56,7 +56,6 @@ class ContentType:
     INTERLINE_EQUATION = 'interline_equation'
     INLINE_EQUATION = 'inline_equation'
     EQUATION = 'equation'
-    CODE = 'code'
     HYPERLINK = 'hyperlink'
     SEAL = 'seal'