mkcontent.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396
  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, type, element, 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. if type == "image":
  151. content_node = {
  152. "type": "image",
  153. "img_path": element.get("image_path"),
  154. "img_alt": "",
  155. "img_title": "",
  156. "img_caption": "",
  157. }
  158. elif type == "table":
  159. content_node = {
  160. "type": "table",
  161. "img_path": element.get("image_path"),
  162. "table_latex": element.get("text"),
  163. "table_title": "",
  164. "table_caption": "",
  165. "table_quality": element.get("quality"),
  166. }
  167. content_list.insert(i+1, content_node)
  168. break
  169. else:
  170. logger.error(f"Can't find the location of image {element.get('image_path')} in the markdown file, search target is {text}")
  171. def __insert_before_para(text, type, element, content_list):
  172. """
  173. 在content_list中找到text,将image_path作为一个新的node插入到text前面
  174. """
  175. for i, c in enumerate(content_list):
  176. content_type = c.get("type")
  177. if content_type in UNI_FORMAT_TEXT_TYPE and text in c.get("text", ''):
  178. if type == "image":
  179. content_node = {
  180. "type": "image",
  181. "img_path": element.get("image_path"),
  182. "img_alt": "",
  183. "img_title": "",
  184. "img_caption": "",
  185. }
  186. elif type == "table":
  187. content_node = {
  188. "type": "table",
  189. "img_path": element.get("image_path"),
  190. "table_latex": element.get("text"),
  191. "table_title": "",
  192. "table_caption": "",
  193. "table_quality": element.get("quality"),
  194. }
  195. content_list.insert(i, content_node)
  196. break
  197. else:
  198. logger.error(f"Can't find the location of image {element.get('image_path')} in the markdown file, search target is {text}")
  199. def mk_universal_format(para_dict: dict):
  200. """
  201. 构造统一格式 https://aicarrier.feishu.cn/wiki/FqmMwcH69iIdCWkkyjvcDwNUnTY
  202. """
  203. content_lst = []
  204. for _, page_info in para_dict.items():
  205. page_lst = [] # 一个page内的段落列表
  206. para_blocks = page_info.get("para_blocks")
  207. pymu_raw_blocks = page_info.get("preproc_blocks")
  208. all_page_images = []
  209. all_page_images.extend(page_info.get("images",[]))
  210. all_page_images.extend(page_info.get("image_backup", []) )
  211. # all_page_images.extend(page_info.get("tables",[]))
  212. # all_page_images.extend(page_info.get("table_backup",[]) )
  213. all_page_tables = []
  214. all_page_tables.extend(page_info.get("tables", []))
  215. if not para_blocks or not pymu_raw_blocks: # 只有图片的拼接的场景
  216. for img in all_page_images:
  217. content_node = {
  218. "type": "image",
  219. "img_path": img['image_path'],
  220. "img_alt":"",
  221. "img_title":"",
  222. "img_caption":""
  223. }
  224. page_lst.append(content_node) # TODO 图片顺序
  225. for table in all_page_tables:
  226. content_node = {
  227. "type": "table",
  228. "img_path": table['image_path'],
  229. "table_latex": table.get("text"),
  230. "table_title": "",
  231. "table_caption": "",
  232. "table_quality": table.get("quality"),
  233. }
  234. page_lst.append(content_node) # TODO 图片顺序
  235. else:
  236. for block in para_blocks:
  237. item = block["paras"]
  238. for _, p in item.items():
  239. font_type = p['para_font_type']# 对于文本来说,要么是普通文本,要么是个行间公式
  240. if font_type == TYPE_INTERLINE_EQUATION:
  241. content_node = {
  242. "type": "equation",
  243. "latex": p["para_text"]
  244. }
  245. page_lst.append(content_node)
  246. else:
  247. para_text = p["para_text"]
  248. is_title = p["is_para_title"]
  249. title_level = p['para_title_level']
  250. if is_title:
  251. content_node = {
  252. "type": f"h{title_level}",
  253. "text": para_text
  254. }
  255. page_lst.append(content_node)
  256. else:
  257. content_node = {
  258. "type": "text",
  259. "text": para_text
  260. }
  261. page_lst.append(content_node)
  262. content_lst.extend(page_lst)
  263. """插入图片"""
  264. for img in all_page_images:
  265. insert_img_or_table("image", img, pymu_raw_blocks, content_lst)
  266. """插入表格"""
  267. for table in all_page_tables:
  268. insert_img_or_table("table", table, pymu_raw_blocks, content_lst)
  269. # end for
  270. return content_lst
  271. def insert_img_or_table(type, element, pymu_raw_blocks, content_lst):
  272. element_bbox = element['bbox']
  273. # 先看在哪个block内
  274. for block in pymu_raw_blocks:
  275. bbox = block['bbox']
  276. if bbox[0] - 1 <= element_bbox[0] < bbox[2] + 1 and bbox[1] - 1 <= element_bbox[1] < bbox[
  277. 3] + 1: # 确定在这个大的block内,然后进入逐行比较距离
  278. for l in block['lines']:
  279. line_box = l['bbox']
  280. if line_box[0] - 1 <= element_bbox[0] < line_box[2] + 1 and line_box[1] - 1 <= element_bbox[1] < line_box[
  281. 3] + 1: # 在line内的,插入line前面
  282. line_txt = "".join([s['text'] for s in l['spans']])
  283. __insert_before_para(line_txt, type, element, content_lst)
  284. break
  285. break
  286. else: # 在行与行之间
  287. # 找到图片x0,y0与line的x0,y0最近的line
  288. min_distance = 100000
  289. min_line = None
  290. for l in block['lines']:
  291. line_box = l['bbox']
  292. distance = math.sqrt((line_box[0] - element_bbox[0]) ** 2 + (line_box[1] - element_bbox[1]) ** 2)
  293. if distance < min_distance:
  294. min_distance = distance
  295. min_line = l
  296. if min_line:
  297. line_txt = "".join([s['text'] for s in min_line['spans']])
  298. img_h = element_bbox[3] - element_bbox[1]
  299. if min_distance < img_h: # 文字在图片前面
  300. __insert_after_para(line_txt, type, element, content_lst)
  301. else:
  302. __insert_before_para(line_txt, type, element, content_lst)
  303. break
  304. else:
  305. logger.error(f"Can't find the location of image {element.get('image_path')} in the markdown file #1")
  306. else: # 应当在两个block之间
  307. # 找到上方最近的block,如果上方没有就找大下方最近的block
  308. top_txt_block = find_top_nearest_text_bbox(pymu_raw_blocks, element_bbox)
  309. if top_txt_block:
  310. line_txt = "".join([s['text'] for s in top_txt_block['lines'][-1]['spans']])
  311. __insert_after_para(line_txt, type, element, content_lst)
  312. else:
  313. bottom_txt_block = find_bottom_nearest_text_bbox(pymu_raw_blocks, element_bbox)
  314. if bottom_txt_block:
  315. line_txt = "".join([s['text'] for s in bottom_txt_block['lines'][0]['spans']])
  316. __insert_before_para(line_txt, type, element, content_lst)
  317. else: # TODO ,图片可能独占一列,这种情况上下是没有图片的
  318. logger.error(f"Can't find the location of image {element.get('image_path')} in the markdown file #2")
  319. def mk_mm_markdown(content_list):
  320. """
  321. 基于同一格式的内容列表,构造markdown,含图片
  322. """
  323. content_md = []
  324. for c in content_list:
  325. content_type = c.get("type")
  326. if content_type == "text":
  327. content_md.append(c.get("text"))
  328. elif content_type == "equation":
  329. content = c.get("latex")
  330. if content.startswith("$$") and content.endswith("$$"):
  331. content_md.append(content)
  332. else:
  333. content_md.append(f"\n$$\n{c.get('latex')}\n$$\n")
  334. elif content_type in UNI_FORMAT_TEXT_TYPE:
  335. content_md.append(f"{'#'*int(content_type[1])} {c.get('text')}")
  336. elif content_type == "image":
  337. content_md.append(f"![]({c.get('img_path')})")
  338. return "\n\n".join(content_md)
  339. def mk_nlp_markdown(content_list):
  340. """
  341. 基于同一格式的内容列表,构造markdown,不含图片
  342. """
  343. content_md = []
  344. for c in content_list:
  345. content_type = c.get("type")
  346. if content_type == "text":
  347. content_md.append(c.get("text"))
  348. elif content_type == "equation":
  349. content_md.append(f"$$\n{c.get('latex')}\n$$")
  350. elif content_type == "table":
  351. content_md.append(f"$$$\n{c.get('table_latex')}\n$$$")
  352. elif content_type in UNI_FORMAT_TEXT_TYPE:
  353. content_md.append(f"{'#'*int(content_type[1])} {c.get('text')}")
  354. return "\n\n".join(content_md)