para_split.py 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476
  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_or_part_overlap
  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 __detect_line_align_direction(line, new_layout_bboxes):
  31. """
  32. 探测line是左对齐,还是右对齐,还是居中。
  33. """
  34. lbox = line['bbox']
  35. x0, x1 = lbox[0], lbox[2]
  36. layout_x0, layout_x1 = new_layout_bboxes[0], new_layout_bboxes[2]
  37. if x0 <= layout_x0 and x1 < layout_x1:
  38. return "left"
  39. elif x0 > layout_x0 and x1 >= layout_x1:
  40. return "right"
  41. else:
  42. return "center"
  43. def __detect_line_group_align_direction(lines, new_layout_bboxes):
  44. """
  45. 首先把lines按照行距离分成几部分。针对每一部分分别探测。
  46. 最后返回[(dir, lines), (dir, lines), ...]
  47. """
  48. pass
  49. def __detect_list_lines(lines, new_layout_bboxes, lang='en'):
  50. """
  51. 探测是否包含了列表,并且把列表的行分开.
  52. 这样的段落特点是,顶格字母大写/数字,紧跟着几行缩进的。缩进的行首字母含小写的。
  53. """
  54. def find_repeating_patterns(lst):
  55. indices = []
  56. ones_indices = []
  57. i = 0
  58. while i < len(lst) - 1: # 确保余下元素至少有2个
  59. if lst[i] == 1 and lst[i+1] in [2, 3]: # 额外检查以防止连续出现的1
  60. start = i
  61. ones_in_this_interval = [i]
  62. i += 1
  63. while i < len(lst) and lst[i] in [2, 3]:
  64. i += 1
  65. # 验证下一个序列是否符合条件
  66. if i < len(lst) - 1 and lst[i] == 1 and lst[i+1] in [2, 3] and lst[i-1] in [2, 3]:
  67. while i < len(lst) and lst[i] in [1, 2, 3]:
  68. if lst[i] == 1:
  69. ones_in_this_interval.append(i)
  70. i += 1
  71. indices.append((start, i - 1))
  72. ones_indices.append(ones_in_this_interval)
  73. else:
  74. i += 1
  75. else:
  76. i += 1
  77. return indices, ones_indices
  78. """===================="""
  79. def split_indices(slen, index_array):
  80. result = []
  81. last_end = 0
  82. for start, end in sorted(index_array):
  83. if start > last_end:
  84. # 前一个区间结束到下一个区间开始之间的部分标记为"text"
  85. result.append(('text', last_end, start - 1))
  86. # 区间内标记为"list"
  87. result.append(('list', start, end))
  88. last_end = end + 1
  89. if last_end < slen:
  90. # 如果最后一个区间结束后还有剩余的字符串,将其标记为"text"
  91. result.append(('text', last_end, slen - 1))
  92. return result
  93. """===================="""
  94. if lang!='en':
  95. return lines, None
  96. else:
  97. total_lines = len(lines)
  98. line_fea_encode = []
  99. """
  100. 对每一行进行特征编码,编码规则如下:
  101. 1. 如果行顶格,且大写字母开头或者数字开头,编码为1
  102. 2. 如果顶格,其他非大写开头编码为4
  103. 3. 如果非顶格,首字符大写,编码为2
  104. 4. 如果非顶格,首字符非大写编码为3
  105. """
  106. for l in lines:
  107. first_char = __get_span_text(l['spans'][0])[0]
  108. layout_left = __find_layout_bbox_by_line(l['bbox'], new_layout_bboxes)[0]
  109. if l['bbox'][0] == layout_left:
  110. if first_char.isupper() or first_char.isdigit():
  111. line_fea_encode.append(1)
  112. else:
  113. line_fea_encode.append(4)
  114. else:
  115. if first_char.isupper():
  116. line_fea_encode.append(2)
  117. else:
  118. line_fea_encode.append(3)
  119. # 然后根据编码进行分段, 选出来 1,2,3连续出现至少2次的行,认为是列表。
  120. list_indice, list_start_idx = find_repeating_patterns(line_fea_encode)
  121. if len(list_indice)>0:
  122. logger.info(f"发现了列表,列表行数:{list_indice}, {list_start_idx}")
  123. # TODO check一下这个特列表里缩进的行左侧是不是对齐的。
  124. segments = []
  125. for start, end in list_indice:
  126. for i in range(start, end+1):
  127. if i>0:
  128. if line_fea_encode[i] == 4:
  129. logger.info(f"列表行的第{i}行不是顶格的")
  130. break
  131. else:
  132. logger.info(f"列表行的第{start}到第{end}行是列表")
  133. return split_indices(total_lines, list_indice), list_start_idx
  134. def __valign_lines(blocks, layout_bboxes):
  135. """
  136. 在一个layoutbox内对齐行的左侧和右侧。
  137. 扫描行的左侧和右侧,如果x0, x1差距不超过一个阈值,就强行对齐到所处layout的左右两侧(和layout有一段距离)。
  138. 3是个经验值,TODO,计算得来,可以设置为1.5个正文字符。
  139. """
  140. min_distance = 3
  141. min_sample = 2
  142. new_layout_bboxes = []
  143. for layout_box in layout_bboxes:
  144. blocks_in_layoutbox = [b for b in blocks if _is_in_or_part_overlap(b['bbox'], layout_box['layout_bbox'])]
  145. if len(blocks_in_layoutbox)==0:
  146. continue
  147. x0_lst = np.array([[line['bbox'][0], 0] for block in blocks_in_layoutbox for line in block['lines']])
  148. x1_lst = np.array([[line['bbox'][2], 0] for block in blocks_in_layoutbox for line in block['lines']])
  149. x0_clusters = DBSCAN(eps=min_distance, min_samples=min_sample).fit(x0_lst)
  150. x1_clusters = DBSCAN(eps=min_distance, min_samples=min_sample).fit(x1_lst)
  151. x0_uniq_label = np.unique(x0_clusters.labels_)
  152. x1_uniq_label = np.unique(x1_clusters.labels_)
  153. x0_2_new_val = {} # 存储旧值对应的新值映射
  154. x1_2_new_val = {}
  155. for label in x0_uniq_label:
  156. if label==-1:
  157. continue
  158. x0_index_of_label = np.where(x0_clusters.labels_==label)
  159. x0_raw_val = x0_lst[x0_index_of_label][:,0]
  160. x0_new_val = np.min(x0_lst[x0_index_of_label][:,0])
  161. x0_2_new_val.update({idx: x0_new_val for idx in x0_raw_val})
  162. for label in x1_uniq_label:
  163. if label==-1:
  164. continue
  165. x1_index_of_label = np.where(x1_clusters.labels_==label)
  166. x1_raw_val = x1_lst[x1_index_of_label][:,0]
  167. x1_new_val = np.max(x1_lst[x1_index_of_label][:,0])
  168. x1_2_new_val.update({idx: x1_new_val for idx in x1_raw_val})
  169. for block in blocks_in_layoutbox:
  170. for line in block['lines']:
  171. x0, x1 = line['bbox'][0], line['bbox'][2]
  172. if x0 in x0_2_new_val:
  173. line['bbox'][0] = int(x0_2_new_val[x0])
  174. if x1 in x1_2_new_val:
  175. line['bbox'][2] = int(x1_2_new_val[x1])
  176. # 其余对不齐的保持不动
  177. # 由于修改了block里的line长度,现在需要重新计算block的bbox
  178. for block in blocks_in_layoutbox:
  179. block['bbox'] = [min([line['bbox'][0] for line in block['lines']]),
  180. min([line['bbox'][1] for line in block['lines']]),
  181. max([line['bbox'][2] for line in block['lines']]),
  182. max([line['bbox'][3] for line in block['lines']])]
  183. """新计算layout的bbox,因为block的bbox变了。"""
  184. layout_x0 = min([block['bbox'][0] for block in blocks_in_layoutbox])
  185. layout_y0 = min([block['bbox'][1] for block in blocks_in_layoutbox])
  186. layout_x1 = max([block['bbox'][2] for block in blocks_in_layoutbox])
  187. layout_y1 = max([block['bbox'][3] for block in blocks_in_layoutbox])
  188. new_layout_bboxes.append([layout_x0, layout_y0, layout_x1, layout_y1])
  189. return new_layout_bboxes
  190. def __common_pre_proc(blocks, layout_bboxes):
  191. """
  192. 不分语言的,对文本进行预处理
  193. """
  194. #__add_line_period(blocks, layout_bboxes)
  195. aligned_layout_bboxes = __valign_lines(blocks, layout_bboxes)
  196. return aligned_layout_bboxes
  197. def __pre_proc_zh_blocks(blocks, layout_bboxes):
  198. """
  199. 对中文文本进行分段预处理
  200. """
  201. pass
  202. def __pre_proc_en_blocks(blocks, layout_bboxes):
  203. """
  204. 对英文文本进行分段预处理
  205. """
  206. pass
  207. def __group_line_by_layout(blocks, layout_bboxes, lang="en"):
  208. """
  209. 每个layout内的行进行聚合
  210. """
  211. # 因为只是一个block一行目前, 一个block就是一个段落
  212. lines_group = []
  213. for lyout in layout_bboxes:
  214. lines = [line for block in blocks if _is_in_or_part_overlap(block['bbox'], lyout['layout_bbox']) for line in block['lines']]
  215. lines_group.append(lines)
  216. return lines_group
  217. def __split_para_in_layoutbox(lines_group, new_layout_bbox, lang="en", char_avg_len=10):
  218. """
  219. lines_group 进行行分段——layout内部进行分段。lines_group内每个元素是一个Layoutbox内的所有行。
  220. 1. 先计算每个group的左右边界。
  221. 2. 然后根据行末尾特征进行分段。
  222. 末尾特征:以句号等结束符结尾。并且距离右侧边界有一定距离。
  223. 且下一行开头不留空白。
  224. """
  225. line_group_end_with_list = [] # 这个layout最后是不是列表,用于跨layout列表合并
  226. paras = []
  227. right_tail_distance = 1.5 * char_avg_len
  228. for lines in lines_group:
  229. total_lines = len(lines)
  230. if total_lines<=1: # 0行无需处理。1行无法分段。
  231. continue
  232. """在进入到真正的分段之前,要对文字块从统计维度进行对齐方式的探测,
  233. 对齐方式分为以下:
  234. 1. 左对齐的文本块(特点是左侧顶格,或者左侧不顶格但是右侧顶格的行数大于非顶格的行数,顶格的首字母有大写也有小写)
  235. 1) 右侧对齐的行,单独成一段
  236. 2) 中间对齐的行,按照字体/行高聚合成一段
  237. 2. 左对齐的列表块(其特点是左侧顶格的行数小于等于非顶格的行数,非定格首字母会有小写,顶格90%是大写。并且左侧顶格行数大于1,大于1是为了这种模式连续出现才能称之为列表)
  238. 这样的文本块,顶格的为一个段落开头,紧随其后非顶格的行属于这个段落。
  239. """
  240. text_segments, list_start_line = __detect_list_lines(lines, new_layout_bbox, lang)
  241. """根据list_range,把lines分成几个部分
  242. """
  243. layout_right = __find_layout_bbox_by_line(lines[0]['bbox'], new_layout_bbox)[2]
  244. layout_left = __find_layout_bbox_by_line(lines[0]['bbox'], new_layout_bbox)[0]
  245. para = [] # 元素是line
  246. is_lines_end_with_list = False
  247. for content_type, start, end in text_segments:
  248. if content_type == 'list':
  249. for i, line in enumerate(lines[start:end+1]):
  250. line_x0 = line['bbox'][0]
  251. if line_x0 == layout_left: # 列表开头
  252. if len(para)>0:
  253. paras.append(para)
  254. para = []
  255. para.append(line)
  256. else:
  257. para.append(line)
  258. if len(para)>0:
  259. paras.append(para)
  260. para = []
  261. is_lines_end_with_list = True
  262. else:
  263. for i, line in enumerate(lines[start:end+1]):
  264. # 如果i有下一行,那么就要根据下一行位置综合判断是否要分段。如果i之后没有行,那么只需要判断一下行结尾特征。
  265. cur_line_type = line['spans'][-1]['type']
  266. next_line = lines[i+1] if i<total_lines-1 else None
  267. if cur_line_type in [TEXT, INLINE_EQUATION]:
  268. if line['bbox'][2] < layout_right - right_tail_distance:
  269. para.append(line)
  270. paras.append(para)
  271. para = []
  272. elif line['bbox'][2] >= layout_right - right_tail_distance and next_line and next_line['bbox'][0] == layout_left: # 现在这行到了行尾沾满,下一行存在且顶格。
  273. para.append(line)
  274. else:
  275. para.append(line)
  276. paras.append(para)
  277. para = []
  278. else: # 其他,图片、表格、行间公式,各自占一段
  279. if len(para)>0: # 先把之前的段落加入到结果中
  280. paras.append(para)
  281. para = []
  282. paras.append([line]) # 再把当前行加入到结果中。当前行为行间公式、图、表等。
  283. para = []
  284. if len(para)>0:
  285. paras.append(para)
  286. para = []
  287. is_lines_end_with_list = False
  288. line_group_end_with_list.append(is_lines_end_with_list)
  289. return paras, line_group_end_with_list
  290. def __find_layout_bbox_by_line(line_bbox, layout_bboxes):
  291. """
  292. 根据line找到所在的layout
  293. """
  294. for layout in layout_bboxes:
  295. if _is_in_or_part_overlap(line_bbox, layout):
  296. return layout
  297. return None
  298. def __connect_para_inter_layoutbox(layout_paras, new_layout_bbox, line_group_end_with_list, lang="en"):
  299. """
  300. layout之间进行分段。
  301. 主要是计算前一个layOut的最后一行和后一个layout的第一行是否可以连接。
  302. 连接的条件需要同时满足:
  303. 1. 上一个layout的最后一行沾满整个行。并且没有结尾符号。
  304. 2. 下一行开头不留空白。
  305. """
  306. connected_layout_paras = []
  307. for i, para in enumerate(layout_paras):
  308. if i==0:
  309. connected_layout_paras.append(para)
  310. continue
  311. pre_last_line = layout_paras[i-1][-1]
  312. next_first_line = layout_paras[i][0]
  313. pre_last_line_text = ''.join([__get_span_text(span) for span in pre_last_line['spans']])
  314. pre_last_line_type = pre_last_line['spans'][-1]['type']
  315. next_first_line_text = ''.join([__get_span_text(span) for span in next_first_line['spans']])
  316. next_first_line_type = next_first_line['spans'][0]['type']
  317. if pre_last_line_type not in [TEXT, INLINE_EQUATION] or next_first_line_type not in [TEXT, INLINE_EQUATION]: # TODO,真的要做好,要考虑跨table, image, 行间的情况
  318. connected_layout_paras.append(para)
  319. continue
  320. pre_x2_max = __find_layout_bbox_by_line(pre_last_line['bbox'], new_layout_bbox)[2]
  321. next_x0_min = __find_layout_bbox_by_line(next_first_line['bbox'], new_layout_bbox)[0]
  322. pre_last_line_text = pre_last_line_text.strip()
  323. next_first_line_text = next_first_line_text.strip()
  324. 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: # 前面一行沾满了整个行,并且没有结尾符号.下一行没有空白开头。
  325. """连接段落条件成立,将前一个layout的段落和后一个layout的段落连接。"""
  326. connected_layout_paras[-1].extend(para)
  327. else:
  328. """连接段落条件不成立,将前一个layout的段落加入到结果中。"""
  329. connected_layout_paras.append(para)
  330. return connected_layout_paras
  331. def __connect_para_inter_page(pre_page_paras, next_page_paras, pre_page_layout_bbox, next_page_layout_bbox, lang):
  332. """
  333. 连接起来相邻两个页面的段落——前一个页面最后一个段落和后一个页面的第一个段落。
  334. 是否可以连接的条件:
  335. 1. 前一个页面的最后一个段落最后一行沾满整个行。并且没有结尾符号。
  336. 2. 后一个页面的第一个段落第一行没有空白开头。
  337. """
  338. # 有的页面可能压根没有文字
  339. if len(pre_page_paras)==0 or len(next_page_paras)==0:
  340. return False
  341. pre_last_para = pre_page_paras[-1]
  342. next_first_para = next_page_paras[0]
  343. pre_last_line = pre_last_para[-1]
  344. next_first_line = next_first_para[0]
  345. pre_last_line_text = ''.join([__get_span_text(span) for span in pre_last_line['spans']])
  346. pre_last_line_type = pre_last_line['spans'][-1]['type']
  347. next_first_line_text = ''.join([__get_span_text(span) for span in next_first_line['spans']])
  348. next_first_line_type = next_first_line['spans'][0]['type']
  349. if pre_last_line_type not in [TEXT, INLINE_EQUATION] or next_first_line_type not in [TEXT, INLINE_EQUATION]: # TODO,真的要做好,要考虑跨table, image, 行间的情况
  350. # 不是文本,不连接
  351. return False
  352. pre_x2_max = __find_layout_bbox_by_line(pre_last_line['bbox'], pre_page_layout_bbox)[2]
  353. next_x0_min = __find_layout_bbox_by_line(next_first_line['bbox'], next_page_layout_bbox)[0]
  354. pre_last_line_text = pre_last_line_text.strip()
  355. next_first_line_text = next_first_line_text.strip()
  356. 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: # 前面一行沾满了整个行,并且没有结尾符号.下一行没有空白开头。
  357. """连接段落条件成立,将前一个layout的段落和后一个layout的段落连接。"""
  358. pre_page_paras[-1].extend(next_first_para)
  359. next_page_paras.pop(0) # 删除后一个页面的第一个段落, 因为他已经被合并到前一个页面的最后一个段落了。
  360. return True
  361. else:
  362. return False
  363. def __do_split(blocks, layout_bboxes, new_layout_bbox, lang="en"):
  364. """
  365. 根据line和layout情况进行分段
  366. 先实现一个根据行末尾特征分段的简单方法。
  367. """
  368. """
  369. 算法思路:
  370. 1. 扫描layout里每一行,找出来行尾距离layout有边界有一定距离的行。
  371. 2. 从上述行中找到末尾是句号等可作为断行标志的行。
  372. 3. 参照上述行尾特征进行分段。
  373. 4. 图、表,目前独占一行,不考虑分段。
  374. """
  375. lines_group = __group_line_by_layout(blocks, layout_bboxes, lang) # block内分段
  376. layout_paras, line_group_end_with_list = __split_para_in_layoutbox(lines_group, new_layout_bbox, lang) # layout内分段
  377. connected_layout_paras = __connect_para_inter_layoutbox(layout_paras, new_layout_bbox, line_group_end_with_list, lang) # layout间链接段落
  378. return connected_layout_paras
  379. def para_split(pdf_info_dict, lang="en"):
  380. """
  381. 根据line和layout情况进行分段
  382. """
  383. new_layout_of_pages = [] # 数组的数组,每个元素是一个页面的layoutS
  384. for _, page in pdf_info_dict.items():
  385. blocks = page['preproc_blocks']
  386. layout_bboxes = page['layout_bboxes']
  387. new_layout_bbox = __common_pre_proc(blocks, layout_bboxes)
  388. new_layout_of_pages.append(new_layout_bbox)
  389. splited_blocks = __do_split(blocks, layout_bboxes, new_layout_bbox, lang)
  390. page['para_blocks'] = splited_blocks
  391. """连接页面与页面之间的可能合并的段落"""
  392. pdf_infos = list(pdf_info_dict.values())
  393. for i, page in enumerate(pdf_info_dict.values()):
  394. if i==0:
  395. continue
  396. pre_page_paras = pdf_infos[i-1]['para_blocks']
  397. next_page_paras = pdf_infos[i]['para_blocks']
  398. pre_page_layout_bbox = new_layout_of_pages[i-1]
  399. next_page_layout_bbox = new_layout_of_pages[i]
  400. is_conn= __connect_para_inter_page(pre_page_paras, next_page_paras, pre_page_layout_bbox, next_page_layout_bbox, lang)
  401. if is_conn:
  402. logger.info(f"连接了第{i-1}页和第{i}页的段落")