mkcontent.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250
  1. import re
  2. import math
  3. from loguru import logger
  4. from libs.boxbase import find_bottom_nearest_text_bbox, find_top_nearest_text_bbox
  5. def mk_nlp_markdown(para_dict: dict):
  6. """
  7. 对排序后的bboxes拼接内容
  8. """
  9. content_lst = []
  10. for _, page_info in para_dict.items():
  11. para_blocks = page_info.get("para_blocks")
  12. if not para_blocks:
  13. continue
  14. for block in para_blocks:
  15. item = block["paras"]
  16. for _, p in item.items():
  17. para_text = p["para_text"]
  18. is_title = p["is_para_title"]
  19. title_level = p['para_title_level']
  20. md_title_prefix = "#"*title_level
  21. if is_title:
  22. content_lst.append(f"{md_title_prefix} {para_text}")
  23. else:
  24. content_lst.append(para_text)
  25. content_text = "\n\n".join(content_lst)
  26. return content_text
  27. # 找到目标字符串在段落中的索引
  28. def __find_index(paragraph, target):
  29. index = paragraph.find(target)
  30. if index != -1:
  31. return index
  32. else:
  33. return None
  34. def __insert_string(paragraph, target, postion):
  35. new_paragraph = paragraph[:postion] + target + paragraph[postion:]
  36. return new_paragraph
  37. def __insert_after(content, image_content, target):
  38. """
  39. 在content中找到target,将image_content插入到target后面
  40. """
  41. index = content.find(target)
  42. if index != -1:
  43. content = content[:index+len(target)] + "\n\n" + image_content + "\n\n" + content[index+len(target):]
  44. else:
  45. logger.error(f"Can't find the location of image {image_content} in the markdown file, search target is {target}")
  46. return content
  47. def __insert_before(content, image_content, target):
  48. """
  49. 在content中找到target,将image_content插入到target前面
  50. """
  51. index = content.find(target)
  52. if index != -1:
  53. content = content[:index] + "\n\n" + image_content + "\n\n" + content[index:]
  54. else:
  55. logger.error(f"Can't find the location of image {image_content} in the markdown file, search target is {target}")
  56. return content
  57. def mk_mm_markdown(para_dict: dict):
  58. """拼装多模态markdown"""
  59. content_lst = []
  60. for _, page_info in para_dict.items():
  61. page_lst = [] # 一个page内的段落列表
  62. para_blocks = page_info.get("para_blocks")
  63. pymu_raw_blocks = page_info.get("preproc_blocks")
  64. all_page_images = []
  65. all_page_images.extend(page_info.get("images",[]))
  66. all_page_images.extend(page_info.get("image_backup", []) )
  67. all_page_images.extend(page_info.get("tables",[]))
  68. all_page_images.extend(page_info.get("table_backup",[]) )
  69. if not para_blocks or not pymu_raw_blocks: # 只有图片的拼接的场景
  70. for img in all_page_images:
  71. page_lst.append(f"![]({img['image_path']})") # TODO 图片顺序
  72. page_md = "\n\n".join(page_lst)
  73. else:
  74. for block in para_blocks:
  75. item = block["paras"]
  76. for _, p in item.items():
  77. para_text = p["para_text"]
  78. is_title = p["is_para_title"]
  79. title_level = p['para_title_level']
  80. md_title_prefix = "#"*title_level
  81. if is_title:
  82. page_lst.append(f"{md_title_prefix} {para_text}")
  83. else:
  84. page_lst.append(para_text)
  85. """拼装成一个页面的文本"""
  86. page_md = "\n\n".join(page_lst)
  87. """插入图片"""
  88. for img in all_page_images:
  89. imgbox = img['bbox']
  90. img_content = f"![]({img['image_path']})"
  91. # 先看在哪个block内
  92. for block in pymu_raw_blocks:
  93. bbox = block['bbox']
  94. if bbox[0]-1 <= imgbox[0] < bbox[2]+1 and bbox[1]-1 <= imgbox[1] < bbox[3]+1:# 确定在block内
  95. for l in block['lines']:
  96. line_box = l['bbox']
  97. if line_box[0]-1 <= imgbox[0] < line_box[2]+1 and line_box[1]-1 <= imgbox[1] < line_box[3]+1: # 在line内的,插入line前面
  98. line_txt = "".join([s['text'] for s in l['spans']])
  99. page_md = __insert_before(page_md, img_content, line_txt)
  100. break
  101. break
  102. else:# 在行与行之间
  103. # 找到图片x0,y0与line的x0,y0最近的line
  104. min_distance = 100000
  105. min_line = None
  106. for l in block['lines']:
  107. line_box = l['bbox']
  108. distance = math.sqrt((line_box[0] - imgbox[0])**2 + (line_box[1] - imgbox[1])**2)
  109. if distance < min_distance:
  110. min_distance = distance
  111. min_line = l
  112. if min_line:
  113. line_txt = "".join([s['text'] for s in min_line['spans']])
  114. img_h = imgbox[3] - imgbox[1]
  115. if min_distance<img_h: # 文字在图片前面
  116. page_md = __insert_after(page_md, img_content, line_txt)
  117. else:
  118. page_md = __insert_before(page_md, img_content, line_txt)
  119. else:
  120. logger.error(f"Can't find the location of image {img['image_path']} in the markdown file")
  121. else:# 应当在两个block之间
  122. # 找到上方最近的block,如果上方没有就找大下方最近的block
  123. top_txt_block = find_top_nearest_text_bbox(pymu_raw_blocks, imgbox)
  124. if top_txt_block:
  125. line_txt = "".join([s['text'] for s in top_txt_block['lines'][-1]['spans']])
  126. page_md = __insert_after(page_md, img_content, line_txt)
  127. else:
  128. bottom_txt_block = find_bottom_nearest_text_bbox(pymu_raw_blocks, imgbox)
  129. if bottom_txt_block:
  130. line_txt = "".join([s['text'] for s in bottom_txt_block['lines'][0]['spans']])
  131. page_md = __insert_before(page_md, img_content, line_txt)
  132. else:
  133. logger.error(f"Can't find the location of image {img['image_path']} in the markdown file")
  134. content_lst.append(page_md)
  135. """拼装成全部页面的文本"""
  136. content_text = "\n\n".join(content_lst)
  137. return content_text
  138. @DeprecationWarning
  139. def mk_mm_markdown_1(para_dict: dict):
  140. """
  141. 得到images和tables变量
  142. """
  143. image_all_list = []
  144. for _, page_info in para_dict.items():
  145. images = page_info.get("images",[])
  146. tables = page_info.get("tables",[])
  147. image_backup = page_info.get("image_backup", [])
  148. table_backup = page_info.get("table_backup",[])
  149. all_page_images = []
  150. all_page_images.extend(images)
  151. all_page_images.extend(image_backup)
  152. all_page_images.extend(tables)
  153. all_page_images.extend(table_backup)
  154. pymu_raw_blocks = page_info.get("pymu_raw_blocks")
  155. # 提取每个图片所在位置
  156. for image_info in all_page_images:
  157. x0_image, y0_image, x1_image, y1_image = image_info['bbox'][:4]
  158. image_path = image_info['image_path']
  159. # 判断图片处于原始PDF中哪个模块之间
  160. image_internal_dict = {}
  161. image_external_dict = {}
  162. between_dict = {}
  163. for block in pymu_raw_blocks:
  164. x0, y0, x1, y1 = block['bbox'][:4]
  165. # 在某个模块内部
  166. if x0 <= x0_image < x1 and y0 <= y0_image < y1:
  167. image_internal_dict['bbox'] = [x0_image, y0_image, x1_image, y1_image]
  168. image_internal_dict['path'] = image_path
  169. # 确定图片在哪句文本之前
  170. y_pre = 0
  171. for line in block['lines']:
  172. x0, y0, x1, y1 = line['spans'][0]['bbox']
  173. if x0 <= x0_image < x1 and y_pre <= y0_image < y0:
  174. text = line['spans']['text']
  175. image_internal_dict['text'] = text
  176. image_internal_dict['markdown_image'] = f'![image_path]({image_path})'
  177. break
  178. else:
  179. y_pre = y0
  180. # 在某两个模块之间
  181. elif x0 <= x0_image < x1:
  182. distance = math.sqrt((x1_image - x0)**2 + (y1_image - y0)**2)
  183. between_dict[block['number']] = distance
  184. # 找到与定位点距离最小的文本block
  185. if between_dict:
  186. min_key = min(between_dict, key=between_dict.get)
  187. spans_list = []
  188. for span in pymu_raw_blocks[min_key]['lines']:
  189. for text_piece in span['spans']:
  190. # 防止索引定位文本内容过多
  191. if len(spans_list) < 60:
  192. spans_list.append(text_piece['text'])
  193. text1 = ''.join(spans_list)
  194. image_external_dict['bbox'] = [x0_image, y0_image, x1_image, y1_image]
  195. image_external_dict['path'] = image_path
  196. image_external_dict['text'] = text1
  197. image_external_dict['markdown_image'] = f'![image_path]({image_path})'
  198. # 将内部图片或外部图片存入当页所有图片的列表
  199. if len(image_internal_dict) != 0:
  200. image_all_list.append(image_internal_dict)
  201. elif len(image_external_dict) != 0:
  202. image_all_list.append(image_external_dict)
  203. else:
  204. logger.error(f"Can't find the location of image {image_path} in the markdown file")
  205. content_text = mk_nlp_markdown(para_dict)
  206. for image_info_extract in image_all_list:
  207. loc = __find_index(content_text, image_info_extract['text'])
  208. if loc is not None:
  209. content_text = __insert_string(content_text, image_info_extract['markdown_image'], loc)
  210. else:
  211. logger.error(f"Can't find the location of image {image_info_extract['path']} in the markdown file")
  212. return content_text