para_split.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318
  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 __get_span_text(span):
  11. c = span.get('content', '')
  12. if len(c)==0:
  13. c = span.get('image_path', '')
  14. return c
  15. def __add_line_period(blocks, layout_bboxes):
  16. """
  17. 为每行添加句号
  18. 如果这个行
  19. 1. 以行内公式结尾,但没有任何标点符号,此时加个句号,认为他就是段落结尾。
  20. """
  21. for block in blocks:
  22. for line in block['lines']:
  23. last_span = line['spans'][-1]
  24. span_type = last_span['type']
  25. if span_type in [INLINE_EQUATION]:
  26. span_content = last_span['content'].strip()
  27. if span_type==INLINE_EQUATION and span_content[-1] not in LINE_STOP_FLAG:
  28. if span_type in [INLINE_EQUATION, INTERLINE_EQUATION]:
  29. last_span['content'] = span_content + '.'
  30. def __valign_lines(blocks, layout_bboxes):
  31. """
  32. 在一个layoutbox内对齐行的左侧和右侧。
  33. 扫描行的左侧和右侧,如果x0, x1差距不超过一个阈值,就强行对齐到所处layout的左右两侧(和layout有一段距离)。
  34. 3是个经验值,TODO,计算得来,可以设置为1.5个正文字符。
  35. """
  36. min_distance = 3
  37. min_sample = 2
  38. new_layout_bboxes = []
  39. for layout_box in layout_bboxes:
  40. blocks_in_layoutbox = [b for b in blocks if _is_in(b['bbox'], layout_box['layout_bbox'])]
  41. if len(blocks_in_layoutbox)==0:
  42. continue
  43. x0_lst = np.array([[line['bbox'][0], 0] for block in blocks_in_layoutbox for line in block['lines']])
  44. x1_lst = np.array([[line['bbox'][2], 0] for block in blocks_in_layoutbox for line in block['lines']])
  45. x0_clusters = DBSCAN(eps=min_distance, min_samples=min_sample).fit(x0_lst)
  46. x1_clusters = DBSCAN(eps=min_distance, min_samples=min_sample).fit(x1_lst)
  47. x0_uniq_label = np.unique(x0_clusters.labels_)
  48. x1_uniq_label = np.unique(x1_clusters.labels_)
  49. x0_2_new_val = {} # 存储旧值对应的新值映射
  50. x1_2_new_val = {}
  51. for label in x0_uniq_label:
  52. if label==-1:
  53. continue
  54. x0_index_of_label = np.where(x0_clusters.labels_==label)
  55. x0_raw_val = x0_lst[x0_index_of_label][:,0]
  56. x0_new_val = np.min(x0_lst[x0_index_of_label][:,0])
  57. x0_2_new_val.update({idx: x0_new_val for idx in x0_raw_val})
  58. for label in x1_uniq_label:
  59. if label==-1:
  60. continue
  61. x1_index_of_label = np.where(x1_clusters.labels_==label)
  62. x1_raw_val = x1_lst[x1_index_of_label][:,0]
  63. x1_new_val = np.max(x1_lst[x1_index_of_label][:,0])
  64. x1_2_new_val.update({idx: x1_new_val for idx in x1_raw_val})
  65. for block in blocks_in_layoutbox:
  66. for line in block['lines']:
  67. x0, x1 = line['bbox'][0], line['bbox'][2]
  68. if x0 in x0_2_new_val:
  69. line['bbox'][0] = int(x0_2_new_val[x0])
  70. if x1 in x1_2_new_val:
  71. line['bbox'][2] = int(x1_2_new_val[x1])
  72. # 其余对不齐的保持不动
  73. # 由于修改了block里的line长度,现在需要重新计算block的bbox
  74. for block in blocks_in_layoutbox:
  75. block['bbox'] = [min([line['bbox'][0] for line in block['lines']]),
  76. min([line['bbox'][1] for line in block['lines']]),
  77. max([line['bbox'][2] for line in block['lines']]),
  78. max([line['bbox'][3] for line in block['lines']])]
  79. """新计算layout的bbox,因为block的bbox变了。"""
  80. layout_x0 = min([block['bbox'][0] for block in blocks_in_layoutbox])
  81. layout_y0 = min([block['bbox'][1] for block in blocks_in_layoutbox])
  82. layout_x1 = max([block['bbox'][2] for block in blocks_in_layoutbox])
  83. layout_y1 = max([block['bbox'][3] for block in blocks_in_layoutbox])
  84. new_layout_bboxes.append([layout_x0, layout_y0, layout_x1, layout_y1])
  85. return new_layout_bboxes
  86. def __common_pre_proc(blocks, layout_bboxes):
  87. """
  88. 不分语言的,对文本进行预处理
  89. """
  90. #__add_line_period(blocks, layout_bboxes)
  91. aligned_layout_bboxes = __valign_lines(blocks, layout_bboxes)
  92. return aligned_layout_bboxes
  93. def __pre_proc_zh_blocks(blocks, layout_bboxes):
  94. """
  95. 对中文文本进行分段预处理
  96. """
  97. pass
  98. def __pre_proc_en_blocks(blocks, layout_bboxes):
  99. """
  100. 对英文文本进行分段预处理
  101. """
  102. pass
  103. def __group_line_by_layout(blocks, layout_bboxes, lang="en"):
  104. """
  105. 每个layout内的行进行聚合
  106. """
  107. # 因为只是一个block一行目前, 一个block就是一个段落
  108. lines_group = []
  109. for lyout in layout_bboxes:
  110. lines = [line for block in blocks if _is_in(block['bbox'], lyout['layout_bbox']) for line in block['lines']]
  111. lines_group.append(lines)
  112. return lines_group
  113. def __split_para_in_layoutbox(lines_group, layout_bboxes, lang="en", char_avg_len=10):
  114. """
  115. lines_group 进行行分段——layout内部进行分段。
  116. 1. 先计算每个group的左右边界。
  117. 2. 然后根据行末尾特征进行分段。
  118. 末尾特征:以句号等结束符结尾。并且距离右侧边界有一定距离。
  119. """
  120. paras = []
  121. right_tail_distance = 1.5 * char_avg_len
  122. for lines in lines_group:
  123. if len(lines)==0:
  124. continue
  125. layout_right = max([line['bbox'][2] for line in lines])
  126. para = [] # 元素是line
  127. for line in lines:
  128. line_text = ''.join([__get_span_text(span) for span in line['spans']])
  129. #logger.info(line_text)
  130. last_span_type = line['spans'][-1]['type']
  131. if last_span_type in [TEXT, INLINE_EQUATION]:
  132. last_char = line['spans'][-1]['content'][-1]
  133. if last_char in LINE_STOP_FLAG or line['bbox'][2] < layout_right - right_tail_distance:
  134. para.append(line)
  135. paras.append(para)
  136. # para_text = ''.join([span['content'] for line in para for span in line['spans']])
  137. # logger.info(para_text)
  138. para = []
  139. else:
  140. para.append(line)
  141. else: # 其他,图片、表格、行间公式,各自占一段
  142. if len(para)>0: # 先把之前的段落加入到结果中
  143. paras.append(para)
  144. para = []
  145. paras.append([line]) # 再把当前行加入到结果中。当前行为行间公式、图、表等。
  146. para = []
  147. # para_text = ''.join([get_span_text(span) for line in para for span in line['spans']])
  148. # logger.info(para_text)
  149. if len(para)>0:
  150. paras.append(para)
  151. # para_text = ''.join([get_span_text(span) for line in para for span in line['spans']])
  152. # logger.info(para_text)
  153. para = []
  154. return paras
  155. def __find_layout_bbox_by_line(line_bbox, layout_bboxes):
  156. """
  157. 根据line找到所在的layout
  158. """
  159. for layout in layout_bboxes:
  160. if _is_in(line_bbox, layout):
  161. return layout
  162. return None
  163. def __connect_para_inter_layoutbox(layout_paras, new_layout_bbox, lang="en"):
  164. """
  165. layout之间进行分段。
  166. 主要是计算前一个layOut的最后一行和后一个layout的第一行是否可以连接。
  167. 连接的条件需要同时满足:
  168. 1. 上一个layout的最后一行沾满整个行。并且没有结尾符号。
  169. 2. 下一行开头不留空白。
  170. """
  171. connected_layout_paras = []
  172. for i, para in enumerate(layout_paras):
  173. if i==0:
  174. connected_layout_paras.append(para)
  175. continue
  176. pre_last_line = layout_paras[i-1][-1]
  177. next_first_line = layout_paras[i][0]
  178. pre_last_line_text = ''.join([__get_span_text(span) for span in pre_last_line['spans']])
  179. pre_last_line_type = pre_last_line['spans'][-1]['type']
  180. next_first_line_text = ''.join([__get_span_text(span) for span in next_first_line['spans']])
  181. next_first_line_type = next_first_line['spans'][0]['type']
  182. if pre_last_line_type not in [TEXT, INLINE_EQUATION] or next_first_line_type not in [TEXT, INLINE_EQUATION]: # TODO,真的要做好,要考虑跨table, image, 行间的情况
  183. connected_layout_paras.append(para)
  184. continue
  185. pre_x2_max = __find_layout_bbox_by_line(pre_last_line['bbox'], new_layout_bbox)[2]
  186. next_x0_min = __find_layout_bbox_by_line(next_first_line['bbox'], new_layout_bbox)[0]
  187. pre_last_line_text = pre_last_line_text.strip()
  188. next_first_line_text = next_first_line_text.strip()
  189. if pre_last_line['bbox'][2] == pre_x2_max and pre_last_line_text[-1] not in LINE_STOP_FLAG and next_first_line['bbox'][0]==next_x0_min: # 前面一行沾满了整个行,并且没有结尾符号.下一行没有空白开头。
  190. """连接段落条件成立,将前一个layout的段落和后一个layout的段落连接。"""
  191. connected_layout_paras[-1].extend(para)
  192. else:
  193. """连接段落条件不成立,将前一个layout的段落加入到结果中。"""
  194. connected_layout_paras.append(para)
  195. return connected_layout_paras
  196. def __connect_para_inter_page(pre_page_paras, next_page_paras, pre_page_layout_bbox, next_page_layout_bbox, lang):
  197. """
  198. 连接起来相邻两个页面的段落——前一个页面最后一个段落和后一个页面的第一个段落。
  199. 是否可以连接的条件:
  200. 1. 前一个页面的最后一个段落最后一行沾满整个行。并且没有结尾符号。
  201. 2. 后一个页面的第一个段落第一行没有空白开头。
  202. """
  203. pre_last_para = pre_page_paras[-1]
  204. next_first_para = next_page_paras[0]
  205. pre_last_line = pre_last_para[-1]
  206. next_first_line = next_first_para[0]
  207. pre_last_line_text = ''.join([__get_span_text(span) for span in pre_last_line['spans']])
  208. pre_last_line_type = pre_last_line['spans'][-1]['type']
  209. next_first_line_text = ''.join([__get_span_text(span) for span in next_first_line['spans']])
  210. next_first_line_type = next_first_line['spans'][0]['type']
  211. if pre_last_line_type not in [TEXT, INLINE_EQUATION] or next_first_line_type not in [TEXT, INLINE_EQUATION]: # TODO,真的要做好,要考虑跨table, image, 行间的情况
  212. # 不是文本,不连接
  213. return False
  214. pre_x2_max = __find_layout_bbox_by_line(pre_last_line['bbox'], pre_page_layout_bbox)[2]
  215. next_x0_min = __find_layout_bbox_by_line(next_first_line['bbox'], next_page_layout_bbox)[0]
  216. pre_last_line_text = pre_last_line_text.strip()
  217. next_first_line_text = next_first_line_text.strip()
  218. if pre_last_line['bbox'][2] == pre_x2_max and pre_last_line_text[-1] not in LINE_STOP_FLAG and next_first_line['bbox'][0]==next_x0_min: # 前面一行沾满了整个行,并且没有结尾符号.下一行没有空白开头。
  219. """连接段落条件成立,将前一个layout的段落和后一个layout的段落连接。"""
  220. pre_page_paras[-1].extend(next_first_para)
  221. next_page_paras.pop(0) # 删除后一个页面的第一个段落, 因为他已经被合并到前一个页面的最后一个段落了。
  222. return True
  223. else:
  224. return False
  225. def __do_split(blocks, layout_bboxes, new_layout_bbox, lang="en"):
  226. """
  227. 根据line和layout情况进行分段
  228. 先实现一个根据行末尾特征分段的简单方法。
  229. """
  230. """
  231. 算法思路:
  232. 1. 扫描layout里每一行,找出来行尾距离layout有边界有一定距离的行。
  233. 2. 从上述行中找到末尾是句号等可作为断行标志的行。
  234. 3. 参照上述行尾特征进行分段。
  235. 4. 图、表,目前独占一行,不考虑分段。
  236. """
  237. lines_group = __group_line_by_layout(blocks, layout_bboxes, lang) # block内分段
  238. layout_paras = __split_para_in_layoutbox(lines_group, layout_bboxes, lang) # layout内分段
  239. connected_layout_paras = __connect_para_inter_layoutbox(layout_paras, new_layout_bbox, lang) # layout间链接段落
  240. return connected_layout_paras
  241. def para_split(pdf_info_dict, lang="en"):
  242. """
  243. 根据line和layout情况进行分段
  244. """
  245. new_layout_of_pages = [] # 数组的数组,每个元素是一个页面的layoutS
  246. for _, page in pdf_info_dict.items():
  247. blocks = page['preproc_blocks']
  248. layout_bboxes = page['layout_bboxes']
  249. new_layout_bbox = __common_pre_proc(blocks, layout_bboxes)
  250. new_layout_of_pages.append(new_layout_bbox)
  251. splited_blocks = __do_split(blocks, layout_bboxes, new_layout_bbox, lang)
  252. page['para_blocks'] = splited_blocks
  253. """连接页面与页面之间的可能合并的段落"""
  254. pdf_infos = list(pdf_info_dict.values())
  255. for i, page in enumerate(pdf_info_dict.values()):
  256. if i==0:
  257. continue
  258. pre_page_paras = pdf_infos[i-1]['para_blocks']
  259. next_page_paras = pdf_infos[i]['para_blocks']
  260. pre_page_layout_bbox = new_layout_of_pages[i-1]
  261. next_page_layout_bbox = new_layout_of_pages[i]
  262. is_conn= __connect_para_inter_page(pre_page_paras, next_page_paras, pre_page_layout_bbox, next_page_layout_bbox, lang)
  263. if is_conn:
  264. logger.info(f"连接了第{i-1}页和第{i}页的段落")