diff --git a/mineru/backend/pipeline/model_json_to_middle_json.py b/mineru/backend/pipeline/model_json_to_middle_json.py index 53f8136a..5f66e6f5 100644 --- a/mineru/backend/pipeline/model_json_to_middle_json.py +++ b/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()) diff --git a/mineru/utils/draw_bbox.py b/mineru/utils/draw_bbox.py index 77ccf3bd..6fa56848 100644 --- a/mineru/utils/draw_bbox.py +++ b/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']: diff --git a/mineru/utils/enum_class.py b/mineru/utils/enum_class.py index a45aee39..2c84ed00 100644 --- a/mineru/utils/enum_class.py +++ b/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'