ocr_mkcontent.py 16 KB

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