ocr_mkcontent.py 12 KB

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