ocr_mkcontent.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267
  1. import re
  2. from loguru import logger
  3. from magic_pdf.libs.commons import join_path
  4. from magic_pdf.libs.language import detect_lang
  5. from magic_pdf.libs.MakeContentConfig import DropMode, MakeMode
  6. from magic_pdf.libs.markdown_utils import ocr_escape_special_markdown_char
  7. from magic_pdf.libs.ocr_content_type import BlockType, ContentType
  8. from magic_pdf.para.para_split_v3 import ListLineTag
  9. def __is_hyphen_at_line_end(line):
  10. """
  11. Check if a line ends with one or more letters followed by a hyphen.
  12. Args:
  13. line (str): The line of text to check.
  14. Returns:
  15. bool: True if the line ends with one or more letters followed by a hyphen, False otherwise.
  16. """
  17. # Use regex to check if the line ends with one or more letters followed by a hyphen
  18. return bool(re.search(r'[A-Za-z]+-\s*$', line))
  19. def ocr_mk_mm_markdown_with_para_and_pagination(pdf_info_dict: list,
  20. img_buket_path):
  21. markdown_with_para_and_pagination = []
  22. page_no = 0
  23. for page_info in pdf_info_dict:
  24. paras_of_layout = page_info.get('para_blocks')
  25. if not paras_of_layout:
  26. continue
  27. page_markdown = ocr_mk_markdown_with_para_core_v2(
  28. paras_of_layout, 'mm', img_buket_path)
  29. markdown_with_para_and_pagination.append({
  30. 'page_no':
  31. page_no,
  32. 'md_content':
  33. '\n\n'.join(page_markdown)
  34. })
  35. page_no += 1
  36. return markdown_with_para_and_pagination
  37. def ocr_mk_markdown_with_para_core_v2(paras_of_layout,
  38. mode,
  39. img_buket_path='',
  40. parse_type="auto",
  41. lang=None
  42. ):
  43. page_markdown = []
  44. for para_block in paras_of_layout:
  45. para_text = ''
  46. para_type = para_block['type']
  47. if para_type in [BlockType.Text, BlockType.List, BlockType.Index]:
  48. para_text = merge_para_with_text(para_block, parse_type=parse_type, lang=lang)
  49. elif para_type == BlockType.Title:
  50. para_text = f'# {merge_para_with_text(para_block, parse_type=parse_type, lang=lang)}'
  51. elif para_type == BlockType.InterlineEquation:
  52. para_text = merge_para_with_text(para_block, parse_type=parse_type, lang=lang)
  53. elif para_type == BlockType.Image:
  54. if mode == 'nlp':
  55. continue
  56. elif mode == 'mm':
  57. for block in para_block['blocks']: # 1st.拼image_body
  58. if block['type'] == BlockType.ImageBody:
  59. for line in block['lines']:
  60. for span in line['spans']:
  61. if span['type'] == ContentType.Image:
  62. para_text += f"\n![]({join_path(img_buket_path, span['image_path'])}) \n"
  63. for block in para_block['blocks']: # 2nd.拼image_caption
  64. if block['type'] == BlockType.ImageCaption:
  65. para_text += merge_para_with_text(block, parse_type=parse_type, lang=lang)
  66. for block in para_block['blocks']: # 2nd.拼image_caption
  67. if block['type'] == BlockType.ImageFootnote:
  68. para_text += merge_para_with_text(block, parse_type=parse_type, lang=lang)
  69. elif para_type == BlockType.Table:
  70. if mode == 'nlp':
  71. continue
  72. elif mode == 'mm':
  73. for block in para_block['blocks']: # 1st.拼table_caption
  74. if block['type'] == BlockType.TableCaption:
  75. para_text += merge_para_with_text(block, parse_type=parse_type, lang=lang)
  76. for block in para_block['blocks']: # 2nd.拼table_body
  77. if block['type'] == BlockType.TableBody:
  78. for line in block['lines']:
  79. for span in line['spans']:
  80. if span['type'] == ContentType.Table:
  81. # if processed by table model
  82. if span.get('latex', ''):
  83. para_text += f"\n\n$\n {span['latex']}\n$\n\n"
  84. elif span.get('html', ''):
  85. para_text += f"\n\n{span['html']}\n\n"
  86. else:
  87. para_text += f"\n![]({join_path(img_buket_path, span['image_path'])}) \n"
  88. for block in para_block['blocks']: # 3rd.拼table_footnote
  89. if block['type'] == BlockType.TableFootnote:
  90. para_text += merge_para_with_text(block, parse_type=parse_type, lang=lang)
  91. if para_text.strip() == '':
  92. continue
  93. else:
  94. page_markdown.append(para_text.strip() + ' ')
  95. return page_markdown
  96. def detect_language(text):
  97. en_pattern = r'[a-zA-Z]+'
  98. en_matches = re.findall(en_pattern, text)
  99. en_length = sum(len(match) for match in en_matches)
  100. if len(text) > 0:
  101. if en_length / len(text) >= 0.5:
  102. return 'en'
  103. else:
  104. return 'unknown'
  105. else:
  106. return 'empty'
  107. def merge_para_with_text(para_block, parse_type="auto", lang=None):
  108. para_text = ''
  109. for i, line in enumerate(para_block['lines']):
  110. if i >= 1 and line.get(ListLineTag.IS_LIST_START_LINE, False):
  111. para_text += ' \n'
  112. line_text = ''
  113. line_lang = ''
  114. for span in line['spans']:
  115. span_type = span['type']
  116. if span_type == ContentType.Text:
  117. line_text += span['content'].strip()
  118. if line_text != '':
  119. line_lang = detect_lang(line_text)
  120. for span in line['spans']:
  121. span_type = span['type']
  122. content = ''
  123. if span_type == ContentType.Text:
  124. content = ocr_escape_special_markdown_char(span['content'])
  125. elif span_type == ContentType.InlineEquation:
  126. content = f" ${span['content']}$ "
  127. elif span_type == ContentType.InterlineEquation:
  128. content = f"\n$$\n{span['content']}\n$$\n"
  129. if content != '':
  130. langs = ['zh', 'ja', 'ko']
  131. if line_lang in langs: # 遇到一些一个字一个span的文档,这种单字语言判断不准,需要用整行文本判断
  132. para_text += content # 中文/日语/韩文语境下,content间不需要空格分隔
  133. elif line_lang == 'en':
  134. # 如果是前一行带有-连字符,那么末尾不应该加空格
  135. if __is_hyphen_at_line_end(content):
  136. para_text += content[:-1]
  137. else:
  138. para_text += content + ' '
  139. else:
  140. para_text += content + ' ' # 西方文本语境下 content间需要空格分隔
  141. return para_text
  142. def para_to_standard_format_v2(para_block, img_buket_path, page_idx, parse_type="auto", lang=None, drop_reason=None):
  143. para_type = para_block['type']
  144. para_content = {}
  145. if para_type == BlockType.Text:
  146. para_content = {
  147. 'type': 'text',
  148. 'text': merge_para_with_text(para_block, parse_type=parse_type, lang=lang),
  149. }
  150. elif para_type == BlockType.Title:
  151. para_content = {
  152. 'type': 'text',
  153. 'text': merge_para_with_text(para_block, parse_type=parse_type, lang=lang),
  154. 'text_level': 1,
  155. }
  156. elif para_type == BlockType.InterlineEquation:
  157. para_content = {
  158. 'type': 'equation',
  159. 'text': merge_para_with_text(para_block, parse_type=parse_type, lang=lang),
  160. 'text_format': 'latex',
  161. }
  162. elif para_type == BlockType.Image:
  163. para_content = {'type': 'image'}
  164. for block in para_block['blocks']:
  165. if block['type'] == BlockType.ImageBody:
  166. para_content['img_path'] = join_path(
  167. img_buket_path,
  168. block['lines'][0]['spans'][0]['image_path'])
  169. if block['type'] == BlockType.ImageCaption:
  170. para_content['img_caption'] = merge_para_with_text(block, parse_type=parse_type, lang=lang)
  171. if block['type'] == BlockType.ImageFootnote:
  172. para_content['img_footnote'] = merge_para_with_text(block, parse_type=parse_type, lang=lang)
  173. elif para_type == BlockType.Table:
  174. para_content = {'type': 'table'}
  175. for block in para_block['blocks']:
  176. if block['type'] == BlockType.TableBody:
  177. if block["lines"][0]["spans"][0].get('latex', ''):
  178. para_content['table_body'] = f"\n\n$\n {block['lines'][0]['spans'][0]['latex']}\n$\n\n"
  179. elif block["lines"][0]["spans"][0].get('html', ''):
  180. para_content['table_body'] = f"\n\n{block['lines'][0]['spans'][0]['html']}\n\n"
  181. para_content['img_path'] = join_path(img_buket_path, block["lines"][0]["spans"][0]['image_path'])
  182. if block['type'] == BlockType.TableCaption:
  183. para_content['table_caption'] = merge_para_with_text(block, parse_type=parse_type, lang=lang)
  184. if block['type'] == BlockType.TableFootnote:
  185. para_content['table_footnote'] = merge_para_with_text(block, parse_type=parse_type, lang=lang)
  186. para_content['page_idx'] = page_idx
  187. if drop_reason is not None:
  188. para_content['drop_reason'] = drop_reason
  189. return para_content
  190. def union_make(pdf_info_dict: list,
  191. make_mode: str,
  192. drop_mode: str,
  193. img_buket_path: str = '',
  194. parse_type: str = "auto",
  195. lang=None):
  196. output_content = []
  197. for page_info in pdf_info_dict:
  198. drop_reason_flag = False
  199. drop_reason = None
  200. if page_info.get('need_drop', False):
  201. drop_reason = page_info.get('drop_reason')
  202. if drop_mode == DropMode.NONE:
  203. pass
  204. elif drop_mode == DropMode.NONE_WITH_REASON:
  205. drop_reason_flag = True
  206. elif drop_mode == DropMode.WHOLE_PDF:
  207. raise Exception((f'drop_mode is {DropMode.WHOLE_PDF} ,'
  208. f'drop_reason is {drop_reason}'))
  209. elif drop_mode == DropMode.SINGLE_PAGE:
  210. logger.warning((f'drop_mode is {DropMode.SINGLE_PAGE} ,'
  211. f'drop_reason is {drop_reason}'))
  212. continue
  213. else:
  214. raise Exception('drop_mode can not be null')
  215. paras_of_layout = page_info.get('para_blocks')
  216. page_idx = page_info.get('page_idx')
  217. if not paras_of_layout:
  218. continue
  219. if make_mode == MakeMode.MM_MD:
  220. page_markdown = ocr_mk_markdown_with_para_core_v2(
  221. paras_of_layout, 'mm', img_buket_path, parse_type=parse_type, lang=lang)
  222. output_content.extend(page_markdown)
  223. elif make_mode == MakeMode.NLP_MD:
  224. page_markdown = ocr_mk_markdown_with_para_core_v2(
  225. paras_of_layout, 'nlp', parse_type=parse_type, lang=lang)
  226. output_content.extend(page_markdown)
  227. elif make_mode == MakeMode.STANDARD_FORMAT:
  228. for para_block in paras_of_layout:
  229. if drop_reason_flag:
  230. para_content = para_to_standard_format_v2(
  231. para_block, img_buket_path, page_idx, parse_type=parse_type, lang=lang, drop_reason=drop_reason)
  232. else:
  233. para_content = para_to_standard_format_v2(
  234. para_block, img_buket_path, page_idx, parse_type=parse_type, lang=lang)
  235. output_content.append(para_content)
  236. if make_mode in [MakeMode.MM_MD, MakeMode.NLP_MD]:
  237. return '\n\n'.join(output_content)
  238. elif make_mode == MakeMode.STANDARD_FORMAT:
  239. return output_content