ocr_mkcontent.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356
  1. from loguru import logger
  2. from magic_pdf.libs.MakeContentConfig import DropMode, MakeMode
  3. from magic_pdf.libs.commons import join_path
  4. from magic_pdf.libs.language import detect_lang
  5. from magic_pdf.libs.markdown_utils import ocr_escape_special_markdown_char
  6. from magic_pdf.libs.ocr_content_type import ContentType, BlockType
  7. import wordninja
  8. import re
  9. def split_long_words(text):
  10. segments = text.split(' ')
  11. for i in range(len(segments)):
  12. words = re.findall(r'\w+|[^\w]', segments[i], re.UNICODE)
  13. for j in range(len(words)):
  14. if len(words[j]) > 15:
  15. words[j] = ' '.join(wordninja.split(words[j]))
  16. segments[i] = ''.join(words)
  17. return ' '.join(segments)
  18. def ocr_mk_mm_markdown_with_para(pdf_info_list: list, img_buket_path):
  19. markdown = []
  20. for page_info in pdf_info_list:
  21. paras_of_layout = page_info.get("para_blocks")
  22. page_markdown = ocr_mk_markdown_with_para_core_v2(paras_of_layout, "mm", img_buket_path)
  23. markdown.extend(page_markdown)
  24. return '\n\n'.join(markdown)
  25. def ocr_mk_nlp_markdown_with_para(pdf_info_dict: list):
  26. markdown = []
  27. for page_info in pdf_info_dict:
  28. paras_of_layout = page_info.get("para_blocks")
  29. page_markdown = ocr_mk_markdown_with_para_core_v2(paras_of_layout, "nlp")
  30. markdown.extend(page_markdown)
  31. return '\n\n'.join(markdown)
  32. def ocr_mk_mm_markdown_with_para_and_pagination(pdf_info_dict: list, img_buket_path):
  33. markdown_with_para_and_pagination = []
  34. page_no = 0
  35. for page_info in pdf_info_dict:
  36. paras_of_layout = page_info.get("para_blocks")
  37. if not paras_of_layout:
  38. continue
  39. page_markdown = ocr_mk_markdown_with_para_core_v2(paras_of_layout, "mm", img_buket_path)
  40. markdown_with_para_and_pagination.append({
  41. 'page_no': page_no,
  42. 'md_content': '\n\n'.join(page_markdown)
  43. })
  44. page_no += 1
  45. return markdown_with_para_and_pagination
  46. def ocr_mk_markdown_with_para_core(paras_of_layout, mode, img_buket_path=""):
  47. page_markdown = []
  48. for paras in paras_of_layout:
  49. for para in paras:
  50. para_text = ''
  51. for line in para:
  52. for span in line['spans']:
  53. span_type = span.get('type')
  54. content = ''
  55. language = ''
  56. if span_type == ContentType.Text:
  57. content = span['content']
  58. language = detect_lang(content)
  59. if language == 'en': # 只对英文长词进行分词处理,中文分词会丢失文本
  60. content = ocr_escape_special_markdown_char(split_long_words(content))
  61. else:
  62. content = ocr_escape_special_markdown_char(content)
  63. elif span_type == ContentType.InlineEquation:
  64. content = f"${span['content']}$"
  65. elif span_type == ContentType.InterlineEquation:
  66. content = f"\n$$\n{span['content']}\n$$\n"
  67. elif span_type in [ContentType.Image, ContentType.Table]:
  68. if mode == 'mm':
  69. content = f"\n![]({join_path(img_buket_path, span['image_path'])})\n"
  70. elif mode == 'nlp':
  71. pass
  72. if content != '':
  73. if language == 'en': # 英文语境下 content间需要空格分隔
  74. para_text += content + ' '
  75. else: # 中文语境下,content间不需要空格分隔
  76. para_text += content
  77. if para_text.strip() == '':
  78. continue
  79. else:
  80. page_markdown.append(para_text.strip() + ' ')
  81. return page_markdown
  82. def ocr_mk_markdown_with_para_core_v2(paras_of_layout, mode, img_buket_path=""):
  83. page_markdown = []
  84. for para_block in paras_of_layout:
  85. para_text = ''
  86. para_type = para_block['type']
  87. if para_type == BlockType.Text:
  88. para_text = merge_para_with_text(para_block)
  89. elif para_type == BlockType.Title:
  90. para_text = f"# {merge_para_with_text(para_block)}"
  91. elif para_type == BlockType.InterlineEquation:
  92. para_text = merge_para_with_text(para_block)
  93. elif para_type == BlockType.Image:
  94. if mode == 'nlp':
  95. continue
  96. elif mode == 'mm':
  97. for block in para_block['blocks']: # 1st.拼image_body
  98. if block['type'] == BlockType.ImageBody:
  99. for line in block['lines']:
  100. for span in line['spans']:
  101. if span['type'] == ContentType.Image:
  102. para_text += f"\n![]({join_path(img_buket_path, span['image_path'])})\n"
  103. for block in para_block['blocks']: # 2nd.拼image_caption
  104. if block['type'] == BlockType.ImageCaption:
  105. para_text += merge_para_with_text(block)
  106. elif para_type == BlockType.Table:
  107. if mode == 'nlp':
  108. continue
  109. elif mode == 'mm':
  110. for block in para_block['blocks']: # 1st.拼table_caption
  111. if block['type'] == BlockType.TableCaption:
  112. para_text += merge_para_with_text(block)
  113. for block in para_block['blocks']: # 2nd.拼table_body
  114. if block['type'] == BlockType.TableBody:
  115. for line in block['lines']:
  116. for span in line['spans']:
  117. if span['type'] == ContentType.Table:
  118. para_text += f"\n![]({join_path(img_buket_path, span['image_path'])})\n"
  119. for block in para_block['blocks']: # 3rd.拼table_footnote
  120. if block['type'] == BlockType.TableFootnote:
  121. para_text += merge_para_with_text(block)
  122. if para_text.strip() == '':
  123. continue
  124. else:
  125. page_markdown.append(para_text.strip() + ' ')
  126. return page_markdown
  127. def merge_para_with_text(para_block):
  128. para_text = ''
  129. for line in para_block['lines']:
  130. for span in line['spans']:
  131. span_type = span['type']
  132. content = ''
  133. language = ''
  134. if span_type == ContentType.Text:
  135. content = span['content']
  136. language = detect_lang(content)
  137. if language == 'en': # 只对英文长词进行分词处理,中文分词会丢失文本
  138. content = ocr_escape_special_markdown_char(split_long_words(content))
  139. else:
  140. content = ocr_escape_special_markdown_char(content)
  141. elif span_type == ContentType.InlineEquation:
  142. content = f"${span['content']}$"
  143. elif span_type == ContentType.InterlineEquation:
  144. content = f"\n$$\n{span['content']}\n$$\n"
  145. if content != '':
  146. if 'zh' in language:
  147. para_text += content # 中文语境下,content间不需要空格分隔
  148. else:
  149. para_text += content + ' ' # 英文语境下 content间需要空格分隔
  150. return para_text
  151. def para_to_standard_format(para, img_buket_path):
  152. para_content = {}
  153. if len(para) == 1:
  154. para_content = line_to_standard_format(para[0], img_buket_path)
  155. elif len(para) > 1:
  156. para_text = ''
  157. inline_equation_num = 0
  158. for line in para:
  159. for span in line['spans']:
  160. language = ''
  161. span_type = span.get('type')
  162. content = ""
  163. if span_type == ContentType.Text:
  164. content = span['content']
  165. language = detect_lang(content)
  166. if language == 'en': # 只对英文长词进行分词处理,中文分词会丢失文本
  167. content = ocr_escape_special_markdown_char(split_long_words(content))
  168. else:
  169. content = ocr_escape_special_markdown_char(content)
  170. elif span_type == ContentType.InlineEquation:
  171. content = f"${span['content']}$"
  172. inline_equation_num += 1
  173. if language == 'en': # 英文语境下 content间需要空格分隔
  174. para_text += content + ' '
  175. else: # 中文语境下,content间不需要空格分隔
  176. para_text += content
  177. para_content = {
  178. 'type': 'text',
  179. 'text': para_text,
  180. 'inline_equation_num': inline_equation_num
  181. }
  182. return para_content
  183. def para_to_standard_format_v2(para_block, img_buket_path):
  184. para_type = para_block['type']
  185. if para_type == BlockType.Text:
  186. para_content = {
  187. 'type': 'text',
  188. 'text': merge_para_with_text(para_block),
  189. }
  190. elif para_type == BlockType.Title:
  191. para_content = {
  192. 'type': 'text',
  193. 'text': merge_para_with_text(para_block),
  194. 'text_level': 1
  195. }
  196. elif para_type == BlockType.InterlineEquation:
  197. para_content = {
  198. 'type': 'equation',
  199. 'text': merge_para_with_text(para_block),
  200. 'text_format': "latex"
  201. }
  202. elif para_type == BlockType.Image:
  203. para_content = {
  204. 'type': 'image',
  205. }
  206. for block in para_block['blocks']:
  207. if block['type'] == BlockType.ImageBody:
  208. para_content['img_path'] = join_path(img_buket_path, block["lines"][0]["spans"][0]['image_path'])
  209. if block['type'] == BlockType.ImageCaption:
  210. para_content['img_caption'] = merge_para_with_text(block)
  211. elif para_type == BlockType.Table:
  212. para_content = {
  213. 'type': 'table',
  214. }
  215. for block in para_block['blocks']:
  216. if block['type'] == BlockType.TableBody:
  217. para_content['img_path'] = join_path(img_buket_path, block["lines"][0]["spans"][0]['image_path'])
  218. if block['type'] == BlockType.TableCaption:
  219. para_content['table_caption'] = merge_para_with_text(block)
  220. if block['type'] == BlockType.TableFootnote:
  221. para_content['table_footnote'] = merge_para_with_text(block)
  222. return para_content
  223. def make_standard_format_with_para(pdf_info_dict: list, img_buket_path: str):
  224. content_list = []
  225. for page_info in pdf_info_dict:
  226. paras_of_layout = page_info.get("para_blocks")
  227. if not paras_of_layout:
  228. continue
  229. for para_block in paras_of_layout:
  230. para_content = para_to_standard_format_v2(para_block, img_buket_path)
  231. content_list.append(para_content)
  232. return content_list
  233. def line_to_standard_format(line, img_buket_path):
  234. line_text = ""
  235. inline_equation_num = 0
  236. for span in line['spans']:
  237. if not span.get('content'):
  238. if not span.get('image_path'):
  239. continue
  240. else:
  241. if span['type'] == ContentType.Image:
  242. content = {
  243. 'type': 'image',
  244. 'img_path': join_path(img_buket_path, span['image_path'])
  245. }
  246. return content
  247. elif span['type'] == ContentType.Table:
  248. content = {
  249. 'type': 'table',
  250. 'img_path': join_path(img_buket_path, span['image_path'])
  251. }
  252. return content
  253. else:
  254. if span['type'] == ContentType.InterlineEquation:
  255. interline_equation = span['content']
  256. content = {
  257. 'type': 'equation',
  258. 'latex': f"$$\n{interline_equation}\n$$"
  259. }
  260. return content
  261. elif span['type'] == ContentType.InlineEquation:
  262. inline_equation = span['content']
  263. line_text += f"${inline_equation}$"
  264. inline_equation_num += 1
  265. elif span['type'] == ContentType.Text:
  266. text_content = ocr_escape_special_markdown_char(span['content']) # 转义特殊符号
  267. line_text += text_content
  268. content = {
  269. 'type': 'text',
  270. 'text': line_text,
  271. 'inline_equation_num': inline_equation_num
  272. }
  273. return content
  274. def ocr_mk_mm_standard_format(pdf_info_dict: list):
  275. """
  276. content_list
  277. type string image/text/table/equation(行间的单独拿出来,行内的和text合并)
  278. latex string latex文本字段。
  279. text string 纯文本格式的文本数据。
  280. md string markdown格式的文本数据。
  281. img_path string s3://full/path/to/img.jpg
  282. """
  283. content_list = []
  284. for page_info in pdf_info_dict:
  285. blocks = page_info.get("preproc_blocks")
  286. if not blocks:
  287. continue
  288. for block in blocks:
  289. for line in block['lines']:
  290. content = line_to_standard_format(line)
  291. content_list.append(content)
  292. return content_list
  293. def union_make(pdf_info_dict: list, make_mode: str, drop_mode: str, img_buket_path: str = ""):
  294. output_content = []
  295. for page_info in pdf_info_dict:
  296. if page_info.get("need_drop", False):
  297. drop_reason = page_info.get("drop_reason")
  298. if drop_mode == DropMode.NONE:
  299. pass
  300. elif drop_mode == DropMode.WHOLE_PDF:
  301. raise Exception(f"drop_mode is {DropMode.WHOLE_PDF} , drop_reason is {drop_reason}")
  302. elif drop_mode == DropMode.SINGLE_PAGE:
  303. logger.warning(f"drop_mode is {DropMode.SINGLE_PAGE} , drop_reason is {drop_reason}")
  304. continue
  305. else:
  306. raise Exception(f"drop_mode can not be null")
  307. paras_of_layout = page_info.get("para_blocks")
  308. if not paras_of_layout:
  309. continue
  310. if make_mode == MakeMode.MM_MD:
  311. page_markdown = ocr_mk_markdown_with_para_core_v2(paras_of_layout, "mm", img_buket_path)
  312. output_content.extend(page_markdown)
  313. elif make_mode == MakeMode.NLP_MD:
  314. page_markdown = ocr_mk_markdown_with_para_core_v2(paras_of_layout, "nlp")
  315. output_content.extend(page_markdown)
  316. elif make_mode == MakeMode.STANDARD_FORMAT:
  317. for para_block in paras_of_layout:
  318. para_content = para_to_standard_format_v2(para_block, img_buket_path)
  319. output_content.append(para_content)
  320. if make_mode in [MakeMode.MM_MD, MakeMode.NLP_MD]:
  321. return '\n\n'.join(output_content)
  322. elif make_mode == MakeMode.STANDARD_FORMAT:
  323. return output_content