ocr_mkcontent.py 16 KB

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