ocr_mkcontent.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363
  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. line_text = ""
  131. line_lang = ""
  132. for span in line['spans']:
  133. span_type = span['type']
  134. if span_type == ContentType.Text:
  135. line_text += span['content'].strip()
  136. if line_text != "":
  137. line_lang = detect_lang(line_text)
  138. for span in line['spans']:
  139. span_type = span['type']
  140. content = ''
  141. if span_type == ContentType.Text:
  142. content = span['content']
  143. language = detect_lang(content)
  144. if language == 'en': # 只对英文长词进行分词处理,中文分词会丢失文本
  145. content = ocr_escape_special_markdown_char(split_long_words(content))
  146. else:
  147. content = ocr_escape_special_markdown_char(content)
  148. elif span_type == ContentType.InlineEquation:
  149. content = f"${span['content']}$"
  150. elif span_type == ContentType.InterlineEquation:
  151. content = f"\n$$\n{span['content']}\n$$\n"
  152. if content != '':
  153. if 'zh' in line_lang: # 遇到一些一个字一个span的文档,这种单字语言判断不准,需要用整行文本判断
  154. para_text += content # 中文语境下,content间不需要空格分隔
  155. else:
  156. para_text += content + ' ' # 英文语境下 content间需要空格分隔
  157. return para_text
  158. def para_to_standard_format(para, img_buket_path):
  159. para_content = {}
  160. if len(para) == 1:
  161. para_content = line_to_standard_format(para[0], img_buket_path)
  162. elif len(para) > 1:
  163. para_text = ''
  164. inline_equation_num = 0
  165. for line in para:
  166. for span in line['spans']:
  167. language = ''
  168. span_type = span.get('type')
  169. content = ""
  170. if span_type == ContentType.Text:
  171. content = span['content']
  172. language = detect_lang(content)
  173. if language == 'en': # 只对英文长词进行分词处理,中文分词会丢失文本
  174. content = ocr_escape_special_markdown_char(split_long_words(content))
  175. else:
  176. content = ocr_escape_special_markdown_char(content)
  177. elif span_type == ContentType.InlineEquation:
  178. content = f"${span['content']}$"
  179. inline_equation_num += 1
  180. if language == 'en': # 英文语境下 content间需要空格分隔
  181. para_text += content + ' '
  182. else: # 中文语境下,content间不需要空格分隔
  183. para_text += content
  184. para_content = {
  185. 'type': 'text',
  186. 'text': para_text,
  187. 'inline_equation_num': inline_equation_num
  188. }
  189. return para_content
  190. def para_to_standard_format_v2(para_block, img_buket_path):
  191. para_type = para_block['type']
  192. if para_type == BlockType.Text:
  193. para_content = {
  194. 'type': 'text',
  195. 'text': merge_para_with_text(para_block),
  196. }
  197. elif para_type == BlockType.Title:
  198. para_content = {
  199. 'type': 'text',
  200. 'text': merge_para_with_text(para_block),
  201. 'text_level': 1
  202. }
  203. elif para_type == BlockType.InterlineEquation:
  204. para_content = {
  205. 'type': 'equation',
  206. 'text': merge_para_with_text(para_block),
  207. 'text_format': "latex"
  208. }
  209. elif para_type == BlockType.Image:
  210. para_content = {
  211. 'type': 'image',
  212. }
  213. for block in para_block['blocks']:
  214. if block['type'] == BlockType.ImageBody:
  215. para_content['img_path'] = join_path(img_buket_path, block["lines"][0]["spans"][0]['image_path'])
  216. if block['type'] == BlockType.ImageCaption:
  217. para_content['img_caption'] = merge_para_with_text(block)
  218. elif para_type == BlockType.Table:
  219. para_content = {
  220. 'type': 'table',
  221. }
  222. for block in para_block['blocks']:
  223. if block['type'] == BlockType.TableBody:
  224. para_content['img_path'] = join_path(img_buket_path, block["lines"][0]["spans"][0]['image_path'])
  225. if block['type'] == BlockType.TableCaption:
  226. para_content['table_caption'] = merge_para_with_text(block)
  227. if block['type'] == BlockType.TableFootnote:
  228. para_content['table_footnote'] = merge_para_with_text(block)
  229. return para_content
  230. def make_standard_format_with_para(pdf_info_dict: list, img_buket_path: str):
  231. content_list = []
  232. for page_info in pdf_info_dict:
  233. paras_of_layout = page_info.get("para_blocks")
  234. if not paras_of_layout:
  235. continue
  236. for para_block in paras_of_layout:
  237. para_content = para_to_standard_format_v2(para_block, img_buket_path)
  238. content_list.append(para_content)
  239. return content_list
  240. def line_to_standard_format(line, img_buket_path):
  241. line_text = ""
  242. inline_equation_num = 0
  243. for span in line['spans']:
  244. if not span.get('content'):
  245. if not span.get('image_path'):
  246. continue
  247. else:
  248. if span['type'] == ContentType.Image:
  249. content = {
  250. 'type': 'image',
  251. 'img_path': join_path(img_buket_path, span['image_path'])
  252. }
  253. return content
  254. elif span['type'] == ContentType.Table:
  255. content = {
  256. 'type': 'table',
  257. 'img_path': join_path(img_buket_path, span['image_path'])
  258. }
  259. return content
  260. else:
  261. if span['type'] == ContentType.InterlineEquation:
  262. interline_equation = span['content']
  263. content = {
  264. 'type': 'equation',
  265. 'latex': f"$$\n{interline_equation}\n$$"
  266. }
  267. return content
  268. elif span['type'] == ContentType.InlineEquation:
  269. inline_equation = span['content']
  270. line_text += f"${inline_equation}$"
  271. inline_equation_num += 1
  272. elif span['type'] == ContentType.Text:
  273. text_content = ocr_escape_special_markdown_char(span['content']) # 转义特殊符号
  274. line_text += text_content
  275. content = {
  276. 'type': 'text',
  277. 'text': line_text,
  278. 'inline_equation_num': inline_equation_num
  279. }
  280. return content
  281. def ocr_mk_mm_standard_format(pdf_info_dict: list):
  282. """
  283. content_list
  284. type string image/text/table/equation(行间的单独拿出来,行内的和text合并)
  285. latex string latex文本字段。
  286. text string 纯文本格式的文本数据。
  287. md string markdown格式的文本数据。
  288. img_path string s3://full/path/to/img.jpg
  289. """
  290. content_list = []
  291. for page_info in pdf_info_dict:
  292. blocks = page_info.get("preproc_blocks")
  293. if not blocks:
  294. continue
  295. for block in blocks:
  296. for line in block['lines']:
  297. content = line_to_standard_format(line)
  298. content_list.append(content)
  299. return content_list
  300. def union_make(pdf_info_dict: list, make_mode: str, drop_mode: str, img_buket_path: str = ""):
  301. output_content = []
  302. for page_info in pdf_info_dict:
  303. if page_info.get("need_drop", False):
  304. drop_reason = page_info.get("drop_reason")
  305. if drop_mode == DropMode.NONE:
  306. pass
  307. elif drop_mode == DropMode.WHOLE_PDF:
  308. raise Exception(f"drop_mode is {DropMode.WHOLE_PDF} , drop_reason is {drop_reason}")
  309. elif drop_mode == DropMode.SINGLE_PAGE:
  310. logger.warning(f"drop_mode is {DropMode.SINGLE_PAGE} , drop_reason is {drop_reason}")
  311. continue
  312. else:
  313. raise Exception(f"drop_mode can not be null")
  314. paras_of_layout = page_info.get("para_blocks")
  315. if not paras_of_layout:
  316. continue
  317. if make_mode == MakeMode.MM_MD:
  318. page_markdown = ocr_mk_markdown_with_para_core_v2(paras_of_layout, "mm", img_buket_path)
  319. output_content.extend(page_markdown)
  320. elif make_mode == MakeMode.NLP_MD:
  321. page_markdown = ocr_mk_markdown_with_para_core_v2(paras_of_layout, "nlp")
  322. output_content.extend(page_markdown)
  323. elif make_mode == MakeMode.STANDARD_FORMAT:
  324. for para_block in paras_of_layout:
  325. para_content = para_to_standard_format_v2(para_block, img_buket_path)
  326. output_content.append(para_content)
  327. if make_mode in [MakeMode.MM_MD, MakeMode.NLP_MD]:
  328. return '\n\n'.join(output_content)
  329. elif make_mode == MakeMode.STANDARD_FORMAT:
  330. return output_content