ocr_mkcontent.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142
  1. from magic_pdf.libs.commons import s3_image_save_path, join_path
  2. from magic_pdf.libs.markdown_utils import ocr_escape_special_markdown_char
  3. from magic_pdf.libs.ocr_content_type import ContentType
  4. def ocr_mk_nlp_markdown(pdf_info_dict: dict):
  5. markdown = []
  6. for _, page_info in pdf_info_dict.items():
  7. blocks = page_info.get("preproc_blocks")
  8. if not blocks:
  9. continue
  10. for block in blocks:
  11. for line in block['lines']:
  12. line_text = ''
  13. for span in line['spans']:
  14. if not span.get('content'):
  15. continue
  16. content = ocr_escape_special_markdown_char(span['content']) # 转义特殊符号
  17. if span['type'] == ContentType.InlineEquation:
  18. content = f"${content}$"
  19. elif span['type'] == ContentType.InterlineEquation:
  20. content = f"$$\n{content}\n$$"
  21. line_text += content + ' '
  22. # 在行末添加两个空格以强制换行
  23. markdown.append(line_text.strip() + ' ')
  24. return '\n'.join(markdown)
  25. def ocr_mk_mm_markdown(pdf_info_dict: dict):
  26. markdown = []
  27. for _, page_info in pdf_info_dict.items():
  28. blocks = page_info.get("preproc_blocks")
  29. if not blocks:
  30. continue
  31. for block in blocks:
  32. for line in block['lines']:
  33. line_text = ''
  34. for span in line['spans']:
  35. if not span.get('content'):
  36. if not span.get('image_path'):
  37. continue
  38. else:
  39. content = f"![]({join_path(s3_image_save_path, span['image_path'])})"
  40. else:
  41. content = ocr_escape_special_markdown_char(span['content']) # 转义特殊符号
  42. if span['type'] == ContentType.InlineEquation:
  43. content = f"${content}$"
  44. elif span['type'] == ContentType.InterlineEquation:
  45. content = f"$$\n{content}\n$$"
  46. line_text += content + ' '
  47. # 在行末添加两个空格以强制换行
  48. markdown.append(line_text.strip() + ' ')
  49. return '\n'.join(markdown)
  50. def ocr_mk_mm_markdown_with_para(pdf_info_dict: dict):
  51. markdown = []
  52. for _, page_info in pdf_info_dict.items():
  53. paras = page_info.get("para_blocks")
  54. if not paras:
  55. continue
  56. for para in paras:
  57. para_text = ''
  58. for line in para:
  59. for span in line['spans']:
  60. span_type = span.get('type')
  61. if span_type == ContentType.Text:
  62. para_text += span['content']
  63. elif span_type == ContentType.InlineEquation:
  64. para_text += f" ${span['content']}$ "
  65. elif span_type == ContentType.InterlineEquation:
  66. para_text += f"$$\n{span['content']}\n$$ "
  67. elif span_type == ContentType.Image:
  68. para_text += f"![]({join_path(s3_image_save_path, span['image_path'])})"
  69. markdown.append(para_text)
  70. return '\n\n'.join(markdown)
  71. def line_to_standard_format(line):
  72. line_text = ""
  73. inline_equation_num = 0
  74. for span in line['spans']:
  75. if not span.get('content'):
  76. if not span.get('image_path'):
  77. continue
  78. else:
  79. if span['type'] == ContentType.Image:
  80. content = {
  81. 'type': 'image',
  82. 'img_path': join_path(s3_image_save_path, span['image_path'])
  83. }
  84. return content
  85. elif span['type'] == ContentType.Table:
  86. content = {
  87. 'type': 'table',
  88. 'img_path': join_path(s3_image_save_path, span['image_path'])
  89. }
  90. return content
  91. else:
  92. if span['type'] == ContentType.InterlineEquation:
  93. interline_equation = ocr_escape_special_markdown_char(span['content']) # 转义特殊符号
  94. content = {
  95. 'type': 'equation',
  96. 'latex': f"$$\n{interline_equation}\n$$"
  97. }
  98. return content
  99. elif span['type'] == ContentType.InlineEquation:
  100. inline_equation = ocr_escape_special_markdown_char(span['content']) # 转义特殊符号
  101. line_text += f"${inline_equation}$"
  102. inline_equation_num += 1
  103. elif span['type'] == ContentType.Text:
  104. line_text += span['content']
  105. content = {
  106. 'type': 'text',
  107. 'text': line_text,
  108. 'inline_equation_num': inline_equation_num
  109. }
  110. return content
  111. def ocr_mk_mm_standard_format(pdf_info_dict: dict):
  112. '''
  113. content_list
  114. type string image/text/table/equation(行间的单独拿出来,行内的和text合并)
  115. latex string latex文本字段。
  116. text string 纯文本格式的文本数据。
  117. md string markdown格式的文本数据。
  118. img_path string s3://full/path/to/img.jpg
  119. '''
  120. content_list = []
  121. for _, page_info in pdf_info_dict.items():
  122. blocks = page_info.get("preproc_blocks")
  123. if not blocks:
  124. continue
  125. for block in blocks:
  126. for line in block['lines']:
  127. content = line_to_standard_format(line)
  128. content_list.append(content)
  129. return content_list