ocr_mkcontent.py 13 KB

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