para_split.py.bak 32 KB

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