263 lines
11 KiB
Python
263 lines
11 KiB
Python
|
# Copyright (c) Opendatalab. All rights reserved.
|
|||
|
|
import os
|
|||
|
|
import time
|
|||
|
|
|
|||
|
|
from loguru import logger
|
|||
|
|
from tqdm import tqdm
|
|||
|
|
|
|||
|
|
from mineru.backend.utils import cross_page_table_merge
|
|||
|
|
from mineru.utils.config_reader import get_device, get_llm_aided_config, get_formula_enable
|
|||
|
|
from mineru.backend.pipeline.model_init import AtomModelSingleton
|
|||
|
|
from mineru.backend.pipeline.para_split import para_split
|
|||
|
|
from mineru.utils.block_pre_proc import prepare_block_bboxes, process_groups
|
|||
|
|
from mineru.utils.block_sort import sort_blocks_by_bbox
|
|||
|
|
from mineru.utils.boxbase import calculate_overlap_area_in_bbox1_area_ratio
|
|||
|
|
from mineru.utils.cut_image import cut_image_and_table
|
|||
|
|
from mineru.utils.enum_class import ContentType
|
|||
|
|
from mineru.utils.llm_aided import llm_aided_title
|
|||
|
|
from mineru.utils.model_utils import clean_memory
|
|||
|
|
from mineru.backend.pipeline.pipeline_magic_model import MagicModel
|
|||
|
|
from mineru.utils.ocr_utils import OcrConfidence
|
|||
|
|
from mineru.utils.span_block_fix import fill_spans_in_blocks, fix_discarded_block, fix_block_spans
|
|||
|
|
from mineru.utils.span_pre_proc import remove_outside_spans, remove_overlaps_low_confidence_spans, \
|
|||
|
|
remove_overlaps_min_spans, txt_spans_extract
|
|||
|
|
from mineru.version import __version__
|
|||
|
|
from mineru.utils.hash_utils import bytes_md5
|
|||
|
|
|
|||
|
|
|
|||
|
|
def page_model_info_to_page_info(page_model_info, image_dict, page, image_writer, page_index, ocr_enable=False, formula_enabled=True):
|
|||
|
|
scale = image_dict["scale"]
|
|||
|
|
page_pil_img = image_dict["img_pil"]
|
|||
|
|
# page_img_md5 = str_md5(image_dict["img_base64"])
|
|||
|
|
page_img_md5 = bytes_md5(page_pil_img.tobytes())
|
|||
|
|
page_w, page_h = map(int, page.get_size())
|
|||
|
|
magic_model = MagicModel(page_model_info, scale)
|
|||
|
|
|
|||
|
|
"""从magic_model对象中获取后面会用到的区块信息"""
|
|||
|
|
discarded_blocks = magic_model.get_discarded()
|
|||
|
|
text_blocks = magic_model.get_text_blocks()
|
|||
|
|
title_blocks = magic_model.get_title_blocks()
|
|||
|
|
inline_equations, interline_equations, interline_equation_blocks = magic_model.get_equations()
|
|||
|
|
|
|||
|
|
img_groups = magic_model.get_imgs()
|
|||
|
|
table_groups = magic_model.get_tables()
|
|||
|
|
|
|||
|
|
"""对image和table的区块分组"""
|
|||
|
|
img_body_blocks, img_caption_blocks, img_footnote_blocks, maybe_text_image_blocks = process_groups(
|
|||
|
|
img_groups, 'image_body', 'image_caption_list', 'image_footnote_list'
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
table_body_blocks, table_caption_blocks, table_footnote_blocks, _ = process_groups(
|
|||
|
|
table_groups, 'table_body', 'table_caption_list', 'table_footnote_list'
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
"""获取所有的spans信息"""
|
|||
|
|
spans = magic_model.get_all_spans()
|
|||
|
|
|
|||
|
|
"""某些图可能是文本块，通过简单的规则判断一下"""
|
|||
|
|
if len(maybe_text_image_blocks) > 0:
|
|||
|
|
for block in maybe_text_image_blocks:
|
|||
|
|
should_add_to_text_blocks = False
|
|||
|
|
|
|||
|
|
if ocr_enable:
|
|||
|
|
# 找到与当前block重叠的text spans
|
|||
|
|
span_in_block_list = [
|
|||
|
|
span for span in spans
|
|||
|
|
if span['type'] == 'text' and
|
|||
|
|
calculate_overlap_area_in_bbox1_area_ratio(span['bbox'], block['bbox']) > 0.7
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
if len(span_in_block_list) > 0:
|
|||
|
|
# 计算spans总面积
|
|||
|
|
spans_area = sum(
|
|||
|
|
(span['bbox'][2] - span['bbox'][0]) * (span['bbox'][3] - span['bbox'][1])
|
|||
|
|
for span in span_in_block_list
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 计算block面积
|
|||
|
|
block_area = (block['bbox'][2] - block['bbox'][0]) * (block['bbox'][3] - block['bbox'][1])
|
|||
|
|
|
|||
|
|
# 判断是否符合文本图条件
|
|||
|
|
if block_area > 0 and spans_area / block_area > 0.25:
|
|||
|
|
should_add_to_text_blocks = True
|
|||
|
|
|
|||
|
|
# 根据条件决定添加到哪个列表
|
|||
|
|
if should_add_to_text_blocks:
|
|||
|
|
block.pop('group_id', None) # 移除group_id
|
|||
|
|
text_blocks.append(block)
|
|||
|
|
else:
|
|||
|
|
img_body_blocks.append(block)
|
|||
|
|
|
|||
|
|
|
|||
|
|
"""将所有区块的bbox整理到一起"""
|
|||
|
|
if formula_enabled:
|
|||
|
|
interline_equation_blocks = []
|
|||
|
|
|
|||
|
|
if len(interline_equation_blocks) > 0:
|
|||
|
|
|
|||
|
|
for block in interline_equation_blocks:
|
|||
|
|
spans.append({
|
|||
|
|
"type": ContentType.INTERLINE_EQUATION,
|
|||
|
|
'score': block['score'],
|
|||
|
|
"bbox": block['bbox'],
|
|||
|
|
"content": "",
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
all_bboxes, all_discarded_blocks, footnote_blocks = prepare_block_bboxes(
|
|||
|
|
img_body_blocks, img_caption_blocks, img_footnote_blocks,
|
|||
|
|
table_body_blocks, table_caption_blocks, table_footnote_blocks,
|
|||
|
|
discarded_blocks,
|
|||
|
|
text_blocks,
|
|||
|
|
title_blocks,
|
|||
|
|
interline_equation_blocks,
|
|||
|
|
page_w,
|
|||
|
|
page_h,
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
all_bboxes, all_discarded_blocks, footnote_blocks = prepare_block_bboxes(
|
|||
|
|
img_body_blocks, img_caption_blocks, img_footnote_blocks,
|
|||
|
|
table_body_blocks, table_caption_blocks, table_footnote_blocks,
|
|||
|
|
discarded_blocks,
|
|||
|
|
text_blocks,
|
|||
|
|
title_blocks,
|
|||
|
|
interline_equations,
|
|||
|
|
page_w,
|
|||
|
|
page_h,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
"""在删除重复span之前，应该通过image_body和table_body的block过滤一下image和table的span"""
|
|||
|
|
"""顺便删除大水印并保留abandon的span"""
|
|||
|
|
spans = remove_outside_spans(spans, all_bboxes, all_discarded_blocks)
|
|||
|
|
|
|||
|
|
"""删除重叠spans中置信度较低的那些"""
|
|||
|
|
spans, dropped_spans_by_confidence = remove_overlaps_low_confidence_spans(spans)
|
|||
|
|
"""删除重叠spans中较小的那些"""
|
|||
|
|
spans, dropped_spans_by_span_overlap = remove_overlaps_min_spans(spans)
|
|||
|
|
|
|||
|
|
"""根据parse_mode，构造spans，主要是文本类的字符填充"""
|
|||
|
|
if ocr_enable:
|
|||
|
|
pass
|
|||
|
|
else:
|
|||
|
|
"""使用新版本的混合ocr方案."""
|
|||
|
|
spans = txt_spans_extract(page, spans, page_pil_img, scale, all_bboxes, all_discarded_blocks)
|
|||
|
|
|
|||
|
|
"""先处理不需要排版的discarded_blocks"""
|
|||
|
|
discarded_block_with_spans, spans = fill_spans_in_blocks(
|
|||
|
|
all_discarded_blocks, spans, 0.4
|
|||
|
|
)
|
|||
|
|
fix_discarded_blocks = fix_discarded_block(discarded_block_with_spans)
|
|||
|
|
|
|||
|
|
"""如果当前页面没有有效的bbox则跳过"""
|
|||
|
|
if len(all_bboxes) == 0 and len(fix_discarded_blocks) == 0:
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
"""对image/table/interline_equation截图"""
|
|||
|
|
for span in spans:
|
|||
|
|
if span['type'] in [ContentType.IMAGE, ContentType.TABLE, ContentType.INTERLINE_EQUATION]:
|
|||
|
|
span = cut_image_and_table(
|
|||
|
|
span, page_pil_img, page_img_md5, page_index, image_writer, scale=scale
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
"""span填充进block"""
|
|||
|
|
block_with_spans, spans = fill_spans_in_blocks(all_bboxes, spans, 0.5)
|
|||
|
|
|
|||
|
|
"""对block进行fix操作"""
|
|||
|
|
fix_blocks = fix_block_spans(block_with_spans)
|
|||
|
|
|
|||
|
|
"""对block进行排序"""
|
|||
|
|
sorted_blocks = sort_blocks_by_bbox(fix_blocks, page_w, page_h, footnote_blocks)
|
|||
|
|
|
|||
|
|
"""构造page_info"""
|
|||
|
|
page_info = make_page_info_dict(sorted_blocks, page_index, page_w, page_h, fix_discarded_blocks)
|
|||
|
|
|
|||
|
|
return page_info
|
|||
|
|
|
|||
|
|
|
|||
|
|
def result_to_middle_json(model_list, images_list, pdf_doc, image_writer, lang=None, ocr_enable=False, formula_enabled=True):
|
|||
|
|
middle_json = {"pdf_info": [], "_backend":"pipeline", "_version_name": __version__}
|
|||
|
|
formula_enabled = get_formula_enable(formula_enabled)
|
|||
|
|
for page_index, page_model_info in tqdm(enumerate(model_list), total=len(model_list), desc="Processing pages"):
|
|||
|
|
page = pdf_doc[page_index]
|
|||
|
|
image_dict = images_list[page_index]
|
|||
|
|
page_info = page_model_info_to_page_info(
|
|||
|
|
page_model_info, image_dict, page, image_writer, page_index, ocr_enable=ocr_enable, formula_enabled=formula_enabled
|
|||
|
|
)
|
|||
|
|
if page_info is None:
|
|||
|
|
page_w, page_h = map(int, page.get_size())
|
|||
|
|
page_info = make_page_info_dict([], page_index, page_w, page_h, [])
|
|||
|
|
middle_json["pdf_info"].append(page_info)
|
|||
|
|
|
|||
|
|
"""后置ocr处理"""
|
|||
|
|
need_ocr_list = []
|
|||
|
|
img_crop_list = []
|
|||
|
|
text_block_list = []
|
|||
|
|
for page_info in middle_json["pdf_info"]:
|
|||
|
|
for block in page_info['preproc_blocks']:
|
|||
|
|
if block['type'] in ['table', 'image']:
|
|||
|
|
for sub_block in block['blocks']:
|
|||
|
|
if sub_block['type'] in ['image_caption', 'image_footnote', 'table_caption', 'table_footnote']:
|
|||
|
|
text_block_list.append(sub_block)
|
|||
|
|
elif block['type'] in ['text', 'title']:
|
|||
|
|
text_block_list.append(block)
|
|||
|
|
for block in page_info['discarded_blocks']:
|
|||
|
|
text_block_list.append(block)
|
|||
|
|
for block in text_block_list:
|
|||
|
|
for line in block['lines']:
|
|||
|
|
for span in line['spans']:
|
|||
|
|
if 'np_img' in span:
|
|||
|
|
need_ocr_list.append(span)
|
|||
|
|
img_crop_list.append(span['np_img'])
|
|||
|
|
span.pop('np_img')
|
|||
|
|
if len(img_crop_list) > 0:
|
|||
|
|
atom_model_manager = AtomModelSingleton()
|
|||
|
|
ocr_model = atom_model_manager.get_atom_model(
|
|||
|
|
atom_model_name='ocr',
|
|||
|
|
det_db_box_thresh=0.3,
|
|||
|
|
lang=lang
|
|||
|
|
)
|
|||
|
|
ocr_res_list = ocr_model.ocr(img_crop_list, det=False, tqdm_enable=True)[0]
|
|||
|
|
assert len(ocr_res_list) == len(
|
|||
|
|
need_ocr_list), f'ocr_res_list: {len(ocr_res_list)}, need_ocr_list: {len(need_ocr_list)}'
|
|||
|
|
for index, span in enumerate(need_ocr_list):
|
|||
|
|
ocr_text, ocr_score = ocr_res_list[index]
|
|||
|
|
if ocr_score > OcrConfidence.min_confidence:
|
|||
|
|
span['content'] = ocr_text
|
|||
|
|
span['score'] = float(f"{ocr_score:.3f}")
|
|||
|
|
else:
|
|||
|
|
span['content'] = ''
|
|||
|
|
span['score'] = 0.0
|
|||
|
|
|
|||
|
|
"""分段"""
|
|||
|
|
para_split(middle_json["pdf_info"])
|
|||
|
|
|
|||
|
|
"""表格跨页合并"""
|
|||
|
|
cross_page_table_merge(middle_json["pdf_info"])
|
|||
|
|
|
|||
|
|
"""llm优化"""
|
|||
|
|
llm_aided_config = get_llm_aided_config()
|
|||
|
|
|
|||
|
|
if llm_aided_config is not None:
|
|||
|
|
"""标题优化"""
|
|||
|
|
title_aided_config = llm_aided_config.get('title_aided', None)
|
|||
|
|
if title_aided_config is not None:
|
|||
|
|
if title_aided_config.get('enable', False):
|
|||
|
|
llm_aided_title_start_time = time.time()
|
|||
|
|
llm_aided_title(middle_json["pdf_info"], title_aided_config)
|
|||
|
|
logger.info(f'llm aided title time: {round(time.time() - llm_aided_title_start_time, 2)}')
|
|||
|
|
|
|||
|
|
"""清理内存"""
|
|||
|
|
pdf_doc.close()
|
|||
|
|
if os.getenv('MINERU_DONOT_CLEAN_MEM') is None and len(model_list) >= 10:
|
|||
|
|
clean_memory(get_device())
|
|||
|
|
|
|||
|
|
return middle_json
|
|||
|
|
|
|||
|
|
|
|||
|
|
def make_page_info_dict(blocks, page_id, page_w, page_h, discarded_blocks):
|
|||
|
|
return_dict = {
|
|||
|
|
'preproc_blocks': blocks,
|
|||
|
|
'page_idx': page_id,
|
|||
|
|
'page_size': [page_w, page_h],
|
|||
|
|
'discarded_blocks': discarded_blocks,
|
|||
|
|
}
|
|||
|
|
return return_dict
|