para_split_v2.py 36 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790
  1. from sklearn.cluster import DBSCAN
  2. import numpy as np
  3. from loguru import logger
  4. import re
  5. from magic_pdf.libs.boxbase import _is_in_or_part_overlap_with_area_ratio as is_in_layout
  6. from magic_pdf.libs.ocr_content_type import ContentType, BlockType
  7. from magic_pdf.model.magic_model import MagicModel
  8. from magic_pdf.libs.Constants import *
  9. LINE_STOP_FLAG = ['.', '!', '?', '。', '!', '?', ":", ":", ")", ")", ";"]
  10. INLINE_EQUATION = ContentType.InlineEquation
  11. INTERLINE_EQUATION = ContentType.InterlineEquation
  12. TEXT = ContentType.Text
  13. debug_able = False
  14. def __get_span_text(span):
  15. c = span.get('content', '')
  16. if len(c) == 0:
  17. c = span.get('image_path', '')
  18. return c
  19. def __detect_list_lines(lines, new_layout_bboxes, lang):
  20. global debug_able
  21. """
  22. 探测是否包含了列表,并且把列表的行分开.
  23. 这样的段落特点是,顶格字母大写/数字,紧跟着几行缩进的。缩进的行首字母含小写的。
  24. """
  25. def find_repeating_patterns2(lst):
  26. indices = []
  27. ones_indices = []
  28. i = 0
  29. while i < len(lst): # Loop through the entire list
  30. if lst[i] == 1: # If we encounter a '1', we might be at the start of a pattern
  31. start = i
  32. ones_in_this_interval = [i]
  33. i += 1
  34. # Traverse elements that are 1, 2 or 3, until we encounter something else
  35. while i < len(lst) and lst[i] in [1, 2, 3]:
  36. if lst[i] == 1:
  37. ones_in_this_interval.append(i)
  38. i += 1
  39. if len(ones_in_this_interval) > 1 or (
  40. start < len(lst) - 1 and ones_in_this_interval and lst[start + 1] in [2, 3]):
  41. indices.append((start, i - 1))
  42. ones_indices.append(ones_in_this_interval)
  43. else:
  44. i += 1
  45. return indices, ones_indices
  46. def find_repeating_patterns(lst):
  47. indices = []
  48. ones_indices = []
  49. i = 0
  50. while i < len(lst) - 1: # 确保余下元素至少有2个
  51. if lst[i] == 1 and lst[i + 1] in [2, 3]: # 额外检查以防止连续出现的1
  52. start = i
  53. ones_in_this_interval = [i]
  54. i += 1
  55. while i < len(lst) and lst[i] in [2, 3]:
  56. i += 1
  57. # 验证下一个序列是否符合条件
  58. if i < len(lst) - 1 and lst[i] == 1 and lst[i + 1] in [2, 3] and lst[i - 1] in [2, 3]:
  59. while i < len(lst) and lst[i] in [1, 2, 3]:
  60. if lst[i] == 1:
  61. ones_in_this_interval.append(i)
  62. i += 1
  63. indices.append((start, i - 1))
  64. ones_indices.append(ones_in_this_interval)
  65. else:
  66. i += 1
  67. else:
  68. i += 1
  69. return indices, ones_indices
  70. """===================="""
  71. def split_indices(slen, index_array):
  72. result = []
  73. last_end = 0
  74. for start, end in sorted(index_array):
  75. if start > last_end:
  76. # 前一个区间结束到下一个区间开始之间的部分标记为"text"
  77. result.append(('text', last_end, start - 1))
  78. # 区间内标记为"list"
  79. result.append(('list', start, end))
  80. last_end = end + 1
  81. if last_end < slen:
  82. # 如果最后一个区间结束后还有剩余的字符串,将其标记为"text"
  83. result.append(('text', last_end, slen - 1))
  84. return result
  85. """===================="""
  86. if lang != 'en':
  87. return lines, None
  88. total_lines = len(lines)
  89. line_fea_encode = []
  90. """
  91. 对每一行进行特征编码,编码规则如下:
  92. 1. 如果行顶格,且大写字母开头或者数字开头,编码为1
  93. 2. 如果顶格,其他非大写开头编码为4
  94. 3. 如果非顶格,首字符大写,编码为2
  95. 4. 如果非顶格,首字符非大写编码为3
  96. """
  97. if len(lines) > 0:
  98. x_map_tag_dict, min_x_tag = cluster_line_x(lines)
  99. for l in lines:
  100. span_text = __get_span_text(l['spans'][0])
  101. if not span_text:
  102. line_fea_encode.append(0)
  103. continue
  104. first_char = span_text[0]
  105. layout = __find_layout_bbox_by_line(l['bbox'], new_layout_bboxes)
  106. if not layout:
  107. line_fea_encode.append(0)
  108. else:
  109. #
  110. if x_map_tag_dict[round(l['bbox'][0])] == min_x_tag:
  111. # if first_char.isupper() or first_char.isdigit() or not first_char.isalnum():
  112. if not first_char.isalnum() or if_match_reference_list(span_text):
  113. line_fea_encode.append(1)
  114. else:
  115. line_fea_encode.append(4)
  116. else:
  117. if first_char.isupper():
  118. line_fea_encode.append(2)
  119. else:
  120. line_fea_encode.append(3)
  121. # 然后根据编码进行分段, 选出来 1,2,3连续出现至少2次的行,认为是列表。
  122. list_indice, list_start_idx = find_repeating_patterns2(line_fea_encode)
  123. if len(list_indice) > 0:
  124. if debug_able:
  125. logger.info(f"发现了列表,列表行数:{list_indice}, {list_start_idx}")
  126. # TODO check一下这个特列表里缩进的行左侧是不是对齐的。
  127. segments = []
  128. for start, end in list_indice:
  129. for i in range(start, end + 1):
  130. if i > 0:
  131. if line_fea_encode[i] == 4:
  132. if debug_able:
  133. logger.info(f"列表行的第{i}行不是顶格的")
  134. break
  135. else:
  136. if debug_able:
  137. logger.info(f"列表行的第{start}到第{end}行是列表")
  138. return split_indices(total_lines, list_indice), list_start_idx
  139. def cluster_line_x(lines: list) -> dict:
  140. """
  141. 对一个block内所有lines的bbox的x0聚类
  142. """
  143. min_distance = 5
  144. min_sample = 1
  145. x0_lst = np.array([[round(line['bbox'][0]), 0] for line in lines])
  146. x0_clusters = DBSCAN(eps=min_distance, min_samples=min_sample).fit(x0_lst)
  147. x0_uniq_label = np.unique(x0_clusters.labels_)
  148. #x1_lst = np.array([[line['bbox'][2], 0] for line in lines])
  149. x0_2_new_val = {} # 存储旧值对应的新值映射
  150. min_x0 = round(lines[0]["bbox"][0])
  151. for label in x0_uniq_label:
  152. if label == -1:
  153. continue
  154. x0_index_of_label = np.where(x0_clusters.labels_ == label)
  155. x0_raw_val = x0_lst[x0_index_of_label][:, 0]
  156. x0_new_val = np.min(x0_lst[x0_index_of_label][:, 0])
  157. x0_2_new_val.update({round(raw_val): round(x0_new_val) for raw_val in x0_raw_val})
  158. if x0_new_val < min_x0:
  159. min_x0 = x0_new_val
  160. return x0_2_new_val, min_x0
  161. def if_match_reference_list(text: str) -> bool:
  162. pattern = re.compile(r'^\d+\..*')
  163. if pattern.match(text):
  164. return True
  165. else:
  166. return False
  167. def __valign_lines(blocks, layout_bboxes):
  168. """
  169. 在一个layoutbox内对齐行的左侧和右侧。
  170. 扫描行的左侧和右侧,如果x0, x1差距不超过一个阈值,就强行对齐到所处layout的左右两侧(和layout有一段距离)。
  171. 3是个经验值,TODO,计算得来,可以设置为1.5个正文字符。
  172. """
  173. min_distance = 3
  174. min_sample = 2
  175. new_layout_bboxes = []
  176. for layout_box in layout_bboxes:
  177. blocks_in_layoutbox = [b for b in blocks if
  178. b["type"] == BlockType.Text and is_in_layout(b['bbox'], layout_box['layout_bbox'])]
  179. if len(blocks_in_layoutbox) == 0 or len(blocks_in_layoutbox[0]["lines"]) == 0:
  180. new_layout_bboxes.append(layout_box['layout_bbox'])
  181. continue
  182. x0_lst = np.array([[line['bbox'][0], 0] for block in blocks_in_layoutbox for line in block['lines']])
  183. x1_lst = np.array([[line['bbox'][2], 0] for block in blocks_in_layoutbox for line in block['lines']])
  184. x0_clusters = DBSCAN(eps=min_distance, min_samples=min_sample).fit(x0_lst)
  185. x1_clusters = DBSCAN(eps=min_distance, min_samples=min_sample).fit(x1_lst)
  186. x0_uniq_label = np.unique(x0_clusters.labels_)
  187. x1_uniq_label = np.unique(x1_clusters.labels_)
  188. x0_2_new_val = {} # 存储旧值对应的新值映射
  189. x1_2_new_val = {}
  190. for label in x0_uniq_label:
  191. if label == -1:
  192. continue
  193. x0_index_of_label = np.where(x0_clusters.labels_ == label)
  194. x0_raw_val = x0_lst[x0_index_of_label][:, 0]
  195. x0_new_val = np.min(x0_lst[x0_index_of_label][:, 0])
  196. x0_2_new_val.update({idx: x0_new_val for idx in x0_raw_val})
  197. for label in x1_uniq_label:
  198. if label == -1:
  199. continue
  200. x1_index_of_label = np.where(x1_clusters.labels_ == label)
  201. x1_raw_val = x1_lst[x1_index_of_label][:, 0]
  202. x1_new_val = np.max(x1_lst[x1_index_of_label][:, 0])
  203. x1_2_new_val.update({idx: x1_new_val for idx in x1_raw_val})
  204. for block in blocks_in_layoutbox:
  205. for line in block['lines']:
  206. x0, x1 = line['bbox'][0], line['bbox'][2]
  207. if x0 in x0_2_new_val:
  208. line['bbox'][0] = int(x0_2_new_val[x0])
  209. if x1 in x1_2_new_val:
  210. line['bbox'][2] = int(x1_2_new_val[x1])
  211. # 其余对不齐的保持不动
  212. # 由于修改了block里的line长度,现在需要重新计算block的bbox
  213. for block in blocks_in_layoutbox:
  214. if len(block["lines"]) > 0:
  215. block['bbox'] = [min([line['bbox'][0] for line in block['lines']]),
  216. min([line['bbox'][1] for line in block['lines']]),
  217. max([line['bbox'][2] for line in block['lines']]),
  218. max([line['bbox'][3] for line in block['lines']])]
  219. """新计算layout的bbox,因为block的bbox变了。"""
  220. layout_x0 = min([block['bbox'][0] for block in blocks_in_layoutbox])
  221. layout_y0 = min([block['bbox'][1] for block in blocks_in_layoutbox])
  222. layout_x1 = max([block['bbox'][2] for block in blocks_in_layoutbox])
  223. layout_y1 = max([block['bbox'][3] for block in blocks_in_layoutbox])
  224. new_layout_bboxes.append([layout_x0, layout_y0, layout_x1, layout_y1])
  225. return new_layout_bboxes
  226. def __align_text_in_layout(blocks, layout_bboxes):
  227. """
  228. 由于ocr出来的line,有时候会在前后有一段空白,这个时候需要对文本进行对齐,超出的部分被layout左右侧截断。
  229. """
  230. for layout in layout_bboxes:
  231. lb = layout['layout_bbox']
  232. blocks_in_layoutbox = [block for block in blocks if
  233. block["type"] == BlockType.Text and is_in_layout(block['bbox'], lb)]
  234. if len(blocks_in_layoutbox) == 0:
  235. continue
  236. for block in blocks_in_layoutbox:
  237. for line in block.get("lines", []):
  238. x0, x1 = line['bbox'][0], line['bbox'][2]
  239. if x0 < lb[0]:
  240. line['bbox'][0] = lb[0]
  241. if x1 > lb[2]:
  242. line['bbox'][2] = lb[2]
  243. def __common_pre_proc(blocks, layout_bboxes):
  244. """
  245. 不分语言的,对文本进行预处理
  246. """
  247. # __add_line_period(blocks, layout_bboxes)
  248. __align_text_in_layout(blocks, layout_bboxes)
  249. aligned_layout_bboxes = __valign_lines(blocks, layout_bboxes)
  250. return aligned_layout_bboxes
  251. def __pre_proc_zh_blocks(blocks, layout_bboxes):
  252. """
  253. 对中文文本进行分段预处理
  254. """
  255. pass
  256. def __pre_proc_en_blocks(blocks, layout_bboxes):
  257. """
  258. 对英文文本进行分段预处理
  259. """
  260. pass
  261. def __group_line_by_layout(blocks, layout_bboxes):
  262. """
  263. 每个layout内的行进行聚合
  264. """
  265. # 因为只是一个block一行目前, 一个block就是一个段落
  266. blocks_group = []
  267. for lyout in layout_bboxes:
  268. blocks_in_layout = [block for block in blocks if is_in_layout(block['bbox'], lyout['layout_bbox'])]
  269. blocks_group.append(blocks_in_layout)
  270. return blocks_group
  271. def __split_para_in_layoutbox(blocks_group, new_layout_bbox, lang="en"):
  272. """
  273. lines_group 进行行分段——layout内部进行分段。lines_group内每个元素是一个Layoutbox内的所有行。
  274. 1. 先计算每个group的左右边界。
  275. 2. 然后根据行末尾特征进行分段。
  276. 末尾特征:以句号等结束符结尾。并且距离右侧边界有一定距离。
  277. 且下一行开头不留空白。
  278. """
  279. list_info = [] # 这个layout最后是不是列表,记录每一个layout里是不是列表开头,列表结尾
  280. for blocks in blocks_group:
  281. is_start_list = None
  282. is_end_list = None
  283. if len(blocks) == 0:
  284. list_info.append([False, False])
  285. continue
  286. if blocks[0]["type"] != BlockType.Text and blocks[-1]["type"] != BlockType.Text:
  287. list_info.append([False, False])
  288. continue
  289. if blocks[0]["type"] != BlockType.Text:
  290. is_start_list = False
  291. if blocks[-1]["type"] != BlockType.Text:
  292. is_end_list = False
  293. lines = [line for block in blocks if
  294. block["type"] == BlockType.Text for line in
  295. block['lines']]
  296. total_lines = len(lines)
  297. if total_lines == 1 or total_lines == 0:
  298. list_info.append([False, False])
  299. continue
  300. """在进入到真正的分段之前,要对文字块从统计维度进行对齐方式的探测,
  301. 对齐方式分为以下:
  302. 1. 左对齐的文本块(特点是左侧顶格,或者左侧不顶格但是右侧顶格的行数大于非顶格的行数,顶格的首字母有大写也有小写)
  303. 1) 右侧对齐的行,单独成一段
  304. 2) 中间对齐的行,按照字体/行高聚合成一段
  305. 2. 左对齐的列表块(其特点是左侧顶格的行数小于等于非顶格的行数,非定格首字母会有小写,顶格90%是大写。并且左侧顶格行数大于1,大于1是为了这种模式连续出现才能称之为列表)
  306. 这样的文本块,顶格的为一个段落开头,紧随其后非顶格的行属于这个段落。
  307. """
  308. text_segments, list_start_line = __detect_list_lines(lines, new_layout_bbox, lang)
  309. """根据list_range,把lines分成几个部分
  310. """
  311. for list_start in list_start_line:
  312. if len(list_start) > 1:
  313. for i in range(0, len(list_start)):
  314. index = list_start[i] - 1
  315. if index >= 0:
  316. if "content" in lines[index]["spans"][-1]:
  317. lines[index]["spans"][-1]["content"] += '\n\n'
  318. layout_list_info = [False, False] # 这个layout最后是不是列表,记录每一个layout里是不是列表开头,列表结尾
  319. for content_type, start, end in text_segments:
  320. if content_type == 'list':
  321. if start == 0 and is_start_list is None:
  322. layout_list_info[0] = True
  323. if end == total_lines - 1 and is_end_list is None:
  324. layout_list_info[1] = True
  325. list_info.append(layout_list_info)
  326. return list_info
  327. def __split_para_lines(lines: list, text_blocks: list) -> list:
  328. text_paras = []
  329. other_paras = []
  330. text_lines = []
  331. for line in lines:
  332. spans_types = [span["type"] for span in line]
  333. if ContentType.Table in spans_types:
  334. other_paras.append([line])
  335. continue
  336. if ContentType.Image in spans_types:
  337. other_paras.append([line])
  338. continue
  339. if ContentType.InterlineEquation in spans_types:
  340. other_paras.append([line])
  341. continue
  342. text_lines.append(line)
  343. for block in text_blocks:
  344. block_bbox = block["bbox"]
  345. para = []
  346. for line in text_lines:
  347. bbox = line["bbox"]
  348. if is_in_layout(bbox, block_bbox):
  349. para.append(line)
  350. if len(para) > 0:
  351. text_paras.append(para)
  352. paras = other_paras.extend(text_paras)
  353. paras_sorted = sorted(paras, key=lambda x: x[0]["bbox"][1])
  354. return paras_sorted
  355. def __connect_list_inter_layout(blocks_group, new_layout_bbox, layout_list_info, page_num, lang):
  356. global debug_able
  357. """
  358. 如果上个layout的最后一个段落是列表,下一个layout的第一个段落也是列表,那么将他们连接起来。 TODO 因为没有区分列表和段落,所以这个方法暂时不实现。
  359. 根据layout_list_info判断是不是列表。,下个layout的第一个段如果不是列表,那么看他们是否有几行都有相同的缩进。
  360. """
  361. if len(blocks_group) == 0 or len(blocks_group) == 0: # 0的时候最后的return 会出错
  362. return blocks_group, [False, False]
  363. for i in range(1, len(blocks_group)):
  364. if len(blocks_group[i]) == 0 or len(blocks_group[i - 1]) == 0:
  365. continue
  366. pre_layout_list_info = layout_list_info[i - 1]
  367. next_layout_list_info = layout_list_info[i]
  368. pre_last_para = blocks_group[i - 1][-1].get("lines", [])
  369. next_paras = blocks_group[i]
  370. next_first_para = next_paras[0]
  371. if pre_layout_list_info[1] and not next_layout_list_info[0] and next_first_para[
  372. "type"] == BlockType.Text: # 前一个是列表结尾,后一个是非列表开头,此时检测是否有相同的缩进
  373. if debug_able:
  374. logger.info(f"连接page {page_num} 内的list")
  375. # 向layout_paras[i] 寻找开头具有相同缩进的连续的行
  376. may_list_lines = []
  377. lines = next_first_para.get("lines", [])
  378. for line in lines:
  379. if line['bbox'][0] > __find_layout_bbox_by_line(line['bbox'], new_layout_bbox)[0]:
  380. may_list_lines.append(line)
  381. else:
  382. break
  383. # 如果这些行的缩进是相等的,那么连到上一个layout的最后一个段落上。
  384. if len(may_list_lines) > 0 and len(set([x['bbox'][0] for x in may_list_lines])) == 1:
  385. pre_last_para.extend(may_list_lines)
  386. next_first_para["lines"] = next_first_para["lines"][len(may_list_lines):]
  387. return blocks_group, [layout_list_info[0][0], layout_list_info[-1][1]] # 同时还返回了这个页面级别的开头、结尾是不是列表的信息
  388. def __connect_list_inter_page(pre_page_paras, next_page_paras, pre_page_layout_bbox, next_page_layout_bbox,
  389. pre_page_list_info, next_page_list_info, page_num, lang):
  390. """
  391. 如果上个layout的最后一个段落是列表,下一个layout的第一个段落也是列表,那么将他们连接起来。 TODO 因为没有区分列表和段落,所以这个方法暂时不实现。
  392. 根据layout_list_info判断是不是列表。,下个layout的第一个段如果不是列表,那么看他们是否有几行都有相同的缩进。
  393. """
  394. if len(pre_page_paras) == 0 or len(next_page_paras) == 0: # 0的时候最后的return 会出错
  395. return False
  396. if len(pre_page_paras[-1]) == 0 or len(next_page_paras[0]) == 0:
  397. return False
  398. if pre_page_paras[-1][-1]["type"] != BlockType.Text or next_page_paras[0][0]["type"] != BlockType.Text:
  399. return False
  400. if pre_page_list_info[1] and not next_page_list_info[0]: # 前一个是列表结尾,后一个是非列表开头,此时检测是否有相同的缩进
  401. if debug_able:
  402. logger.info(f"连接page {page_num} 内的list")
  403. # 向layout_paras[i] 寻找开头具有相同缩进的连续的行
  404. may_list_lines = []
  405. next_page_first_para = next_page_paras[0][0]
  406. if next_page_first_para["type"] == BlockType.Text:
  407. lines = next_page_first_para["lines"]
  408. for line in lines:
  409. if line['bbox'][0] > __find_layout_bbox_by_line(line['bbox'], next_page_layout_bbox)[0]:
  410. may_list_lines.append(line)
  411. else:
  412. break
  413. # 如果这些行的缩进是相等的,那么连到上一个layout的最后一个段落上。
  414. if len(may_list_lines) > 0 and len(set([x['bbox'][0] for x in may_list_lines])) == 1:
  415. #pre_page_paras[-1].append(may_list_lines)
  416. # 下一页合并到上一页最后一段,打一个cross_page的标签
  417. for line in may_list_lines:
  418. for span in line["spans"]:
  419. span[CROSS_PAGE] = True
  420. pre_page_paras[-1][-1]["lines"].extend(may_list_lines)
  421. next_page_first_para["lines"] = next_page_first_para["lines"][len(may_list_lines):]
  422. return True
  423. return False
  424. def __find_layout_bbox_by_line(line_bbox, layout_bboxes):
  425. """
  426. 根据line找到所在的layout
  427. """
  428. for layout in layout_bboxes:
  429. if is_in_layout(line_bbox, layout):
  430. return layout
  431. return None
  432. def __connect_para_inter_layoutbox(blocks_group, new_layout_bbox):
  433. """
  434. layout之间进行分段。
  435. 主要是计算前一个layOut的最后一行和后一个layout的第一行是否可以连接。
  436. 连接的条件需要同时满足:
  437. 1. 上一个layout的最后一行沾满整个行。并且没有结尾符号。
  438. 2. 下一行开头不留空白。
  439. """
  440. connected_layout_blocks = []
  441. if len(blocks_group) == 0:
  442. return connected_layout_blocks
  443. connected_layout_blocks.append(blocks_group[0])
  444. for i in range(1, len(blocks_group)):
  445. try:
  446. if len(blocks_group[i]) == 0:
  447. continue
  448. if len(blocks_group[i - 1]) == 0: # TODO 考虑连接问题,
  449. connected_layout_blocks.append(blocks_group[i])
  450. continue
  451. # text类型的段才需要考虑layout间的合并
  452. if blocks_group[i - 1][-1]["type"] != BlockType.Text or blocks_group[i][0]["type"] != BlockType.Text:
  453. connected_layout_blocks.append(blocks_group[i])
  454. continue
  455. if len(blocks_group[i - 1][-1]["lines"]) == 0 or len(blocks_group[i][0]["lines"]) == 0:
  456. connected_layout_blocks.append(blocks_group[i])
  457. continue
  458. pre_last_line = blocks_group[i - 1][-1]["lines"][-1]
  459. next_first_line = blocks_group[i][0]["lines"][0]
  460. except Exception as e:
  461. logger.error(f"page layout {i} has no line")
  462. continue
  463. pre_last_line_text = ''.join([__get_span_text(span) for span in pre_last_line['spans']])
  464. pre_last_line_type = pre_last_line['spans'][-1]['type']
  465. next_first_line_text = ''.join([__get_span_text(span) for span in next_first_line['spans']])
  466. next_first_line_type = next_first_line['spans'][0]['type']
  467. if pre_last_line_type not in [TEXT, INLINE_EQUATION] or next_first_line_type not in [TEXT, INLINE_EQUATION]:
  468. #connected_layout_paras.append(layout_paras[i])
  469. connected_layout_blocks.append(blocks_group[i])
  470. continue
  471. pre_layout = __find_layout_bbox_by_line(pre_last_line['bbox'], new_layout_bbox)
  472. next_layout = __find_layout_bbox_by_line(next_first_line['bbox'], new_layout_bbox)
  473. pre_x2_max = pre_layout[2] if pre_layout else -1
  474. next_x0_min = next_layout[0] if next_layout else -1
  475. pre_last_line_text = pre_last_line_text.strip()
  476. next_first_line_text = next_first_line_text.strip()
  477. if pre_last_line['bbox'][2] == pre_x2_max and pre_last_line_text and pre_last_line_text[
  478. -1] not in LINE_STOP_FLAG and \
  479. next_first_line['bbox'][0] == next_x0_min: # 前面一行沾满了整个行,并且没有结尾符号.下一行没有空白开头。
  480. """连接段落条件成立,将前一个layout的段落和后一个layout的段落连接。"""
  481. #connected_layout_paras[-1][-1].extend(layout_paras[i][0])
  482. connected_layout_blocks[-1][-1]["lines"].extend(blocks_group[i][0]["lines"])
  483. #layout_paras[i].pop(0) # 删除后一个layout的第一个段落, 因为他已经被合并到前一个layout的最后一个段落了。
  484. blocks_group[i][0]["lines"] = [] #删除后一个layout第一个段落中的lines,因为他已经被合并到前一个layout的最后一个段落了
  485. blocks_group[i][0][LINES_DELETED] = True
  486. # if len(layout_paras[i]) == 0:
  487. # layout_paras.pop(i)
  488. # else:
  489. # connected_layout_paras.append(layout_paras[i])
  490. connected_layout_blocks.append(blocks_group[i])
  491. else:
  492. """连接段落条件不成立,将前一个layout的段落加入到结果中。"""
  493. #connected_layout_paras.append(layout_paras[i])
  494. connected_layout_blocks.append(blocks_group[i])
  495. return connected_layout_blocks
  496. def __connect_para_inter_page(pre_page_paras, next_page_paras, pre_page_layout_bbox, next_page_layout_bbox, page_num,
  497. lang):
  498. """
  499. 连接起来相邻两个页面的段落——前一个页面最后一个段落和后一个页面的第一个段落。
  500. 是否可以连接的条件:
  501. 1. 前一个页面的最后一个段落最后一行沾满整个行。并且没有结尾符号。
  502. 2. 后一个页面的第一个段落第一行没有空白开头。
  503. """
  504. # 有的页面可能压根没有文字
  505. if len(pre_page_paras) == 0 or len(next_page_paras) == 0 or len(pre_page_paras[0]) == 0 or len(
  506. next_page_paras[0]) == 0: # TODO [[]]为什么出现在pre_page_paras里?
  507. return False
  508. pre_last_block = pre_page_paras[-1][-1]
  509. next_first_block = next_page_paras[0][0]
  510. if pre_last_block["type"] != BlockType.Text or next_first_block["type"] != BlockType.Text:
  511. return False
  512. if len(pre_last_block["lines"]) == 0 or len(next_first_block["lines"]) == 0:
  513. return False
  514. pre_last_para = pre_last_block["lines"]
  515. next_first_para = next_first_block["lines"]
  516. pre_last_line = pre_last_para[-1]
  517. next_first_line = next_first_para[0]
  518. pre_last_line_text = ''.join([__get_span_text(span) for span in pre_last_line['spans']])
  519. pre_last_line_type = pre_last_line['spans'][-1]['type']
  520. next_first_line_text = ''.join([__get_span_text(span) for span in next_first_line['spans']])
  521. next_first_line_type = next_first_line['spans'][0]['type']
  522. if pre_last_line_type not in [TEXT, INLINE_EQUATION] or next_first_line_type not in [TEXT,
  523. INLINE_EQUATION]: # TODO,真的要做好,要考虑跨table, image, 行间的情况
  524. # 不是文本,不连接
  525. return False
  526. pre_x2_max_bbox = __find_layout_bbox_by_line(pre_last_line['bbox'], pre_page_layout_bbox)
  527. if not pre_x2_max_bbox:
  528. return False
  529. next_x0_min_bbox = __find_layout_bbox_by_line(next_first_line['bbox'], next_page_layout_bbox)
  530. if not next_x0_min_bbox:
  531. return False
  532. pre_x2_max = pre_x2_max_bbox[2]
  533. next_x0_min = next_x0_min_bbox[0]
  534. pre_last_line_text = pre_last_line_text.strip()
  535. next_first_line_text = next_first_line_text.strip()
  536. if pre_last_line['bbox'][2] == pre_x2_max and pre_last_line_text[-1] not in LINE_STOP_FLAG and \
  537. next_first_line['bbox'][0] == next_x0_min: # 前面一行沾满了整个行,并且没有结尾符号.下一行没有空白开头。
  538. """连接段落条件成立,将前一个layout的段落和后一个layout的段落连接。"""
  539. # 下一页合并到上一页最后一段,打一个cross_page的标签
  540. for line in next_first_para:
  541. for span in line["spans"]:
  542. span[CROSS_PAGE] = True
  543. pre_last_para.extend(next_first_para)
  544. #next_page_paras[0].pop(0) # 删除后一个页面的第一个段落, 因为他已经被合并到前一个页面的最后一个段落了。
  545. next_page_paras[0][0]["lines"] = []
  546. next_page_paras[0][0][LINES_DELETED] = True
  547. return True
  548. else:
  549. return False
  550. def find_consecutive_true_regions(input_array):
  551. start_index = None # 连续True区域的起始索引
  552. regions = [] # 用于保存所有连续True区域的起始和结束索引
  553. for i in range(len(input_array)):
  554. # 如果我们找到了一个True值,并且当前并没有在连续True区域中
  555. if input_array[i] and start_index is None:
  556. start_index = i # 记录连续True区域的起始索引
  557. # 如果我们找到了一个False值,并且当前在连续True区域中
  558. elif not input_array[i] and start_index is not None:
  559. # 如果连续True区域长度大于1,那么将其添加到结果列表中
  560. if i - start_index > 1:
  561. regions.append((start_index, i - 1))
  562. start_index = None # 重置起始索引
  563. # 如果最后一个元素是True,那么需要将最后一个连续True区域加入到结果列表中
  564. if start_index is not None and len(input_array) - start_index > 1:
  565. regions.append((start_index, len(input_array) - 1))
  566. return regions
  567. def __connect_middle_align_text(page_paras, new_layout_bbox, page_num, lang):
  568. global debug_able
  569. """
  570. 找出来中间对齐的连续单行文本,如果连续行高度相同,那么合并为一个段落。
  571. 一个line居中的条件是:
  572. 1. 水平中心点跨越layout的中心点。
  573. 2. 左右两侧都有空白
  574. """
  575. for layout_i, layout_para in enumerate(page_paras):
  576. layout_box = new_layout_bbox[layout_i]
  577. single_line_paras_tag = []
  578. for i in range(len(layout_para)):
  579. #single_line_paras_tag.append(len(layout_para[i]) == 1 and layout_para[i][0]['spans'][0]['type'] == TEXT)
  580. single_line_paras_tag.append(layout_para[i]['type'] == BlockType.Text and len(layout_para[i]["lines"]) == 1)
  581. """找出来连续的单行文本,如果连续行高度相同,那么合并为一个段落。"""
  582. consecutive_single_line_indices = find_consecutive_true_regions(single_line_paras_tag)
  583. if len(consecutive_single_line_indices) > 0:
  584. #index_offset = 0
  585. """检查这些行是否是高度相同的,居中的"""
  586. for start, end in consecutive_single_line_indices:
  587. #start += index_offset
  588. #end += index_offset
  589. line_hi = np.array([block["lines"][0]['bbox'][3] - block["lines"][0]['bbox'][1] for block in
  590. layout_para[start:end + 1]])
  591. first_line_text = ''.join([__get_span_text(span) for span in layout_para[start]["lines"][0]['spans']])
  592. if "Table" in first_line_text or "Figure" in first_line_text:
  593. pass
  594. if debug_able:
  595. logger.info(line_hi.std())
  596. if line_hi.std() < 2:
  597. """行高度相同,那么判断是否居中"""
  598. all_left_x0 = [block["lines"][0]['bbox'][0] for block in layout_para[start:end + 1]]
  599. all_right_x1 = [block["lines"][0]['bbox'][2] for block in layout_para[start:end + 1]]
  600. layout_center = (layout_box[0] + layout_box[2]) / 2
  601. if all([x0 < layout_center < x1 for x0, x1 in zip(all_left_x0, all_right_x1)]) \
  602. and not all([x0 == layout_box[0] for x0 in all_left_x0]) \
  603. and not all([x1 == layout_box[2] for x1 in all_right_x1]):
  604. merge_para = [block["lines"][0] for block in layout_para[start:end + 1]]
  605. para_text = ''.join([__get_span_text(span) for line in merge_para for span in line['spans']])
  606. if debug_able:
  607. logger.info(para_text)
  608. layout_para[start]["lines"] = merge_para
  609. for i_para in range(start + 1, end + 1):
  610. layout_para[i_para]["lines"] = []
  611. layout_para[i_para][LINES_DELETED] = True
  612. #layout_para[start:end + 1] = [merge_para]
  613. #index_offset -= end - start
  614. return
  615. def __merge_signle_list_text(page_paras, new_layout_bbox, page_num, lang):
  616. """
  617. 找出来连续的单行文本,如果首行顶格,接下来的几个单行段落缩进对齐,那么合并为一个段落。
  618. """
  619. pass
  620. def __do_split_page(blocks, layout_bboxes, new_layout_bbox, page_num, lang):
  621. """
  622. 根据line和layout情况进行分段
  623. 先实现一个根据行末尾特征分段的简单方法。
  624. """
  625. """
  626. 算法思路:
  627. 1. 扫描layout里每一行,找出来行尾距离layout有边界有一定距离的行。
  628. 2. 从上述行中找到末尾是句号等可作为断行标志的行。
  629. 3. 参照上述行尾特征进行分段。
  630. 4. 图、表,目前独占一行,不考虑分段。
  631. """
  632. blocks_group = __group_line_by_layout(blocks, layout_bboxes) # block内分段
  633. layout_list_info = __split_para_in_layoutbox(blocks_group, new_layout_bbox, lang) # layout内分段
  634. blocks_group, page_list_info = __connect_list_inter_layout(blocks_group, new_layout_bbox, layout_list_info,
  635. page_num, lang) # layout之间连接列表段落
  636. connected_layout_blocks = __connect_para_inter_layoutbox(blocks_group, new_layout_bbox) # layout间链接段落
  637. return connected_layout_blocks, page_list_info
  638. def para_split(pdf_info_dict, debug_mode, lang="en"):
  639. global debug_able
  640. debug_able = debug_mode
  641. new_layout_of_pages = [] # 数组的数组,每个元素是一个页面的layoutS
  642. all_page_list_info = [] # 保存每个页面开头和结尾是否是列表
  643. for page_num, page in pdf_info_dict.items():
  644. blocks = page['preproc_blocks']
  645. layout_bboxes = page['layout_bboxes']
  646. new_layout_bbox = __common_pre_proc(blocks, layout_bboxes)
  647. new_layout_of_pages.append(new_layout_bbox)
  648. splited_blocks, page_list_info = __do_split_page(blocks, layout_bboxes, new_layout_bbox, page_num, lang)
  649. all_page_list_info.append(page_list_info)
  650. page['para_blocks'] = splited_blocks
  651. """连接页面与页面之间的可能合并的段落"""
  652. pdf_infos = list(pdf_info_dict.values())
  653. for page_num, page in enumerate(pdf_info_dict.values()):
  654. if page_num == 0:
  655. continue
  656. pre_page_paras = pdf_infos[page_num - 1]['para_blocks']
  657. next_page_paras = pdf_infos[page_num]['para_blocks']
  658. pre_page_layout_bbox = new_layout_of_pages[page_num - 1]
  659. next_page_layout_bbox = new_layout_of_pages[page_num]
  660. is_conn = __connect_para_inter_page(pre_page_paras, next_page_paras, pre_page_layout_bbox,
  661. next_page_layout_bbox, page_num, lang)
  662. if debug_able:
  663. if is_conn:
  664. logger.info(f"连接了第{page_num - 1}页和第{page_num}页的段落")
  665. is_list_conn = __connect_list_inter_page(pre_page_paras, next_page_paras, pre_page_layout_bbox,
  666. next_page_layout_bbox, all_page_list_info[page_num - 1],
  667. all_page_list_info[page_num], page_num, lang)
  668. if debug_able:
  669. if is_list_conn:
  670. logger.info(f"连接了第{page_num - 1}页和第{page_num}页的列表段落")
  671. """接下来可能会漏掉一些特别的一些可以合并的内容,对他们进行段落连接
  672. 1. 正文中有时出现一个行顶格,接下来几行缩进的情况。
  673. 2. 居中的一些连续单行,如果高度相同,那么可能是一个段落。
  674. """
  675. for page_num, page in enumerate(pdf_info_dict.values()):
  676. page_paras = page['para_blocks']
  677. new_layout_bbox = new_layout_of_pages[page_num]
  678. __connect_middle_align_text(page_paras, new_layout_bbox, page_num, lang)
  679. __merge_signle_list_text(page_paras, new_layout_bbox, page_num, lang)
  680. # layout展平
  681. for page_num, page in enumerate(pdf_info_dict.values()):
  682. page_paras = page['para_blocks']
  683. page_blocks = [block for layout in page_paras for block in layout]
  684. page["para_blocks"] = page_blocks