layout_sort.py 33 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732
  1. """
  2. 对pdf上的box进行layout识别,并对内部组成的box进行排序
  3. """
  4. from loguru import logger
  5. from pdf_tools.layout.bbox_sort import CONTENT_IDX, CONTENT_TYPE_IDX, X0_EXT_IDX, X0_IDX, X1_EXT_IDX, X1_IDX, Y0_EXT_IDX, Y0_IDX, Y1_EXT_IDX, Y1_IDX, paper_bbox_sort
  6. from pdf_tools.layout.layout_det_utils import find_all_left_bbox_direct, find_all_right_bbox_direct, find_bottom_bbox_direct_from_left_edge, find_bottom_bbox_direct_from_right_edge, find_top_bbox_direct_from_left_edge, find_top_bbox_direct_from_right_edge, find_all_top_bbox_direct, find_all_bottom_bbox_direct, get_left_edge_bboxes, get_right_edge_bboxes
  7. from pdf_tools.libs.boxbase import get_bbox_in_boundry
  8. LAYOUT_V = "V"
  9. LAYOUT_H = "H"
  10. LAYOUT_UNPROC = "U"
  11. LAYOUT_BAD = "B"
  12. def _is_single_line_text(bbox):
  13. """
  14. 检查bbox里面的文字是否只有一行
  15. """
  16. return True # TODO
  17. box_type = bbox[CONTENT_TYPE_IDX]
  18. if box_type != 'text':
  19. return False
  20. paras = bbox[CONTENT_IDX]["paras"]
  21. text_content = ""
  22. for para_id, para in paras.items(): # 拼装内部的段落文本
  23. is_title = para['is_title']
  24. if is_title!=0:
  25. text_content += f"## {para['text']}"
  26. else:
  27. text_content += para["text"]
  28. text_content += "\n\n"
  29. return bbox[CONTENT_TYPE_IDX] == 'text' and len(text_content.split("\n\n")) <= 1
  30. def _horizontal_split(bboxes:list, boundry:tuple, avg_font_size=20)-> list:
  31. """
  32. 对bboxes进行水平切割
  33. 方法是:找到左侧和右侧都没有被直接遮挡的box,然后进行扩展,之后进行切割
  34. return:
  35. 返回几个大的Layout区域 [[x0, y0, x1, y1, "h|u|v"], ], h代表水平,u代表未探测的,v代表垂直布局
  36. """
  37. sorted_layout_blocks = [] # 这是要最终返回的值
  38. bound_x0, bound_y0, bound_x1, bound_y1 = boundry
  39. all_bboxes = get_bbox_in_boundry(bboxes, boundry)
  40. #all_bboxes = paper_bbox_sort(all_bboxes, abs(bound_x1-bound_x0), abs(bound_y1-bound_x0)) # 大致拍下序, 这个是基于直接遮挡的。
  41. """
  42. 首先在水平方向上扩展独占一行的bbox
  43. """
  44. last_h_split_line_y1 = bound_y0 #记录下上次的水平分割线
  45. for i, bbox in enumerate(all_bboxes):
  46. left_nearest_bbox = find_all_left_bbox_direct(bbox, all_bboxes) # 非扩展线
  47. right_nearest_bbox = find_all_right_bbox_direct(bbox, all_bboxes)
  48. if left_nearest_bbox is None and right_nearest_bbox is None: # 独占一行
  49. """
  50. 然而,如果只是孤立的一行文字,那么就还要满足以下几个条件才可以:
  51. 1. bbox和中心线相交。或者
  52. 2. 上方或者下方也存在同类水平的独占一行的bbox。 或者
  53. 3. TODO 加强条件:这个bbox上方和下方是同一列column,那么就不能算作独占一行
  54. """
  55. # 先检查这个bbox里是否只包含一行文字
  56. is_single_line = _is_single_line_text(bbox)
  57. """
  58. 这里有个点需要注意,当页面内容不是居中的时候,第一次调用传递的是page的boundry,这个时候mid_x就不是中心线了.
  59. 所以这里计算出最紧致的boundry,然后再计算mid_x
  60. """
  61. boundry_real_x0, boundry_real_x1 = min([bbox[X0_IDX] for bbox in all_bboxes]), max([bbox[X1_IDX] for bbox in all_bboxes])
  62. mid_x = (boundry_real_x0+boundry_real_x1)/2
  63. # 检查这个box是否内容在中心线有交
  64. # 必须跨过去2个字符的宽度
  65. is_cross_boundry_mid_line = min(mid_x-bbox[X0_IDX], bbox[X1_IDX]-mid_x) > avg_font_size*2
  66. """
  67. 检查条件2
  68. """
  69. is_belong_to_col = False
  70. """
  71. 检查是否能被上方col吸收,方法是:
  72. 1. 上方非空且不是独占一行的,并且
  73. 2. 从上个水平分割的最大y=y1开始到当前bbox,最左侧的bbox的[min_x0, max_x1],能够覆盖当前box的[x0, x1]
  74. """
  75. """
  76. 以迭代的方式向上找,查找范围是[bound_x0, last_h_sp, bound_x1, bbox[Y0_IDX]]
  77. """
  78. #先确定上方的y0, y0
  79. b_y0, b_y1 = last_h_split_line_y1, bbox[Y0_IDX]
  80. #然后从box开始逐个向上找到所有与box在x上有交集的box
  81. box_to_check = [bound_x0, b_y0, bound_x1, b_y1]
  82. bbox_in_bound_check = get_bbox_in_boundry(all_bboxes, box_to_check)
  83. bboxes_on_top = []
  84. virtual_box = bbox
  85. while True:
  86. b_on_top = find_all_top_bbox_direct(virtual_box, bbox_in_bound_check)
  87. if b_on_top is not None:
  88. bboxes_on_top.append(b_on_top)
  89. virtual_box = [min([virtual_box[X0_IDX], b_on_top[X0_IDX]]), min(virtual_box[Y0_IDX], b_on_top[Y0_IDX]), max([virtual_box[X1_IDX], b_on_top[X1_IDX]]), b_y1]
  90. else:
  91. break
  92. # 随后确定这些box的最小x0, 最大x1
  93. if len(bboxes_on_top)>0 and len(bboxes_on_top) != len(bbox_in_bound_check):# virtual_box可能会膨胀到占满整个区域,这实际上就不能属于一个col了。
  94. min_x0, max_x1 = virtual_box[X0_IDX], virtual_box[X1_IDX]
  95. # 然后采用一种比较粗糙的方法,看min_x0,max_x1是否与位于[bound_x0, last_h_sp, bound_x1, bbox[Y0_IDX]]之间的box有相交
  96. if not any([b[X0_IDX] <= min_x0-1 <= b[X1_IDX] or b[X0_IDX] <= max_x1+1 <= b[X1_IDX] for b in bbox_in_bound_check]):
  97. # 其上,下都不能被扩展成行,暂时只检查一下上方 TODO
  98. top_nearest_bbox = find_all_top_bbox_direct(bbox, bboxes)
  99. bottom_nearest_bbox = find_all_bottom_bbox_direct(bbox, bboxes)
  100. if not any([
  101. top_nearest_bbox is not None and (find_all_left_bbox_direct(top_nearest_bbox, bboxes) is None and find_all_right_bbox_direct(top_nearest_bbox, bboxes) is None),
  102. bottom_nearest_bbox is not None and (find_all_left_bbox_direct(bottom_nearest_bbox, bboxes) is None and find_all_right_bbox_direct(bottom_nearest_bbox, bboxes) is None),
  103. top_nearest_bbox is None or bottom_nearest_bbox is None
  104. ]):
  105. is_belong_to_col = True
  106. # 检查是否能被下方col吸收 TODO
  107. """
  108. 这里为什么没有is_cross_boundry_mid_line的条件呢?
  109. 确实有些杂志左右两栏宽度不是对称的。
  110. """
  111. if not is_belong_to_col or is_cross_boundry_mid_line:
  112. bbox[X0_EXT_IDX] = bound_x0
  113. bbox[Y0_EXT_IDX] = bbox[Y0_IDX]
  114. bbox[X1_EXT_IDX] = bound_x1
  115. bbox[Y1_EXT_IDX] = bbox[Y1_IDX]
  116. last_h_split_line_y1 = bbox[Y1_IDX] # 更新这条线
  117. else:
  118. continue
  119. """
  120. 此时独占一行的被成功扩展到指定的边界上,这个时候利用边界条件合并连续的bbox,成为一个group
  121. 然后合并所有连续水平方向的bbox.
  122. """
  123. all_bboxes.sort(key=lambda x: x[Y0_IDX])
  124. h_bboxes = []
  125. h_bbox_group = []
  126. for bbox in all_bboxes:
  127. if bbox[X0_EXT_IDX] == bound_x0 and bbox[X1_EXT_IDX] == bound_x1:
  128. h_bbox_group.append(bbox)
  129. else:
  130. if len(h_bbox_group)>0:
  131. h_bboxes.append(h_bbox_group)
  132. h_bbox_group = []
  133. # 最后一个group
  134. if len(h_bbox_group)>0:
  135. h_bboxes.append(h_bbox_group)
  136. """
  137. 现在h_bboxes里面是所有的group了,每个group都是一个list
  138. 对h_bboxes里的每个group进行计算放回到sorted_layouts里
  139. """
  140. h_layouts = []
  141. for gp in h_bboxes:
  142. gp.sort(key=lambda x: x[Y0_IDX])
  143. # 然后计算这个group的layout_bbox,也就是最小的x0,y0, 最大的x1,y1
  144. x0, y0, x1, y1 = gp[0][X0_EXT_IDX], gp[0][Y0_EXT_IDX], gp[-1][X1_EXT_IDX], gp[-1][Y1_EXT_IDX]
  145. h_layouts.append([x0, y0, x1, y1, LAYOUT_H]) # 水平的布局
  146. """
  147. 接下来利用这些连续的水平bbox的layout_bbox的y0, y1,从水平上切分开其余的为几个部分
  148. """
  149. h_split_lines = [bound_y0]
  150. for gp in h_bboxes: # gp是一个list[bbox_list]
  151. y0, y1 = gp[0][1], gp[-1][3]
  152. h_split_lines.append(y0)
  153. h_split_lines.append(y1)
  154. h_split_lines.append(bound_y1)
  155. unsplited_bboxes = []
  156. for i in range(0, len(h_split_lines), 2):
  157. start_y0, start_y1 = h_split_lines[i:i+2]
  158. # 然后找出[start_y0, start_y1]之间的其他bbox,这些组成一个未分割板块
  159. bboxes_in_block = [bbox for bbox in all_bboxes if bbox[Y0_IDX]>=start_y0 and bbox[Y1_IDX]<=start_y1]
  160. unsplited_bboxes.append(bboxes_in_block)
  161. # 接着把未处理的加入到h_layouts里
  162. for bboxes_in_block in unsplited_bboxes:
  163. if len(bboxes_in_block) == 0:
  164. continue
  165. x0, y0, x1, y1 = bound_x0, min([bbox[Y0_IDX] for bbox in bboxes_in_block]), bound_x1, max([bbox[Y1_IDX] for bbox in bboxes_in_block])
  166. h_layouts.append([x0, y0, x1, y1, LAYOUT_UNPROC])
  167. h_layouts.sort(key=lambda x: x[1]) # 按照y0排序, 也就是从上到下的顺序
  168. """
  169. 转换成如下格式返回
  170. """
  171. for layout in h_layouts:
  172. sorted_layout_blocks.append({
  173. "layout_bbox": layout[:4],
  174. "layout_label":layout[4],
  175. "sub_layout":[],
  176. })
  177. return sorted_layout_blocks
  178. ###############################################################################################
  179. #
  180. # 垂直方向的处理
  181. #
  182. #
  183. ###############################################################################################
  184. def _vertical_align_split_v1(bboxes:list, boundry:tuple)-> list:
  185. """
  186. 计算垂直方向上的对齐, 并分割bboxes成layout。负责对一列多行的进行列维度分割。
  187. 如果不能完全分割,剩余部分作为layout_lable为u的layout返回
  188. -----------------------
  189. | | |
  190. | | |
  191. | | |
  192. | | |
  193. -------------------------
  194. 此函数会将:以上布局将会切分出来2列
  195. """
  196. sorted_layout_blocks = [] # 这是要最终返回的值
  197. new_boundry = [boundry[0], boundry[1], boundry[2], boundry[3]]
  198. v_blocks = []
  199. """
  200. 先从左到右切分
  201. """
  202. while True:
  203. all_bboxes = get_bbox_in_boundry(bboxes, new_boundry)
  204. left_edge_bboxes = get_left_edge_bboxes(all_bboxes)
  205. if len(left_edge_bboxes) == 0:
  206. break
  207. right_split_line_x1 = max([bbox[X1_IDX] for bbox in left_edge_bboxes])+1
  208. # 然后检查这条线能不与其他bbox的左边界相交或者重合
  209. if any([bbox[X0_IDX] <= right_split_line_x1 <= bbox[X1_IDX] for bbox in all_bboxes]):
  210. # 垂直切分线与某些box发生相交,说明无法完全垂直方向切分。
  211. break
  212. else: # 说明成功分割出一列
  213. # 找到左侧边界最靠左的bbox作为layout的x0
  214. layout_x0 = min([bbox[X0_IDX] for bbox in left_edge_bboxes]) # 这里主要是为了画出来有一定间距
  215. v_blocks.append([layout_x0, new_boundry[1], right_split_line_x1, new_boundry[3], LAYOUT_V])
  216. new_boundry[0] = right_split_line_x1 # 更新边界
  217. """
  218. 再从右到左切, 此时如果还是无法完全切分,那么剩余部分作为layout_lable为u的layout返回
  219. """
  220. unsplited_block = []
  221. while True:
  222. all_bboxes = get_bbox_in_boundry(bboxes, new_boundry)
  223. right_edge_bboxes = get_right_edge_bboxes(all_bboxes)
  224. if len(right_edge_bboxes) == 0:
  225. break
  226. left_split_line_x0 = min([bbox[X0_IDX] for bbox in right_edge_bboxes])-1
  227. # 然后检查这条线能不与其他bbox的左边界相交或者重合
  228. if any([bbox[X0_IDX] <= left_split_line_x0 <= bbox[X1_IDX] for bbox in all_bboxes]):
  229. # 这里是余下的
  230. unsplited_block.append([new_boundry[0], new_boundry[1], new_boundry[2], new_boundry[3], LAYOUT_UNPROC])
  231. break
  232. else:
  233. # 找到右侧边界最靠右的bbox作为layout的x1
  234. layout_x1 = max([bbox[X1_IDX] for bbox in right_edge_bboxes])
  235. v_blocks.append([left_split_line_x0, new_boundry[1], layout_x1, new_boundry[3], LAYOUT_V])
  236. new_boundry[2] = left_split_line_x0 # 更新右边界
  237. """
  238. 最后拼装成layout格式返回
  239. """
  240. for block in v_blocks:
  241. sorted_layout_blocks.append({
  242. "layout_bbox": block[:4],
  243. "layout_label":block[4],
  244. "sub_layout":[],
  245. })
  246. for block in unsplited_block:
  247. sorted_layout_blocks.append({
  248. "layout_bbox": block[:4],
  249. "layout_label":block[4],
  250. "sub_layout":[],
  251. })
  252. # 按照x0排序
  253. sorted_layout_blocks.sort(key=lambda x: x['layout_bbox'][0])
  254. return sorted_layout_blocks
  255. def _vertical_align_split_v2(bboxes:list, boundry:tuple)-> list:
  256. """
  257. 改进的 _vertical_align_split算法,原算法会因为第二列的box由于左侧没有遮挡被认为是左侧的一部分,导致整个layout多列被识别为一列。
  258. 利用从左上角的box开始向下看的方法,不断扩展w_x0, w_x1,直到不能继续向下扩展,或者到达边界下边界。
  259. """
  260. sorted_layout_blocks = [] # 这是要最终返回的值
  261. new_boundry = [boundry[0], boundry[1], boundry[2], boundry[3]]
  262. bad_boxes = [] # 被割中的box
  263. v_blocks = []
  264. while True:
  265. all_bboxes = get_bbox_in_boundry(bboxes, new_boundry)
  266. if len(all_bboxes) == 0:
  267. break
  268. left_top_box = min(all_bboxes, key=lambda x: (x[X0_IDX],x[Y0_IDX]))# 这里应该加强,检查一下必须是在第一列的 TODO
  269. start_box = [left_top_box[X0_IDX], left_top_box[Y0_IDX], left_top_box[X1_IDX], left_top_box[Y1_IDX]]
  270. w_x0, w_x1 = left_top_box[X0_IDX], left_top_box[X1_IDX]
  271. """
  272. 然后沿着这个box线向下找最近的那个box, 然后扩展w_x0, w_x1
  273. 扩展之后,宽度会增加,随后用x=w_x1来检测在边界内是否有box与相交,如果相交,那么就说明不能再扩展了。
  274. 当不能扩展的时候就要看是否到达下边界:
  275. 1. 达到,那么更新左边界继续分下一个列
  276. 2. 没有达到,那么此时开始从右侧切分进入下面的循环里
  277. """
  278. while left_top_box is not None: # 向下去找
  279. virtual_box = [w_x0, left_top_box[Y0_IDX], w_x1, left_top_box[Y1_IDX]]
  280. left_top_box = find_bottom_bbox_direct_from_left_edge(virtual_box, all_bboxes)
  281. if left_top_box:
  282. w_x0, w_x1 = min(virtual_box[X0_IDX], left_top_box[X0_IDX]), max([virtual_box[X1_IDX], left_top_box[X1_IDX]])
  283. # 万一这个初始的box在column中间,那么还要向上看
  284. start_box = [w_x0, start_box[Y0_IDX], w_x1, start_box[Y1_IDX]] # 扩展一下宽度更鲁棒
  285. left_top_box = find_top_bbox_direct_from_left_edge(start_box, all_bboxes)
  286. while left_top_box is not None: # 向上去找
  287. virtual_box = [w_x0, left_top_box[Y0_IDX], w_x1, left_top_box[Y1_IDX]]
  288. left_top_box = find_top_bbox_direct_from_left_edge(virtual_box, all_bboxes)
  289. if left_top_box:
  290. w_x0, w_x1 = min(virtual_box[X0_IDX], left_top_box[X0_IDX]), max([virtual_box[X1_IDX], left_top_box[X1_IDX]])
  291. # 检查相交
  292. if any([bbox[X0_IDX] <= w_x1+1 <= bbox[X1_IDX] for bbox in all_bboxes]):
  293. for b in all_bboxes:
  294. if b[X0_IDX] <= w_x1+1 <= b[X1_IDX]:
  295. bad_boxes.append([b[X0_IDX], b[Y0_IDX], b[X1_IDX], b[Y1_IDX]])
  296. break
  297. else: # 说明成功分割出一列
  298. v_blocks.append([w_x0, new_boundry[1], w_x1, new_boundry[3], LAYOUT_V])
  299. new_boundry[0] = w_x1 # 更新边界
  300. """
  301. 接着开始从右上角的box扫描
  302. """
  303. w_x0 , w_x1 = 0, 0
  304. unsplited_block = []
  305. while True:
  306. all_bboxes = get_bbox_in_boundry(bboxes, new_boundry)
  307. if len(all_bboxes) == 0:
  308. break
  309. # 先找到X1最大的
  310. bbox_list_sorted = sorted(all_bboxes, key=lambda bbox: bbox[X1_IDX], reverse=True)
  311. # Then, find the boxes with the smallest Y0 value
  312. bigest_x1 = bbox_list_sorted[0][X1_IDX]
  313. boxes_with_bigest_x1 = [bbox for bbox in bbox_list_sorted if bbox[X1_IDX] == bigest_x1] # 也就是最靠右的那些
  314. right_top_box = min(boxes_with_bigest_x1, key=lambda bbox: bbox[Y0_IDX]) # y0最小的那个
  315. start_box = [right_top_box[X0_IDX], right_top_box[Y0_IDX], right_top_box[X1_IDX], right_top_box[Y1_IDX]]
  316. w_x0, w_x1 = right_top_box[X0_IDX], right_top_box[X1_IDX]
  317. while right_top_box is not None:
  318. virtual_box = [w_x0, right_top_box[Y0_IDX], w_x1, right_top_box[Y1_IDX]]
  319. right_top_box = find_bottom_bbox_direct_from_right_edge(virtual_box, all_bboxes)
  320. if right_top_box:
  321. w_x0, w_x1 = min([w_x0, right_top_box[X0_IDX]]), max([w_x1, right_top_box[X1_IDX]])
  322. # 在向上扫描
  323. start_box = [w_x0, start_box[Y0_IDX], w_x1, start_box[Y1_IDX]] # 扩展一下宽度更鲁棒
  324. right_top_box = find_top_bbox_direct_from_right_edge(start_box, all_bboxes)
  325. while right_top_box is not None:
  326. virtual_box = [w_x0, right_top_box[Y0_IDX], w_x1, right_top_box[Y1_IDX]]
  327. right_top_box = find_top_bbox_direct_from_right_edge(virtual_box, all_bboxes)
  328. if right_top_box:
  329. w_x0, w_x1 = min([w_x0, right_top_box[X0_IDX]]), max([w_x1, right_top_box[X1_IDX]])
  330. # 检查是否与其他box相交, 垂直切分线与某些box发生相交,说明无法完全垂直方向切分。
  331. if any([bbox[X0_IDX] <= w_x0-1 <= bbox[X1_IDX] for bbox in all_bboxes]):
  332. unsplited_block.append([new_boundry[0], new_boundry[1], new_boundry[2], new_boundry[3], LAYOUT_UNPROC])
  333. for b in all_bboxes:
  334. if b[X0_IDX] <= w_x0-1 <= b[X1_IDX]:
  335. bad_boxes.append([b[X0_IDX], b[Y0_IDX], b[X1_IDX], b[Y1_IDX]])
  336. break
  337. else: # 说明成功分割出一列
  338. v_blocks.append([w_x0, new_boundry[1], w_x1, new_boundry[3], LAYOUT_V])
  339. new_boundry[2] = w_x0
  340. """转换数据结构"""
  341. for block in v_blocks:
  342. sorted_layout_blocks.append({
  343. "layout_bbox": block[:4],
  344. "layout_label":block[4],
  345. "sub_layout":[],
  346. })
  347. for block in unsplited_block:
  348. sorted_layout_blocks.append({
  349. "layout_bbox": block[:4],
  350. "layout_label":block[4],
  351. "sub_layout":[],
  352. "bad_boxes": bad_boxes # 记录下来,这个box是被割中的
  353. })
  354. # 按照x0排序
  355. sorted_layout_blocks.sort(key=lambda x: x['layout_bbox'][0])
  356. return sorted_layout_blocks
  357. def _try_horizontal_mult_column_split(bboxes:list, boundry:tuple)-> list:
  358. """
  359. 尝试水平切分,如果切分不动,那就当一个BAD_LAYOUT返回
  360. ------------------
  361. | | |
  362. ------------------
  363. | | | | <- 这里是此函数要切分的场景
  364. ------------------
  365. | | |
  366. | | |
  367. """
  368. pass
  369. def _vertical_split(bboxes:list, boundry:tuple)-> list:
  370. """
  371. 从垂直方向进行切割,分block
  372. 这个版本里,如果垂直切分不动,那就当一个BAD_LAYOUT返回
  373. --------------------------
  374. | | |
  375. | | |
  376. | |
  377. 这种列是此函数要切分的 -> | |
  378. | |
  379. | | |
  380. | | |
  381. -------------------------
  382. """
  383. sorted_layout_blocks = [] # 这是要最终返回的值
  384. bound_x0, bound_y0, bound_x1, bound_y1 = boundry
  385. all_bboxes = get_bbox_in_boundry(bboxes, boundry)
  386. """
  387. all_bboxes = fix_vertical_bbox_pos(all_bboxes) # 垂直方向解覆盖
  388. all_bboxes = fix_hor_bbox_pos(all_bboxes) # 水平解覆盖
  389. 这两行代码目前先不执行,因为公式检测,表格检测还不是很成熟,导致非常多的textblock参与了运算,时间消耗太大。
  390. 这两行代码的作用是:
  391. 如果遇到互相重叠的bbox, 那么会把面积较小的box进行压缩,从而避免重叠。对布局切分来说带来正反馈。
  392. """
  393. #all_bboxes = paper_bbox_sort(all_bboxes, abs(bound_x1-bound_x0), abs(bound_y1-bound_x0)) # 大致拍下序, 这个是基于直接遮挡的。
  394. """
  395. 首先在垂直方向上扩展独占一行的bbox
  396. """
  397. for bbox in all_bboxes:
  398. top_nearest_bbox = find_all_top_bbox_direct(bbox, all_bboxes) # 非扩展线
  399. bottom_nearest_bbox = find_all_bottom_bbox_direct(bbox, all_bboxes)
  400. if top_nearest_bbox is None and bottom_nearest_bbox is None and not any([b[X0_IDX]<bbox[X1_IDX]<b[X1_IDX] or b[X0_IDX]<bbox[X0_IDX]<b[X1_IDX] for b in all_bboxes]): # 独占一列, 且不和其他重叠
  401. bbox[X0_EXT_IDX] = bbox[X0_IDX]
  402. bbox[Y0_EXT_IDX] = bound_y0
  403. bbox[X1_EXT_IDX] = bbox[X1_IDX]
  404. bbox[Y1_EXT_IDX] = bound_y1
  405. """
  406. 此时独占一列的被成功扩展到指定的边界上,这个时候利用边界条件合并连续的bbox,成为一个group
  407. 然后合并所有连续垂直方向的bbox.
  408. """
  409. all_bboxes.sort(key=lambda x: x[X0_IDX])
  410. # fix: 这里水平方向的列不要合并成一个行,因为需要保证返回给下游的最小block,总是可以无脑从上到下阅读文字。
  411. v_bboxes = []
  412. for box in all_bboxes:
  413. if box[Y0_EXT_IDX] == bound_y0 and box[Y1_EXT_IDX] == bound_y1:
  414. v_bboxes.append(box)
  415. """
  416. 现在v_bboxes里面是所有的group了,每个group都是一个list
  417. 对v_bboxes里的每个group进行计算放回到sorted_layouts里
  418. """
  419. v_layouts = []
  420. for vbox in v_bboxes:
  421. #gp.sort(key=lambda x: x[X0_IDX])
  422. # 然后计算这个group的layout_bbox,也就是最小的x0,y0, 最大的x1,y1
  423. x0, y0, x1, y1 = vbox[X0_EXT_IDX], vbox[Y0_EXT_IDX], vbox[X1_EXT_IDX], vbox[Y1_EXT_IDX]
  424. v_layouts.append([x0, y0, x1, y1, LAYOUT_V]) # 垂直的布局
  425. """
  426. 接下来利用这些连续的垂直bbox的layout_bbox的x0, x1,从垂直上切分开其余的为几个部分
  427. """
  428. v_split_lines = [bound_x0]
  429. for gp in v_bboxes:
  430. x0, x1 = gp[X0_IDX], gp[X1_IDX]
  431. v_split_lines.append(x0)
  432. v_split_lines.append(x1)
  433. v_split_lines.append(bound_x1)
  434. unsplited_bboxes = []
  435. for i in range(0, len(v_split_lines), 2):
  436. start_x0, start_x1 = v_split_lines[i:i+2]
  437. # 然后找出[start_x0, start_x1]之间的其他bbox,这些组成一个未分割板块
  438. bboxes_in_block = [bbox for bbox in all_bboxes if bbox[X0_IDX]>=start_x0 and bbox[X1_IDX]<=start_x1]
  439. unsplited_bboxes.append(bboxes_in_block)
  440. # 接着把未处理的加入到v_layouts里
  441. for bboxes_in_block in unsplited_bboxes:
  442. if len(bboxes_in_block) == 0:
  443. continue
  444. x0, y0, x1, y1 = min([bbox[X0_IDX] for bbox in bboxes_in_block]), bound_y0, max([bbox[X1_IDX] for bbox in bboxes_in_block]), bound_y1
  445. v_layouts.append([x0, y0, x1, y1, LAYOUT_UNPROC]) # 说明这篇区域未能够分析出可靠的版面
  446. v_layouts.sort(key=lambda x: x[0]) # 按照x0排序, 也就是从左到右的顺序
  447. for layout in v_layouts:
  448. sorted_layout_blocks.append({
  449. "layout_bbox": layout[:4],
  450. "layout_label":layout[4],
  451. "sub_layout":[],
  452. })
  453. """
  454. 至此,垂直方向切成了2种类型,其一是独占一列的,其二是未处理的。
  455. 下面对这些未处理的进行垂直方向切分,这个切分要切出来类似“吕”这种类型的垂直方向的布局
  456. """
  457. for i, layout in enumerate(sorted_layout_blocks):
  458. if layout['layout_label'] == LAYOUT_UNPROC:
  459. x0, y0, x1, y1 = layout['layout_bbox']
  460. v_split_layouts = _vertical_align_split_v2(bboxes, [x0, y0, x1, y1])
  461. sorted_layout_blocks[i] = {
  462. "layout_bbox": [x0, y0, x1, y1],
  463. "layout_label": LAYOUT_H,
  464. "sub_layout": v_split_layouts
  465. }
  466. layout['layout_label'] = LAYOUT_H # 被垂线切分成了水平布局
  467. return sorted_layout_blocks
  468. def split_layout(bboxes:list, boundry:tuple, page_num:int)-> list:
  469. """
  470. 把bboxes切割成layout
  471. return:
  472. [
  473. {
  474. "layout_bbox": [x0, y0, x1, y1],
  475. "layout_label":"u|v|h|b", 未处理|垂直|水平|BAD_LAYOUT
  476. "sub_layout": [] #每个元素都是[x0, y0, x1, y1, block_content, idx_x, idx_y, content_type, ext_x0, ext_y0, ext_x1, ext_y1], 并且顺序就是阅读顺序
  477. }
  478. ]
  479. example:
  480. [
  481. {
  482. "layout_bbox": [0, 0, 100, 100],
  483. "layout_label":"u|v|h|b",
  484. "sub_layout":[
  485. ]
  486. },
  487. {
  488. "layout_bbox": [0, 0, 100, 100],
  489. "layout_label":"u|v|h|b",
  490. "sub_layout":[
  491. {
  492. "layout_bbox": [0, 0, 100, 100],
  493. "layout_label":"u|v|h|b",
  494. "content_bboxes":[
  495. [],
  496. [],
  497. []
  498. ]
  499. },
  500. {
  501. "layout_bbox": [0, 0, 100, 100],
  502. "layout_label":"u|v|h|b",
  503. "sub_layout":[
  504. ]
  505. }
  506. }
  507. ]
  508. """
  509. sorted_layouts = [] # 最终返回的结果
  510. boundry_x0, boundry_y0, boundry_x1, boundry_y1 = boundry
  511. if len(bboxes) <=1:
  512. return [
  513. {
  514. "layout_bbox": [boundry_x0, boundry_y0, boundry_x1, boundry_y1],
  515. "layout_label": LAYOUT_V,
  516. "sub_layout":[]
  517. }
  518. ]
  519. """
  520. 接下来按照先水平后垂直的顺序进行切分
  521. """
  522. bboxes = paper_bbox_sort(bboxes, boundry_x1-boundry_x0, boundry_y1-boundry_y0)
  523. sorted_layouts = _horizontal_split(bboxes, boundry) # 通过水平分割出来的layout
  524. for i, layout in enumerate(sorted_layouts):
  525. x0, y0, x1, y1 = layout['layout_bbox']
  526. layout_type = layout['layout_label']
  527. if layout_type == LAYOUT_UNPROC: # 说明是非独占单行的,这些需要垂直切分
  528. v_split_layouts = _vertical_split(bboxes, [x0, y0, x1, y1])
  529. """
  530. 最后这里有个逻辑问题:如果这个函数只分离出来了一个column layout,那么这个layout分割肯定超出了算法能力范围。因为我们假定的是传进来的
  531. box已经把行全部剥离了,所以这里必须十多个列才可以。如果只剥离出来一个layout,并且是多个box,那么就说明这个layout是无法分割的,标记为LAYOUT_UNPROC
  532. """
  533. layout_label = LAYOUT_V
  534. if len(v_split_layouts) == 1:
  535. if len(v_split_layouts[0]['sub_layout']) == 0:
  536. layout_label = LAYOUT_UNPROC
  537. #logger.warning(f"WARNING: pageno={page_num}, 无法分割的layout: ", v_split_layouts)
  538. """
  539. 组合起来最终的layout
  540. """
  541. sorted_layouts[i] = {
  542. "layout_bbox": [x0, y0, x1, y1],
  543. "layout_label": layout_label,
  544. "sub_layout": v_split_layouts
  545. }
  546. layout['layout_label'] = LAYOUT_H
  547. """
  548. 水平和垂直方向都切分完毕了。此时还有一些未处理的,这些未处理的可能是因为水平和垂直方向都无法切分。
  549. 这些最后调用_try_horizontal_mult_block_split做一次水平多个block的联合切分,如果也不能切分最终就当做BAD_LAYOUT返回
  550. """
  551. # TODO
  552. return sorted_layouts
  553. def get_bboxes_layout(all_boxes:list, boundry:tuple, page_id:int):
  554. """
  555. 对利用layout排序之后的box,进行排序
  556. return:
  557. [
  558. {
  559. "layout_bbox": [x0, y0, x1, y1],
  560. "layout_label":"u|v|h|b", 未处理|垂直|水平|BAD_LAYOUT
  561. },
  562. ]
  563. """
  564. def _preorder_traversal(layout):
  565. """
  566. 对sorted_layouts的叶子节点,也就是len(sub_layout)==0的节点进行排序。排序按照前序遍历的顺序,也就是从上到下,从左到右的顺序
  567. """
  568. sorted_layout_blocks = []
  569. for layout in layout:
  570. sub_layout = layout['sub_layout']
  571. if len(sub_layout) == 0:
  572. sorted_layout_blocks.append(layout)
  573. else:
  574. s = _preorder_traversal(sub_layout)
  575. sorted_layout_blocks.extend(s)
  576. return sorted_layout_blocks
  577. # -------------------------------------------------------------------------------------------------------------------------
  578. sorted_layouts = split_layout(all_boxes, boundry, page_id)# 先切分成layout,得到一个Tree
  579. total_sorted_layout_blocks = _preorder_traversal(sorted_layouts)
  580. return total_sorted_layout_blocks, sorted_layouts
  581. def get_columns_cnt_of_layout(layout_tree):
  582. """
  583. 获取一个layout的宽度
  584. """
  585. max_width_list = [0] # 初始化一个元素,防止max,min函数报错
  586. for items in layout_tree: # 针对每一层(横切)计算列数,横着的算一列
  587. layout_type = items['layout_label']
  588. sub_layouts = items['sub_layout']
  589. if len(sub_layouts)==0:
  590. max_width_list.append(1)
  591. else:
  592. if layout_type == LAYOUT_H:
  593. max_width_list.append(1)
  594. else:
  595. width = 0
  596. for l in sub_layouts:
  597. if len(l['sub_layout']) == 0:
  598. width += 1
  599. else:
  600. for lay in l['sub_layout']:
  601. width += get_columns_cnt_of_layout([lay])
  602. max_width_list.append(width)
  603. return max(max_width_list)
  604. def sort_with_layout(bboxes:list, page_width, page_height) -> (list,list):
  605. """
  606. 输入是一个bbox的list.
  607. 获取到输入之后,先进行layout切分,然后对这些bbox进行排序。返回排序后的bboxes
  608. """
  609. new_bboxes = []
  610. for box in bboxes:
  611. # new_bboxes.append([box[0], box[1], box[2], box[3], None, None, None, 'text', None, None, None, None])
  612. new_bboxes.append([box[0], box[1], box[2], box[3], None, None, None, 'text', None, None, None, None, box[4]])
  613. layout_bboxes, _ = get_bboxes_layout(new_bboxes, [0, 0, page_width, page_height], 0)
  614. if any([lay['layout_label']==LAYOUT_UNPROC for lay in layout_bboxes]):
  615. logger.warning(f"drop this pdf, reason: 复杂版面")
  616. return None,None
  617. sorted_bboxes = []
  618. # 利用layout bbox每次框定一些box,然后排序
  619. for layout in layout_bboxes:
  620. lbox = layout['layout_bbox']
  621. bbox_in_layout = get_bbox_in_boundry(new_bboxes, lbox)
  622. sorted_bbox = paper_bbox_sort(bbox_in_layout, lbox[2]-lbox[0], lbox[3]-lbox[1])
  623. sorted_bboxes.extend(sorted_bbox)
  624. return sorted_bboxes, layout_bboxes
  625. def sort_text_block(text_block, layout_bboxes):
  626. """
  627. 对一页的text_block进行排序
  628. """
  629. sorted_text_bbox = []
  630. all_text_bbox = []
  631. # 做一个box=>text的映射
  632. box_to_text = {}
  633. for blk in text_block:
  634. box = blk['bbox']
  635. box_to_text[(box[0], box[1], box[2], box[3])] = blk
  636. all_text_bbox.append(box)
  637. # text_blocks_to_sort = []
  638. # for box in box_to_text.keys():
  639. # text_blocks_to_sort.append([box[0], box[1], box[2], box[3], None, None, None, 'text', None, None, None, None])
  640. # 按照layout_bboxes的顺序,对text_block进行排序
  641. for layout in layout_bboxes:
  642. layout_box = layout['layout_bbox']
  643. text_bbox_in_layout = get_bbox_in_boundry(all_text_bbox, [layout_box[0]-1, layout_box[1]-1, layout_box[2]+1, layout_box[3]+1])
  644. #sorted_bbox = paper_bbox_sort(text_bbox_in_layout, layout_box[2]-layout_box[0], layout_box[3]-layout_box[1])
  645. text_bbox_in_layout.sort(key = lambda x: x[1]) # 一个layout内部的box,按照y0自上而下排序
  646. #sorted_bbox = [[b] for b in text_blocks_to_sort]
  647. for sb in text_bbox_in_layout:
  648. sorted_text_bbox.append(box_to_text[(sb[0], sb[1], sb[2], sb[3])])
  649. return sorted_text_bbox