para_split.py 7.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209
  1. from sklearn.cluster import DBSCAN
  2. import numpy as np
  3. from loguru import logger
  4. from magic_pdf.libs.boxbase import _is_in
  5. from magic_pdf.libs.ocr_content_type import ContentType
  6. LINE_STOP_FLAG = ['.', '!', '?', '。', '!', '?',":", ":", ")", ")", ";"]
  7. INLINE_EQUATION = ContentType.InlineEquation
  8. INTERLINE_EQUATION = ContentType.InterlineEquation
  9. TEXT = "text"
  10. def __add_line_period(blocks, layout_bboxes):
  11. """
  12. 为每行添加句号
  13. 如果这个行
  14. 1. 以行内公式结尾,但没有任何标点符号,此时加个句号,认为他就是段落结尾。
  15. """
  16. for block in blocks:
  17. for line in block['lines']:
  18. last_span = line['spans'][-1]
  19. span_type = last_span['type']
  20. if span_type in [INLINE_EQUATION]:
  21. span_content = last_span['content'].strip()
  22. if span_type==INLINE_EQUATION and span_content[-1] not in LINE_STOP_FLAG:
  23. if span_type in [INLINE_EQUATION, INTERLINE_EQUATION]:
  24. last_span['content'] = span_content + '.'
  25. def __valign_lines(blocks, layout_bboxes):
  26. """
  27. 在一个layoutbox内对齐行的左侧和右侧。
  28. 扫描行的左侧和右侧,如果x0, x1差距不超过一个阈值,就强行对齐到所处layout的左右两侧(和layout有一段距离)。
  29. 3是个经验值,TODO,计算得来,可以设置为1.5个正文字符。
  30. """
  31. min_distance = 3
  32. min_sample = 2
  33. for layout_box in layout_bboxes:
  34. blocks_in_layoutbox = [b for b in blocks if _is_in(b['bbox'], layout_box['layout_bbox'])]
  35. if len(blocks_in_layoutbox)==0:
  36. continue
  37. x0_lst = np.array([[line['bbox'][0], 0] for block in blocks_in_layoutbox for line in block['lines']])
  38. x1_lst = np.array([[line['bbox'][2], 0] for block in blocks_in_layoutbox for line in block['lines']])
  39. x0_clusters = DBSCAN(eps=min_distance, min_samples=min_sample).fit(x0_lst)
  40. x1_clusters = DBSCAN(eps=min_distance, min_samples=min_sample).fit(x1_lst)
  41. x0_uniq_label = np.unique(x0_clusters.labels_)
  42. x1_uniq_label = np.unique(x1_clusters.labels_)
  43. x0_2_new_val = {} # 存储旧值对应的新值映射
  44. x1_2_new_val = {}
  45. for label in x0_uniq_label:
  46. if label==-1:
  47. continue
  48. x0_index_of_label = np.where(x0_clusters.labels_==label)
  49. x0_raw_val = x0_lst[x0_index_of_label][:,0]
  50. x0_new_val = np.min(x0_lst[x0_index_of_label][:,0])
  51. x0_2_new_val.update({idx: x0_new_val for idx in x0_raw_val})
  52. for label in x1_uniq_label:
  53. if label==-1:
  54. continue
  55. x1_index_of_label = np.where(x1_clusters.labels_==label)
  56. x1_raw_val = x1_lst[x1_index_of_label][:,0]
  57. x1_new_val = np.max(x1_lst[x1_index_of_label][:,0])
  58. x1_2_new_val.update({idx: x1_new_val for idx in x1_raw_val})
  59. for block in blocks_in_layoutbox:
  60. for line in block['lines']:
  61. x0, x1 = line['bbox'][0], line['bbox'][2]
  62. if x0 in x0_2_new_val:
  63. line['bbox'][0] = int(x0_2_new_val[x0])
  64. if x1 in x1_2_new_val:
  65. line['bbox'][2] = int(x1_2_new_val[x1])
  66. # 其余对不齐的保持不动
  67. # 由于修改了block里的line长度,现在需要重新计算block的bbox
  68. for block in blocks_in_layoutbox:
  69. block['bbox'] = [min([line['bbox'][0] for line in block['lines']]),
  70. min([line['bbox'][1] for line in block['lines']]),
  71. max([line['bbox'][2] for line in block['lines']]),
  72. max([line['bbox'][3] for line in block['lines']])]
  73. def __common_pre_proc(blocks, layout_bboxes):
  74. """
  75. 不分语言的,对文本进行预处理
  76. """
  77. __add_line_period(blocks, layout_bboxes)
  78. __valign_lines(blocks, layout_bboxes)
  79. def __pre_proc_zh_blocks(blocks, layout_bboxes):
  80. """
  81. 对中文文本进行分段预处理
  82. """
  83. pass
  84. def __pre_proc_en_blocks(blocks, layout_bboxes):
  85. """
  86. 对英文文本进行分段预处理
  87. """
  88. pass
  89. def __group_line_by_layout(blocks, layout_bboxes, lang="en"):
  90. """
  91. 每个layout内的行进行聚合
  92. """
  93. # 因为只是一个block一行目前, 一个block就是一个段落
  94. lines_group = []
  95. for lyout in layout_bboxes:
  96. lines = [line for block in blocks if _is_in(block['bbox'], lyout['layout_bbox']) for line in block['lines']]
  97. lines_group.append(lines)
  98. return lines_group
  99. def __split_para_in_layoutbox(lines_group, layout_bboxes, lang="en", char_avg_len=10):
  100. """
  101. lines_group 进行行分段——layout内部进行分段。
  102. 1. 先计算每个group的左右边界。
  103. 2. 然后根据行末尾特征进行分段。
  104. 末尾特征:以句号等结束符结尾。并且距离右侧边界有一定距离。
  105. """
  106. def get_span_text(span):
  107. c = span.get('content', '')
  108. if len(c)==0:
  109. c = span.get('image-path', '')
  110. return c
  111. paras = []
  112. right_tail_distance = 1.5 * char_avg_len
  113. for lines in lines_group:
  114. if len(lines)==0:
  115. continue
  116. layout_right = max([line['bbox'][2] for line in lines])
  117. para = [] # 元素是line
  118. for line in lines:
  119. line_text = ''.join([get_span_text(span) for span in line['spans']])
  120. #logger.info(line_text)
  121. last_span_type = line['spans'][-1]['type']
  122. if last_span_type in [TEXT, INLINE_EQUATION]:
  123. last_char = line['spans'][-1]['content'][-1]
  124. if last_char in LINE_STOP_FLAG or line['bbox'][2] < layout_right - right_tail_distance:
  125. para.append(line)
  126. paras.append(para)
  127. # para_text = ''.join([span['content'] for line in para for span in line['spans']])
  128. # logger.info(para_text)
  129. para = []
  130. else:
  131. para.append(line)
  132. else: # 其他,图片、表格、行间公式,各自占一段
  133. if len(para)>0:
  134. paras.append(para)
  135. para = []
  136. else:
  137. paras.append([line])
  138. para = []
  139. # para_text = ''.join([get_span_text(span) for line in para for span in line['spans']])
  140. # logger.info(para_text)
  141. if len(para)>0:
  142. paras.append(para)
  143. # para_text = ''.join([get_span_text(span) for line in para for span in line['spans']])
  144. # logger.info(para_text)
  145. para = []
  146. return paras
  147. def __do_split(blocks, layout_bboxes, lang="en"):
  148. """
  149. 根据line和layout情况进行分段
  150. 先实现一个根据行末尾特征分段的简单方法。
  151. """
  152. """
  153. 算法思路:
  154. 1. 扫描layout里每一行,找出来行尾距离layout有边界有一定距离的行。
  155. 2. 从上述行中找到末尾是句号等可作为断行标志的行。
  156. 3. 参照上述行尾特征进行分段。
  157. 4. 图、表,目前独占一行,不考虑分段。
  158. """
  159. lines_group = __group_line_by_layout(blocks, layout_bboxes, lang) # block内分段
  160. layout_paras = __split_para_in_layoutbox(lines_group, layout_bboxes, lang) # block间连接分段
  161. return layout_paras
  162. def para_split(blocks, layout_bboxes, lang="en"):
  163. """
  164. 根据line和layout情况进行分段
  165. """
  166. __common_pre_proc(blocks, layout_bboxes)
  167. if lang=='en':
  168. __do_split(blocks, layout_bboxes, lang)
  169. elif lang=='zh':
  170. __do_split(blocks, layout_bboxes, lang)
  171. splited_blocks = __do_split(blocks, layout_bboxes, lang)
  172. return splited_blocks