import math from loguru import logger from magic_pdf.config.ocr_content_type import ContentType from magic_pdf.libs.boxbase import (find_bottom_nearest_text_bbox, find_top_nearest_text_bbox) from magic_pdf.libs.commons import join_path TYPE_INLINE_EQUATION = ContentType.InlineEquation TYPE_INTERLINE_EQUATION = ContentType.InterlineEquation UNI_FORMAT_TEXT_TYPE = ['text', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6'] @DeprecationWarning def mk_nlp_markdown_1(para_dict: dict): """对排序后的bboxes拼接内容.""" content_lst = [] for _, page_info in para_dict.items(): para_blocks = page_info.get('para_blocks') if not para_blocks: continue for block in para_blocks: item = block['paras'] for _, p in item.items(): para_text = p['para_text'] is_title = p['is_para_title'] title_level = p['para_title_level'] md_title_prefix = '#' * title_level if is_title: content_lst.append(f'{md_title_prefix} {para_text}') else: content_lst.append(para_text) content_text = '\n\n'.join(content_lst) return content_text # 找到目标字符串在段落中的索引 def __find_index(paragraph, target): index = paragraph.find(target) if index != -1: return index else: return None def __insert_string(paragraph, target, position): new_paragraph = paragraph[:position] + target + paragraph[position:] return new_paragraph def __insert_after(content, image_content, target): """在content中找到target，将image_content插入到target后面.""" index = content.find(target) if index != -1: content = ( content[: index + len(target)] + '\n\n' + image_content + '\n\n' + content[index + len(target) :] ) else: logger.error( f"Can't find the location of image {image_content} in the markdown file, search target is {target}" ) return content def __insert_before(content, image_content, target): """在content中找到target，将image_content插入到target前面.""" index = content.find(target) if index != -1: content = content[:index] + '\n\n' + image_content + '\n\n' + content[index:] else: logger.error( f"Can't find the location of image {image_content} in the markdown file, search target is {target}" ) return content @DeprecationWarning def mk_mm_markdown_1(para_dict: dict): """拼装多模态markdown.""" content_lst = [] for _, page_info in para_dict.items(): page_lst = [] # 一个page内的段落列表 para_blocks = page_info.get('para_blocks') pymu_raw_blocks = page_info.get('preproc_blocks') all_page_images = [] all_page_images.extend(page_info.get('images', [])) all_page_images.extend(page_info.get('image_backup', [])) all_page_images.extend(page_info.get('tables', [])) all_page_images.extend(page_info.get('table_backup', [])) if not para_blocks or not pymu_raw_blocks: # 只有图片的拼接的场景 for img in all_page_images: page_lst.append(f"![]({img['image_path']})") # TODO 图片顺序 page_md = '\n\n'.join(page_lst) else: for block in para_blocks: item = block['paras'] for _, p in item.items(): para_text = p['para_text'] is_title = p['is_para_title'] title_level = p['para_title_level'] md_title_prefix = '#' * title_level if is_title: page_lst.append(f'{md_title_prefix} {para_text}') else: page_lst.append(para_text) """拼装成一个页面的文本""" page_md = '\n\n'.join(page_lst) """插入图片""" for img in all_page_images: imgbox = img['bbox'] img_content = f"![]({img['image_path']})" # 先看在哪个block内 for block in pymu_raw_blocks: bbox = block['bbox'] if ( bbox[0] - 1 <= imgbox[0] < bbox[2] + 1 and bbox[1] - 1 <= imgbox[1] < bbox[3] + 1 ): # 确定在block内 for l in block['lines']: # noqa: E741 line_box = l['bbox'] if ( line_box[0] - 1 <= imgbox[0] < line_box[2] + 1 and line_box[1] - 1 <= imgbox[1] < line_box[3] + 1 ): # 在line内的，插入line前面 line_txt = ''.join([s['text'] for s in l['spans']]) page_md = __insert_before( page_md, img_content, line_txt ) break break else: # 在行与行之间 # 找到图片x0,y0与line的x0,y0最近的line min_distance = 100000 min_line = None for l in block['lines']: # noqa: E741 line_box = l['bbox'] distance = math.sqrt( (line_box[0] - imgbox[0]) ** 2 + (line_box[1] - imgbox[1]) ** 2 ) if distance < min_distance: min_distance = distance min_line = l if min_line: line_txt = ''.join( [s['text'] for s in min_line['spans']] ) img_h = imgbox[3] - imgbox[1] if min_distance < img_h: # 文字在图片前面 page_md = __insert_after( page_md, img_content, line_txt ) else: page_md = __insert_before( page_md, img_content, line_txt ) else: logger.error( f"Can't find the location of image {img['image_path']} in the markdown file #1" ) else: # 应当在两个block之间 # 找到上方最近的block，如果上方没有就找大下方最近的block top_txt_block = find_top_nearest_text_bbox(pymu_raw_blocks, imgbox) if top_txt_block: line_txt = ''.join( [s['text'] for s in top_txt_block['lines'][-1]['spans']] ) page_md = __insert_after(page_md, img_content, line_txt) else: bottom_txt_block = find_bottom_nearest_text_bbox( pymu_raw_blocks, imgbox ) if bottom_txt_block: line_txt = ''.join( [ s['text'] for s in bottom_txt_block['lines'][0]['spans'] ] ) page_md = __insert_before(page_md, img_content, line_txt) else: logger.error( f"Can't find the location of image {img['image_path']} in the markdown file #2" ) content_lst.append(page_md) """拼装成全部页面的文本""" content_text = '\n\n'.join(content_lst) return content_text def __insert_after_para(text, type, element, content_list): """在content_list中找到text，将image_path作为一个新的node插入到text后面.""" for i, c in enumerate(content_list): content_type = c.get('type') if content_type in UNI_FORMAT_TEXT_TYPE and text in c.get('text', ''): if type == 'image': content_node = { 'type': 'image', 'img_path': element.get('image_path'), 'img_alt': '', 'img_title': '', 'img_caption': '', } elif type == 'table': content_node = { 'type': 'table', 'img_path': element.get('image_path'), 'table_latex': element.get('text'), 'table_title': '', 'table_caption': '', 'table_quality': element.get('quality'), } content_list.insert(i + 1, content_node) break else: logger.error( f"Can't find the location of image {element.get('image_path')} in the markdown file, search target is {text}" ) def __insert_before_para(text, type, element, content_list): """在content_list中找到text，将image_path作为一个新的node插入到text前面.""" for i, c in enumerate(content_list): content_type = c.get('type') if content_type in UNI_FORMAT_TEXT_TYPE and text in c.get('text', ''): if type == 'image': content_node = { 'type': 'image', 'img_path': element.get('image_path'), 'img_alt': '', 'img_title': '', 'img_caption': '', } elif type == 'table': content_node = { 'type': 'table', 'img_path': element.get('image_path'), 'table_latex': element.get('text'), 'table_title': '', 'table_caption': '', 'table_quality': element.get('quality'), } content_list.insert(i, content_node) break else: logger.error( f"Can't find the location of image {element.get('image_path')} in the markdown file, search target is {text}" ) def mk_universal_format(pdf_info_list: list, img_buket_path): """构造统一格式 https://aicarrier.feishu.cn/wiki/FqmMwcH69iIdCWkkyjvcDwNUnTY.""" content_lst = [] for page_info in pdf_info_list: page_lst = [] # 一个page内的段落列表 para_blocks = page_info.get('para_blocks') pymu_raw_blocks = page_info.get('preproc_blocks') all_page_images = [] all_page_images.extend(page_info.get('images', [])) all_page_images.extend(page_info.get('image_backup', [])) # all_page_images.extend(page_info.get("tables",[])) # all_page_images.extend(page_info.get("table_backup",[]) ) all_page_tables = [] all_page_tables.extend(page_info.get('tables', [])) if not para_blocks or not pymu_raw_blocks: # 只有图片的拼接的场景 for img in all_page_images: content_node = { 'type': 'image', 'img_path': join_path(img_buket_path, img['image_path']), 'img_alt': '', 'img_title': '', 'img_caption': '', } page_lst.append(content_node) # TODO 图片顺序 for table in all_page_tables: content_node = { 'type': 'table', 'img_path': join_path(img_buket_path, table['image_path']), 'table_latex': table.get('text'), 'table_title': '', 'table_caption': '', 'table_quality': table.get('quality'), } page_lst.append(content_node) # TODO 图片顺序 else: for block in para_blocks: item = block['paras'] for _, p in item.items(): font_type = p[ 'para_font_type' ] # 对于文本来说，要么是普通文本，要么是个行间公式 if font_type == TYPE_INTERLINE_EQUATION: content_node = {'type': 'equation', 'latex': p['para_text']} page_lst.append(content_node) else: para_text = p['para_text'] is_title = p['is_para_title'] title_level = p['para_title_level'] if is_title: content_node = { 'type': f'h{title_level}', 'text': para_text, } page_lst.append(content_node) else: content_node = {'type': 'text', 'text': para_text} page_lst.append(content_node) content_lst.extend(page_lst) """插入图片""" for img in all_page_images: insert_img_or_table('image', img, pymu_raw_blocks, content_lst) """插入表格""" for table in all_page_tables: insert_img_or_table('table', table, pymu_raw_blocks, content_lst) # end for return content_lst def insert_img_or_table(type, element, pymu_raw_blocks, content_lst): element_bbox = element['bbox'] # 先看在哪个block内 for block in pymu_raw_blocks: bbox = block['bbox'] if ( bbox[0] - 1 <= element_bbox[0] < bbox[2] + 1 and bbox[1] - 1 <= element_bbox[1] < bbox[3] + 1 ): # 确定在这个大的block内，然后进入逐行比较距离 for l in block['lines']: # noqa: E741 line_box = l['bbox'] if ( line_box[0] - 1 <= element_bbox[0] < line_box[2] + 1 and line_box[1] - 1 <= element_bbox[1] < line_box[3] + 1 ): # 在line内的，插入line前面 line_txt = ''.join([s['text'] for s in l['spans']]) __insert_before_para(line_txt, type, element, content_lst) break break else: # 在行与行之间 # 找到图片x0,y0与line的x0,y0最近的line min_distance = 100000 min_line = None for l in block['lines']: # noqa: E741 line_box = l['bbox'] distance = math.sqrt( (line_box[0] - element_bbox[0]) ** 2 + (line_box[1] - element_bbox[1]) ** 2 ) if distance < min_distance: min_distance = distance min_line = l if min_line: line_txt = ''.join([s['text'] for s in min_line['spans']]) img_h = element_bbox[3] - element_bbox[1] if min_distance < img_h: # 文字在图片前面 __insert_after_para(line_txt, type, element, content_lst) else: __insert_before_para(line_txt, type, element, content_lst) break else: logger.error( f"Can't find the location of image {element.get('image_path')} in the markdown file #1" ) else: # 应当在两个block之间 # 找到上方最近的block，如果上方没有就找大下方最近的block top_txt_block = find_top_nearest_text_bbox(pymu_raw_blocks, element_bbox) if top_txt_block: line_txt = ''.join([s['text'] for s in top_txt_block['lines'][-1]['spans']]) __insert_after_para(line_txt, type, element, content_lst) else: bottom_txt_block = find_bottom_nearest_text_bbox( pymu_raw_blocks, element_bbox ) if bottom_txt_block: line_txt = ''.join( [s['text'] for s in bottom_txt_block['lines'][0]['spans']] ) __insert_before_para(line_txt, type, element, content_lst) else: # TODO ，图片可能独占一列，这种情况上下是没有图片的 logger.error( f"Can't find the location of image {element.get('image_path')} in the markdown file #2" ) def mk_mm_markdown(content_list): """基于同一格式的内容列表，构造markdown，含图片.""" content_md = [] for c in content_list: content_type = c.get('type') if content_type == 'text': content_md.append(c.get('text')) elif content_type == 'equation': content = c.get('latex') if content.startswith('$$') and content.endswith('$$'): content_md.append(content) else: content_md.append(f"\n$$\n{c.get('latex')}\n$$\n") elif content_type in UNI_FORMAT_TEXT_TYPE: content_md.append(f"{'#'*int(content_type[1])} {c.get('text')}") elif content_type == 'image': content_md.append(f"![]({c.get('img_path')})") return '\n\n'.join(content_md) def mk_nlp_markdown(content_list): """基于同一格式的内容列表，构造markdown，不含图片.""" content_md = [] for c in content_list: content_type = c.get('type') if content_type == 'text': content_md.append(c.get('text')) elif content_type == 'equation': content_md.append(f"$$\n{c.get('latex')}\n$$") elif content_type == 'table': content_md.append(f"$$$\n{c.get('table_latex')}\n$$$") elif content_type in UNI_FORMAT_TEXT_TYPE: content_md.append(f"{'#'*int(content_type[1])} {c.get('text')}") return '\n\n'.join(content_md)