para_split_v3.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313
  1. import copy
  2. from loguru import logger
  3. from magic_pdf.libs.Constants import LINES_DELETED, CROSS_PAGE
  4. from magic_pdf.libs.ocr_content_type import BlockType, ContentType
  5. LINE_STOP_FLAG = ('.', '!', '?', '。', '!', '?', ')', ')', '"', '”', ':', ':', ';', ';')
  6. LIST_END_FLAG = ('.', '。', ';', ';')
  7. class ListLineTag:
  8. IS_LIST_START_LINE = "is_list_start_line"
  9. IS_LIST_END_LINE = "is_list_end_line"
  10. def __process_blocks(blocks):
  11. # 对所有block预处理
  12. # 1.通过title和interline_equation将block分组
  13. # 2.bbox边界根据line信息重置
  14. result = []
  15. current_group = []
  16. for i in range(len(blocks)):
  17. current_block = blocks[i]
  18. # 如果当前块是 text 类型
  19. if current_block['type'] == 'text':
  20. current_block["bbox_fs"] = copy.deepcopy(current_block["bbox"])
  21. if 'lines' in current_block and len(current_block["lines"]) > 0:
  22. current_block['bbox_fs'] = [min([line['bbox'][0] for line in current_block['lines']]),
  23. min([line['bbox'][1] for line in current_block['lines']]),
  24. max([line['bbox'][2] for line in current_block['lines']]),
  25. max([line['bbox'][3] for line in current_block['lines']])]
  26. current_group.append(current_block)
  27. # 检查下一个块是否存在
  28. if i + 1 < len(blocks):
  29. next_block = blocks[i + 1]
  30. # 如果下一个块不是 text 类型且是 title 或 interline_equation 类型
  31. if next_block['type'] in ['title', 'interline_equation']:
  32. result.append(current_group)
  33. current_group = []
  34. # 处理最后一个 group
  35. if current_group:
  36. result.append(current_group)
  37. return result
  38. def __is_list_or_index_block(block):
  39. # 一个block如果是list block 应该同时满足以下特征
  40. # 1.block内有多个line 2.block 内有多个line左侧顶格写 3.block内有多个line 右侧不顶格(狗牙状)
  41. # 1.block内有多个line 2.block 内有多个line左侧顶格写 3.多个line以endflag结尾
  42. # 1.block内有多个line 2.block 内有多个line左侧顶格写 3.block内有多个line 左侧不顶格
  43. # index block 是一种特殊的list block
  44. # 一个block如果是index block 应该同时满足以下特征
  45. # 1.block内有多个line 2.block 内有多个line两侧均顶格写 3.line的开头或者结尾均为数字
  46. if len(block['lines']) >= 2:
  47. first_line = block['lines'][0]
  48. line_height = first_line['bbox'][3] - first_line['bbox'][1]
  49. block_weight = block['bbox_fs'][2] - block['bbox_fs'][0]
  50. left_close_num = 0
  51. left_not_close_num = 0
  52. right_not_close_num = 0
  53. right_close_num = 0
  54. lines_text_list = []
  55. center_close_num = 0
  56. external_sides_not_close_num = 0
  57. multiple_para_flag = False
  58. last_line = block['lines'][-1]
  59. # 如果首行左边不顶格而右边顶格,末行左边顶格而右边不顶格 (第一行可能可以右边不顶格)
  60. if (first_line['bbox'][0] - block['bbox_fs'][0] > line_height / 2 and
  61. # block['bbox_fs'][2] - first_line['bbox'][2] < line_height and
  62. abs(last_line['bbox'][0] - block['bbox_fs'][0]) < line_height / 2 and
  63. block['bbox_fs'][2] - last_line['bbox'][2] > line_height
  64. ):
  65. multiple_para_flag = True
  66. for line in block['lines']:
  67. line_mid_x = (line['bbox'][0] + line['bbox'][2]) / 2
  68. block_mid_x = (block['bbox_fs'][0] + block['bbox_fs'][2]) / 2
  69. if (line['bbox'][0] - block['bbox_fs'][0] > 0.8 * line_height and
  70. block['bbox_fs'][2] - line['bbox'][2] > 0.8 * line_height):
  71. external_sides_not_close_num += 1
  72. if abs(line_mid_x - block_mid_x) < line_height/2:
  73. center_close_num += 1
  74. line_text = ""
  75. for span in line['spans']:
  76. span_type = span['type']
  77. if span_type == ContentType.Text:
  78. line_text += span['content'].strip()
  79. lines_text_list.append(line_text)
  80. # 计算line左侧顶格数量是否大于2,是否顶格用abs(block['bbox_fs'][0] - line['bbox'][0]) < line_height/2 来判断
  81. if abs(block['bbox_fs'][0] - line['bbox'][0]) < line_height / 2:
  82. left_close_num += 1
  83. elif line['bbox'][0] - block['bbox_fs'][0] > line_height:
  84. # logger.info(f"{line_text}, {block['bbox_fs']}, {line['bbox']}")
  85. left_not_close_num += 1
  86. # 计算右侧是否顶格
  87. if abs(block['bbox_fs'][2] - line['bbox'][2]) < line_height:
  88. right_close_num += 1
  89. else:
  90. # 右侧不顶格情况下是否有一段距离,拍脑袋用0.3block宽度做阈值
  91. closed_area = 0.26 * block_weight
  92. # closed_area = 5 * line_height
  93. if block['bbox_fs'][2] - line['bbox'][2] > closed_area:
  94. right_not_close_num += 1
  95. # 判断lines_text_list中的元素是否有超过80%都以LIST_END_FLAG结尾
  96. line_end_flag = False
  97. # 判断lines_text_list中的元素是否有超过80%都以数字开头或都以数字结尾
  98. line_num_flag = False
  99. num_start_count = 0
  100. num_end_count = 0
  101. flag_end_count = 0
  102. if len(lines_text_list) > 0:
  103. for line_text in lines_text_list:
  104. if len(line_text) > 0:
  105. if line_text[-1] in LIST_END_FLAG:
  106. flag_end_count += 1
  107. if line_text[0].isdigit():
  108. num_start_count += 1
  109. if line_text[-1].isdigit():
  110. num_end_count += 1
  111. if flag_end_count / len(lines_text_list) >= 0.8:
  112. line_end_flag = True
  113. if num_start_count / len(lines_text_list) >= 0.8 or num_end_count / len(lines_text_list) >= 0.8:
  114. line_num_flag = True
  115. # 有的目录右侧不贴边, 目前认为左边或者右边有一边全贴边,且符合数字规则极为index
  116. if ((left_close_num/len(block['lines']) >= 0.8 or right_close_num/len(block['lines']) >= 0.8)
  117. and line_num_flag
  118. ):
  119. for line in block['lines']:
  120. line[ListLineTag.IS_LIST_START_LINE] = True
  121. return BlockType.Index
  122. # 全部line都居中的特殊list识别,每行都需要换行,特征是多行,且大多数行都前后not_close,每line中点x坐标接近
  123. elif external_sides_not_close_num >= 2 and center_close_num == len(block['lines']) and external_sides_not_close_num / len(block['lines']) >= 0.5:
  124. for line in block['lines']:
  125. line[ListLineTag.IS_LIST_START_LINE] = True
  126. return BlockType.List
  127. elif left_close_num >= 2 and (
  128. right_not_close_num >= 2 or line_end_flag or left_not_close_num >= 2) and not multiple_para_flag:
  129. # 处理一种特殊的没有缩进的list,所有行都贴左边,通过右边的空隙判断是否是item尾
  130. if left_close_num / len(block['lines']) > 0.8:
  131. # 这种是每个item只有一行,且左边都贴边的短item list
  132. if flag_end_count == 0 and right_close_num / len(block['lines']) < 0.5:
  133. for line in block['lines']:
  134. if abs(block['bbox_fs'][0] - line['bbox'][0]) < line_height / 2:
  135. line[ListLineTag.IS_LIST_START_LINE] = True
  136. # 这种是大部分line item 都有结束标识符的情况,按结束标识符区分不同item
  137. elif line_end_flag:
  138. for i, line in enumerate(block['lines']):
  139. if lines_text_list[i][-1] in LIST_END_FLAG:
  140. line[ListLineTag.IS_LIST_END_LINE] = True
  141. if i + 1 < len(block['lines']):
  142. block['lines'][i+1][ListLineTag.IS_LIST_START_LINE] = True
  143. # line item基本没有结束标识符,而且也没有缩进,按右侧空隙判断哪些是item end
  144. else:
  145. line_start_flag = False
  146. for i, line in enumerate(block['lines']):
  147. if line_start_flag:
  148. line[ListLineTag.IS_LIST_START_LINE] = True
  149. line_start_flag = False
  150. # elif abs(block['bbox_fs'][2] - line['bbox'][2]) > line_height:
  151. if abs(block['bbox_fs'][2] - line['bbox'][2]) > 0.1 * block_weight:
  152. line[ListLineTag.IS_LIST_END_LINE] = True
  153. line_start_flag = True
  154. # 一种有缩进的特殊有序list,start line 左侧不贴边且以数字开头,end line 以 IS_LIST_END_LINE 结尾且数量和start line 一致
  155. elif num_start_count >= 2 and num_start_count == flag_end_count: # 简单一点先不考虑左侧不贴边的情况
  156. for i, line in enumerate(block['lines']):
  157. if lines_text_list[i][0].isdigit():
  158. line[ListLineTag.IS_LIST_START_LINE] = True
  159. if lines_text_list[i][-1] in LIST_END_FLAG:
  160. line[ListLineTag.IS_LIST_END_LINE] = True
  161. else:
  162. # 正常有缩进的list处理
  163. for line in block['lines']:
  164. if abs(block['bbox_fs'][0] - line['bbox'][0]) < line_height / 2:
  165. line[ListLineTag.IS_LIST_START_LINE] = True
  166. if abs(block['bbox_fs'][2] - line['bbox'][2]) > line_height:
  167. line[ListLineTag.IS_LIST_END_LINE] = True
  168. return BlockType.List
  169. else:
  170. return BlockType.Text
  171. else:
  172. return BlockType.Text
  173. def __merge_2_text_blocks(block1, block2):
  174. if len(block1['lines']) > 0:
  175. first_line = block1['lines'][0]
  176. line_height = first_line['bbox'][3] - first_line['bbox'][1]
  177. block1_weight = block1['bbox'][2] - block1['bbox'][0]
  178. block2_weight = block2['bbox'][2] - block2['bbox'][0]
  179. min_block_weight = min(block1_weight, block2_weight)
  180. if abs(block1['bbox_fs'][0] - first_line['bbox'][0]) < line_height / 2:
  181. last_line = block2['lines'][-1]
  182. if len(last_line['spans']) > 0:
  183. last_span = last_line['spans'][-1]
  184. line_height = last_line['bbox'][3] - last_line['bbox'][1]
  185. if (abs(block2['bbox_fs'][2] - last_line['bbox'][2]) < line_height and
  186. not last_span['content'].endswith(LINE_STOP_FLAG) and
  187. # 两个block宽度差距超过2倍也不合并
  188. abs(block1_weight - block2_weight) < min_block_weight
  189. ):
  190. if block1['page_num'] != block2['page_num']:
  191. for line in block1['lines']:
  192. for span in line['spans']:
  193. span[CROSS_PAGE] = True
  194. block2['lines'].extend(block1['lines'])
  195. block1['lines'] = []
  196. block1[LINES_DELETED] = True
  197. return block1, block2
  198. def __merge_2_list_blocks(block1, block2):
  199. if block1['page_num'] != block2['page_num']:
  200. for line in block1['lines']:
  201. for span in line['spans']:
  202. span[CROSS_PAGE] = True
  203. block2['lines'].extend(block1['lines'])
  204. block1['lines'] = []
  205. block1[LINES_DELETED] = True
  206. return block1, block2
  207. def __is_list_group(text_blocks_group):
  208. # list group的特征是一个group内的所有block都满足以下条件
  209. # 1.每个block都不超过3行 2. 每个block 的左边界都比较接近(逻辑简单点先不加这个规则)
  210. for block in text_blocks_group:
  211. if len(block['lines']) > 3:
  212. return False
  213. return True
  214. def __para_merge_page(blocks):
  215. page_text_blocks_groups = __process_blocks(blocks)
  216. for text_blocks_group in page_text_blocks_groups:
  217. if len(text_blocks_group) > 0:
  218. # 需要先在合并前对所有block判断是否为list or index block
  219. for block in text_blocks_group:
  220. block_type = __is_list_or_index_block(block)
  221. block['type'] = block_type
  222. # logger.info(f"{block['type']}:{block}")
  223. if len(text_blocks_group) > 1:
  224. # 在合并前判断这个group 是否是一个 list group
  225. is_list_group = __is_list_group(text_blocks_group)
  226. # 倒序遍历
  227. for i in range(len(text_blocks_group) - 1, -1, -1):
  228. current_block = text_blocks_group[i]
  229. # 检查是否有前一个块
  230. if i - 1 >= 0:
  231. prev_block = text_blocks_group[i - 1]
  232. if current_block['type'] == 'text' and prev_block['type'] == 'text' and not is_list_group:
  233. __merge_2_text_blocks(current_block, prev_block)
  234. elif (
  235. (current_block['type'] == BlockType.List and prev_block['type'] == BlockType.List) or
  236. (current_block['type'] == BlockType.Index and prev_block['type'] == BlockType.Index)
  237. ):
  238. __merge_2_list_blocks(current_block, prev_block)
  239. else:
  240. continue
  241. def para_split(pdf_info_dict, debug_mode=False):
  242. all_blocks = []
  243. for page_num, page in pdf_info_dict.items():
  244. blocks = copy.deepcopy(page['preproc_blocks'])
  245. for block in blocks:
  246. block['page_num'] = page_num
  247. all_blocks.extend(blocks)
  248. __para_merge_page(all_blocks)
  249. for page_num, page in pdf_info_dict.items():
  250. page['para_blocks'] = []
  251. for block in all_blocks:
  252. if block['page_num'] == page_num:
  253. page['para_blocks'].append(block)
  254. if __name__ == '__main__':
  255. input_blocks = []
  256. # 调用函数
  257. groups = __process_blocks(input_blocks)
  258. for group_index, group in enumerate(groups):
  259. print(f"Group {group_index}: {group}")