mkcontent.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353
  1. import math
  2. from loguru import logger
  3. from magic_pdf.libs.boxbase import find_bottom_nearest_text_bbox, find_top_nearest_text_bbox
  4. from magic_pdf.libs.ocr_content_type import ContentType
  5. TYPE_INLINE_EQUATION = ContentType.InlineEquation
  6. TYPE_INTERLINE_EQUATION = ContentType.InterlineEquation
  7. UNI_FORMAT_TEXT_TYPE = ['text', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6']
  8. @DeprecationWarning
  9. def mk_nlp_markdown_1(para_dict: dict):
  10. """
  11. 对排序后的bboxes拼接内容
  12. """
  13. content_lst = []
  14. for _, page_info in para_dict.items():
  15. para_blocks = page_info.get("para_blocks")
  16. if not para_blocks:
  17. continue
  18. for block in para_blocks:
  19. item = block["paras"]
  20. for _, p in item.items():
  21. para_text = p["para_text"]
  22. is_title = p["is_para_title"]
  23. title_level = p['para_title_level']
  24. md_title_prefix = "#"*title_level
  25. if is_title:
  26. content_lst.append(f"{md_title_prefix} {para_text}")
  27. else:
  28. content_lst.append(para_text)
  29. content_text = "\n\n".join(content_lst)
  30. return content_text
  31. # 找到目标字符串在段落中的索引
  32. def __find_index(paragraph, target):
  33. index = paragraph.find(target)
  34. if index != -1:
  35. return index
  36. else:
  37. return None
  38. def __insert_string(paragraph, target, postion):
  39. new_paragraph = paragraph[:postion] + target + paragraph[postion:]
  40. return new_paragraph
  41. def __insert_after(content, image_content, target):
  42. """
  43. 在content中找到target,将image_content插入到target后面
  44. """
  45. index = content.find(target)
  46. if index != -1:
  47. content = content[:index+len(target)] + "\n\n" + image_content + "\n\n" + content[index+len(target):]
  48. else:
  49. logger.error(f"Can't find the location of image {image_content} in the markdown file, search target is {target}")
  50. return content
  51. def __insert_before(content, image_content, target):
  52. """
  53. 在content中找到target,将image_content插入到target前面
  54. """
  55. index = content.find(target)
  56. if index != -1:
  57. content = content[:index] + "\n\n" + image_content + "\n\n" + content[index:]
  58. else:
  59. logger.error(f"Can't find the location of image {image_content} in the markdown file, search target is {target}")
  60. return content
  61. @DeprecationWarning
  62. def mk_mm_markdown_1(para_dict: dict):
  63. """拼装多模态markdown"""
  64. content_lst = []
  65. for _, page_info in para_dict.items():
  66. page_lst = [] # 一个page内的段落列表
  67. para_blocks = page_info.get("para_blocks")
  68. pymu_raw_blocks = page_info.get("preproc_blocks")
  69. all_page_images = []
  70. all_page_images.extend(page_info.get("images",[]))
  71. all_page_images.extend(page_info.get("image_backup", []) )
  72. all_page_images.extend(page_info.get("tables",[]))
  73. all_page_images.extend(page_info.get("table_backup",[]) )
  74. if not para_blocks or not pymu_raw_blocks: # 只有图片的拼接的场景
  75. for img in all_page_images:
  76. page_lst.append(f"![]({img['image_path']})") # TODO 图片顺序
  77. page_md = "\n\n".join(page_lst)
  78. else:
  79. for block in para_blocks:
  80. item = block["paras"]
  81. for _, p in item.items():
  82. para_text = p["para_text"]
  83. is_title = p["is_para_title"]
  84. title_level = p['para_title_level']
  85. md_title_prefix = "#"*title_level
  86. if is_title:
  87. page_lst.append(f"{md_title_prefix} {para_text}")
  88. else:
  89. page_lst.append(para_text)
  90. """拼装成一个页面的文本"""
  91. page_md = "\n\n".join(page_lst)
  92. """插入图片"""
  93. for img in all_page_images:
  94. imgbox = img['bbox']
  95. img_content = f"![]({img['image_path']})"
  96. # 先看在哪个block内
  97. for block in pymu_raw_blocks:
  98. bbox = block['bbox']
  99. if bbox[0]-1 <= imgbox[0] < bbox[2]+1 and bbox[1]-1 <= imgbox[1] < bbox[3]+1:# 确定在block内
  100. for l in block['lines']:
  101. line_box = l['bbox']
  102. if line_box[0]-1 <= imgbox[0] < line_box[2]+1 and line_box[1]-1 <= imgbox[1] < line_box[3]+1: # 在line内的,插入line前面
  103. line_txt = "".join([s['text'] for s in l['spans']])
  104. page_md = __insert_before(page_md, img_content, line_txt)
  105. break
  106. break
  107. else:# 在行与行之间
  108. # 找到图片x0,y0与line的x0,y0最近的line
  109. min_distance = 100000
  110. min_line = None
  111. for l in block['lines']:
  112. line_box = l['bbox']
  113. distance = math.sqrt((line_box[0] - imgbox[0])**2 + (line_box[1] - imgbox[1])**2)
  114. if distance < min_distance:
  115. min_distance = distance
  116. min_line = l
  117. if min_line:
  118. line_txt = "".join([s['text'] for s in min_line['spans']])
  119. img_h = imgbox[3] - imgbox[1]
  120. if min_distance<img_h: # 文字在图片前面
  121. page_md = __insert_after(page_md, img_content, line_txt)
  122. else:
  123. page_md = __insert_before(page_md, img_content, line_txt)
  124. else:
  125. logger.error(f"Can't find the location of image {img['image_path']} in the markdown file #1")
  126. else:# 应当在两个block之间
  127. # 找到上方最近的block,如果上方没有就找大下方最近的block
  128. top_txt_block = find_top_nearest_text_bbox(pymu_raw_blocks, imgbox)
  129. if top_txt_block:
  130. line_txt = "".join([s['text'] for s in top_txt_block['lines'][-1]['spans']])
  131. page_md = __insert_after(page_md, img_content, line_txt)
  132. else:
  133. bottom_txt_block = find_bottom_nearest_text_bbox(pymu_raw_blocks, imgbox)
  134. if bottom_txt_block:
  135. line_txt = "".join([s['text'] for s in bottom_txt_block['lines'][0]['spans']])
  136. page_md = __insert_before(page_md, img_content, line_txt)
  137. else:
  138. logger.error(f"Can't find the location of image {img['image_path']} in the markdown file #2")
  139. content_lst.append(page_md)
  140. """拼装成全部页面的文本"""
  141. content_text = "\n\n".join(content_lst)
  142. return content_text
  143. def __insert_after_para(text, image_path, content_list):
  144. """
  145. 在content_list中找到text,将image_path作为一个新的node插入到text后面
  146. """
  147. for i, c in enumerate(content_list):
  148. content_type = c.get("type")
  149. if content_type in UNI_FORMAT_TEXT_TYPE and text in c.get("text", ''):
  150. img_node = {
  151. "type": "image",
  152. "img_path": image_path,
  153. "img_alt":"",
  154. "img_title":"",
  155. "img_caption":""
  156. }
  157. content_list.insert(i+1, img_node)
  158. break
  159. else:
  160. logger.error(f"Can't find the location of image {image_path} in the markdown file, search target is {text}")
  161. def __insert_before_para(text, image_path, content_list):
  162. """
  163. 在content_list中找到text,将image_path作为一个新的node插入到text前面
  164. """
  165. for i, c in enumerate(content_list):
  166. content_type = c.get("type")
  167. if content_type in UNI_FORMAT_TEXT_TYPE and text in c.get("text", ''):
  168. img_node = {
  169. "type": "image",
  170. "img_path": image_path,
  171. "img_alt":"",
  172. "img_title":"",
  173. "img_caption":""
  174. }
  175. content_list.insert(i, img_node)
  176. break
  177. else:
  178. logger.error(f"Can't find the location of image {image_path} in the markdown file, search target is {text}")
  179. def mk_universal_format(para_dict: dict):
  180. """
  181. 构造统一格式 https://aicarrier.feishu.cn/wiki/FqmMwcH69iIdCWkkyjvcDwNUnTY
  182. """
  183. content_lst = []
  184. for _, page_info in para_dict.items():
  185. page_lst = [] # 一个page内的段落列表
  186. para_blocks = page_info.get("para_blocks")
  187. pymu_raw_blocks = page_info.get("preproc_blocks")
  188. all_page_images = []
  189. all_page_images.extend(page_info.get("images",[]))
  190. all_page_images.extend(page_info.get("image_backup", []) )
  191. all_page_images.extend(page_info.get("tables",[]))
  192. all_page_images.extend(page_info.get("table_backup",[]) )
  193. if not para_blocks or not pymu_raw_blocks: # 只有图片的拼接的场景
  194. for img in all_page_images:
  195. content_node = {
  196. "type": "image",
  197. "img_path": img['image_path'],
  198. "img_alt":"",
  199. "img_title":"",
  200. "img_caption":""
  201. }
  202. page_lst.append(content_node) # TODO 图片顺序
  203. else:
  204. for block in para_blocks:
  205. item = block["paras"]
  206. for _, p in item.items():
  207. font_type = p['para_font_type']# 对于文本来说,要么是普通文本,要么是个行间公式
  208. if font_type == TYPE_INTERLINE_EQUATION:
  209. content_node = {
  210. "type": "equation",
  211. "latex": p["para_text"]
  212. }
  213. page_lst.append(content_node)
  214. else:
  215. para_text = p["para_text"]
  216. is_title = p["is_para_title"]
  217. title_level = p['para_title_level']
  218. if is_title:
  219. content_node = {
  220. "type": f"h{title_level}",
  221. "text": para_text
  222. }
  223. page_lst.append(content_node)
  224. else:
  225. content_node = {
  226. "type": "text",
  227. "text": para_text
  228. }
  229. page_lst.append(content_node)
  230. content_lst.extend(page_lst)
  231. """插入图片"""
  232. for img in all_page_images:
  233. imgbox = img['bbox']
  234. img_content = f"{img['image_path']}"
  235. # 先看在哪个block内
  236. for block in pymu_raw_blocks:
  237. bbox = block['bbox']
  238. if bbox[0]-1 <= imgbox[0] < bbox[2]+1 and bbox[1]-1 <= imgbox[1] < bbox[3]+1:# 确定在这个大的block内,然后进入逐行比较距离
  239. for l in block['lines']:
  240. line_box = l['bbox']
  241. if line_box[0]-1 <= imgbox[0] < line_box[2]+1 and line_box[1]-1 <= imgbox[1] < line_box[3]+1: # 在line内的,插入line前面
  242. line_txt = "".join([s['text'] for s in l['spans']])
  243. __insert_before_para(line_txt, img_content, content_lst)
  244. break
  245. break
  246. else:# 在行与行之间
  247. # 找到图片x0,y0与line的x0,y0最近的line
  248. min_distance = 100000
  249. min_line = None
  250. for l in block['lines']:
  251. line_box = l['bbox']
  252. distance = math.sqrt((line_box[0] - imgbox[0])**2 + (line_box[1] - imgbox[1])**2)
  253. if distance < min_distance:
  254. min_distance = distance
  255. min_line = l
  256. if min_line:
  257. line_txt = "".join([s['text'] for s in min_line['spans']])
  258. img_h = imgbox[3] - imgbox[1]
  259. if min_distance<img_h: # 文字在图片前面
  260. __insert_after_para(line_txt, img_content, content_lst)
  261. else:
  262. __insert_before_para(line_txt, img_content, content_lst)
  263. break
  264. else:
  265. logger.error(f"Can't find the location of image {img['image_path']} in the markdown file #1")
  266. else:# 应当在两个block之间
  267. # 找到上方最近的block,如果上方没有就找大下方最近的block
  268. top_txt_block = find_top_nearest_text_bbox(pymu_raw_blocks, imgbox)
  269. if top_txt_block:
  270. line_txt = "".join([s['text'] for s in top_txt_block['lines'][-1]['spans']])
  271. __insert_after_para(line_txt, img_content, content_lst)
  272. else:
  273. bottom_txt_block = find_bottom_nearest_text_bbox(pymu_raw_blocks, imgbox)
  274. if bottom_txt_block:
  275. line_txt = "".join([s['text'] for s in bottom_txt_block['lines'][0]['spans']])
  276. __insert_before_para(line_txt, img_content, content_lst)
  277. else: # TODO ,图片可能独占一列,这种情况上下是没有图片的
  278. logger.error(f"Can't find the location of image {img['image_path']} in the markdown file #2")
  279. # end for
  280. return content_lst
  281. def mk_mm_markdown(content_list):
  282. """
  283. 基于同一格式的内容列表,构造markdown,含图片
  284. """
  285. content_md = []
  286. for c in content_list:
  287. content_type = c.get("type")
  288. if content_type == "text":
  289. content_md.append(c.get("text"))
  290. elif content_type == "equation":
  291. content = c.get("latex")
  292. if content.startswith("$$") and content.endswith("$$"):
  293. content_md.append(content)
  294. else:
  295. content_md.append(f"\n$$\n{c.get('latex')}\n$$\n")
  296. elif content_type in UNI_FORMAT_TEXT_TYPE:
  297. content_md.append(f"{'#'*int(content_type[1])} {c.get('text')}")
  298. elif content_type == "image":
  299. content_md.append(f"![]({c.get('img_path')})")
  300. return "\n\n".join(content_md)
  301. def mk_nlp_markdown(content_list):
  302. """
  303. 基于同一格式的内容列表,构造markdown,不含图片
  304. """
  305. content_md = []
  306. for c in content_list:
  307. content_type = c.get("type")
  308. if content_type == "text":
  309. content_md.append(c.get("text"))
  310. elif content_type == "equation":
  311. content_md.append(f"$$\n{c.get('latex')}\n$$")
  312. elif content_type in UNI_FORMAT_TEXT_TYPE:
  313. content_md.append(f"{'#'*int(content_type[1])} {c.get('text')}")
  314. return "\n\n".join(content_md)