ocr_mkcontent.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369
  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, page_idx):
  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. 'page_idx': page_idx
  197. }
  198. elif para_type == BlockType.Title:
  199. para_content = {
  200. 'type': 'text',
  201. 'text': merge_para_with_text(para_block),
  202. 'text_level': 1,
  203. 'page_idx': page_idx
  204. }
  205. elif para_type == BlockType.InterlineEquation:
  206. para_content = {
  207. 'type': 'equation',
  208. 'text': merge_para_with_text(para_block),
  209. 'text_format': "latex",
  210. 'page_idx': page_idx
  211. }
  212. elif para_type == BlockType.Image:
  213. para_content = {
  214. 'type': 'image',
  215. 'page_idx': page_idx
  216. }
  217. for block in para_block['blocks']:
  218. if block['type'] == BlockType.ImageBody:
  219. para_content['img_path'] = join_path(img_buket_path, block["lines"][0]["spans"][0]['image_path'])
  220. if block['type'] == BlockType.ImageCaption:
  221. para_content['img_caption'] = merge_para_with_text(block)
  222. elif para_type == BlockType.Table:
  223. para_content = {
  224. 'type': 'table',
  225. 'page_idx': page_idx
  226. }
  227. for block in para_block['blocks']:
  228. if block['type'] == BlockType.TableBody:
  229. para_content['img_path'] = join_path(img_buket_path, block["lines"][0]["spans"][0]['image_path'])
  230. if block['type'] == BlockType.TableCaption:
  231. para_content['table_caption'] = merge_para_with_text(block)
  232. if block['type'] == BlockType.TableFootnote:
  233. para_content['table_footnote'] = merge_para_with_text(block)
  234. return para_content
  235. def make_standard_format_with_para(pdf_info_dict: list, img_buket_path: str):
  236. content_list = []
  237. for page_info in pdf_info_dict:
  238. paras_of_layout = page_info.get("para_blocks")
  239. if not paras_of_layout:
  240. continue
  241. for para_block in paras_of_layout:
  242. para_content = para_to_standard_format_v2(para_block, img_buket_path)
  243. content_list.append(para_content)
  244. return content_list
  245. def line_to_standard_format(line, img_buket_path):
  246. line_text = ""
  247. inline_equation_num = 0
  248. for span in line['spans']:
  249. if not span.get('content'):
  250. if not span.get('image_path'):
  251. continue
  252. else:
  253. if span['type'] == ContentType.Image:
  254. content = {
  255. 'type': 'image',
  256. 'img_path': join_path(img_buket_path, span['image_path'])
  257. }
  258. return content
  259. elif span['type'] == ContentType.Table:
  260. content = {
  261. 'type': 'table',
  262. 'img_path': join_path(img_buket_path, span['image_path'])
  263. }
  264. return content
  265. else:
  266. if span['type'] == ContentType.InterlineEquation:
  267. interline_equation = span['content']
  268. content = {
  269. 'type': 'equation',
  270. 'latex': f"$$\n{interline_equation}\n$$"
  271. }
  272. return content
  273. elif span['type'] == ContentType.InlineEquation:
  274. inline_equation = span['content']
  275. line_text += f"${inline_equation}$"
  276. inline_equation_num += 1
  277. elif span['type'] == ContentType.Text:
  278. text_content = ocr_escape_special_markdown_char(span['content']) # 转义特殊符号
  279. line_text += text_content
  280. content = {
  281. 'type': 'text',
  282. 'text': line_text,
  283. 'inline_equation_num': inline_equation_num
  284. }
  285. return content
  286. def ocr_mk_mm_standard_format(pdf_info_dict: list):
  287. """
  288. content_list
  289. type string image/text/table/equation(行间的单独拿出来,行内的和text合并)
  290. latex string latex文本字段。
  291. text string 纯文本格式的文本数据。
  292. md string markdown格式的文本数据。
  293. img_path string s3://full/path/to/img.jpg
  294. """
  295. content_list = []
  296. for page_info in pdf_info_dict:
  297. blocks = page_info.get("preproc_blocks")
  298. if not blocks:
  299. continue
  300. for block in blocks:
  301. for line in block['lines']:
  302. content = line_to_standard_format(line)
  303. content_list.append(content)
  304. return content_list
  305. def union_make(pdf_info_dict: list, make_mode: str, drop_mode: str, img_buket_path: str = ""):
  306. output_content = []
  307. for page_info in pdf_info_dict:
  308. if page_info.get("need_drop", False):
  309. drop_reason = page_info.get("drop_reason")
  310. if drop_mode == DropMode.NONE:
  311. pass
  312. elif drop_mode == DropMode.WHOLE_PDF:
  313. raise Exception(f"drop_mode is {DropMode.WHOLE_PDF} , drop_reason is {drop_reason}")
  314. elif drop_mode == DropMode.SINGLE_PAGE:
  315. logger.warning(f"drop_mode is {DropMode.SINGLE_PAGE} , drop_reason is {drop_reason}")
  316. continue
  317. else:
  318. raise Exception(f"drop_mode can not be null")
  319. paras_of_layout = page_info.get("para_blocks")
  320. page_idx = page_info.get("page_idx")
  321. if not paras_of_layout:
  322. continue
  323. if make_mode == MakeMode.MM_MD:
  324. page_markdown = ocr_mk_markdown_with_para_core_v2(paras_of_layout, "mm", img_buket_path)
  325. output_content.extend(page_markdown)
  326. elif make_mode == MakeMode.NLP_MD:
  327. page_markdown = ocr_mk_markdown_with_para_core_v2(paras_of_layout, "nlp")
  328. output_content.extend(page_markdown)
  329. elif make_mode == MakeMode.STANDARD_FORMAT:
  330. for para_block in paras_of_layout:
  331. para_content = para_to_standard_format_v2(para_block, img_buket_path, page_idx)
  332. output_content.append(para_content)
  333. if make_mode in [MakeMode.MM_MD, MakeMode.NLP_MD]:
  334. return '\n\n'.join(output_content)
  335. elif make_mode == MakeMode.STANDARD_FORMAT:
  336. return output_content