Skip to content

GitLab

Unverified Commit a7e615ff authored by Kaiwen Liu's avatar Kaiwen Liu Committed by GitHub
Browse files

fix: resolve inaccuracy of drawing layout box caused by paragraphs combination #384 (#542) 

* fix: resolve inaccuracy of drawing layout box caused by paragraphs combination

* fix: resolve inaccuracy of drawing layout box caused by paragraphs combination #384

* fix: resolve inaccuracy of drawing layout box caused by paragraphs combination #384

* fix: resolve inaccuracy of drawing layout box caused by paragraphs combination #384
parent 244d122c
Hide whitespace changes
Inline Side-by-side
Showing with 25 additions and 26 deletions
magic_pdf/para/para_split_v2.py
View file @ a7e615ff
import copy
from sklearn.cluster import DBSCAN
import numpy as np
from loguru import logger
......@@ -167,7 +169,7 @@ def cluster_line_x(lines: list) -> dict:
x0_lst = np.array([[round(line['bbox'][0]), 0] for line in lines])
x0_clusters = DBSCAN(eps=min_distance, min_samples=min_sample).fit(x0_lst)
x0_uniq_label = np.unique(x0_clusters.labels_)
#x1_lst = np.array([[line['bbox'][2], 0] for line in lines])
# x1_lst = np.array([[line['bbox'][2], 0] for line in lines])
x0_2_new_val = {} # 存储旧值对应的新值映射
min_x0 = round(lines[0]["bbox"][0])
for label in x0_uniq_label:
......@@ -245,16 +247,17 @@ def __valign_lines(blocks, layout_bboxes):
# 由于修改了block里的line长度，现在需要重新计算block的bbox
for block in blocks_in_layoutbox:
if len(block["lines"]) > 0:
block['bbox'] = [min([line['bbox'][0] for line in block['lines']]),
min([line['bbox'][1] for line in block['lines']]),
max([line['bbox'][2] for line in block['lines']]),
max([line['bbox'][3] for line in block['lines']])]
block['bbox_fs'] = [min([line['bbox'][0] for line in block['lines']]),
min([line['bbox'][1] for line in block['lines']]),
max([line['bbox'][2] for line in block['lines']]),
max([line['bbox'][3] for line in block['lines']])]
else:
block['bbox_fs'] = copy.deepcopy(block['bbox'])
"""新计算layout的bbox，因为block的bbox变了。"""
layout_x0 = min([block['bbox'][0] for block in blocks_in_layoutbox])
layout_y0 = min([block['bbox'][1] for block in blocks_in_layoutbox])
layout_x1 = max([block['bbox'][2] for block in blocks_in_layoutbox])
layout_y1 = max([block['bbox'][3] for block in blocks_in_layoutbox])
layout_x0 = min([block['bbox_fs'][0] for block in blocks_in_layoutbox])
layout_y0 = min([block['bbox_fs'][1] for block in blocks_in_layoutbox])
layout_x1 = max([block['bbox_fs'][2] for block in blocks_in_layoutbox])
layout_y1 = max([block['bbox_fs'][3] for block in blocks_in_layoutbox])
new_layout_bboxes.append([layout_x0, layout_y0, layout_x1, layout_y1])
return new_layout_bboxes
......@@ -312,7 +315,7 @@ def __group_line_by_layout(blocks, layout_bboxes):
# 因为只是一个block一行目前, 一个block就是一个段落
blocks_group = []
for lyout in layout_bboxes:
blocks_in_layout = [block for block in blocks if is_in_layout(block['bbox'], lyout['layout_bbox'])]
blocks_in_layout = [block for block in blocks if is_in_layout(block.get('bbox_fs', None), lyout['layout_bbox'])]
blocks_group.append(blocks_in_layout)
return blocks_group
......@@ -365,7 +368,8 @@ def __split_para_in_layoutbox(blocks_group, new_layout_bbox, lang="en"):
for i in range(0, len(list_start)):
index = list_start[i] - 1
if index >= 0:
if "content" in lines[index]["spans"][-1]:
if "content" in lines[index]["spans"][-1] and lines[index]["spans"][-1].get('type', '') not in [
ContentType.InlineEquation, ContentType.InterlineEquation]:
lines[index]["spans"][-1]["content"] += '\n\n'
layout_list_info = [False, False] # 这个layout最后是不是列表,记录每一个layout里是不是列表开头，列表结尾
for content_type, start, end in text_segments:
......@@ -477,7 +481,7 @@ def __connect_list_inter_page(pre_page_paras, next_page_paras, pre_page_layout_b
break
# 如果这些行的缩进是相等的，那么连到上一个layout的最后一个段落上。
if len(may_list_lines) > 0 and len(set([x['bbox'][0] for x in may_list_lines])) == 1:
#pre_page_paras[-1].append(may_list_lines)
# pre_page_paras[-1].append(may_list_lines)
# 下一页合并到上一页最后一段，打一个cross_page的标签
for line in may_list_lines:
for span in line["spans"]:
......@@ -537,7 +541,6 @@ def __connect_para_inter_layoutbox(blocks_group, new_layout_bbox):
next_first_line_text = ''.join([__get_span_text(span) for span in next_first_line['spans']])
next_first_line_type = next_first_line['spans'][0]['type']
if pre_last_line_type not in [TEXT, INLINE_EQUATION] or next_first_line_type not in [TEXT, INLINE_EQUATION]:
#connected_layout_paras.append(layout_paras[i])
connected_layout_blocks.append(blocks_group[i])
continue
pre_layout = __find_layout_bbox_by_line(pre_last_line['bbox'], new_layout_bbox)
......@@ -552,10 +555,8 @@ def __connect_para_inter_layoutbox(blocks_group, new_layout_bbox):
-1] not in LINE_STOP_FLAG and \
next_first_line['bbox'][0] == next_x0_min: # 前面一行沾满了整个行，并且没有结尾符号.下一行没有空白开头。
"""连接段落条件成立，将前一个layout的段落和后一个layout的段落连接。"""
#connected_layout_paras[-1][-1].extend(layout_paras[i][0])
connected_layout_blocks[-1][-1]["lines"].extend(blocks_group[i][0]["lines"])
#layout_paras[i].pop(0) # 删除后一个layout的第一个段落， 因为他已经被合并到前一个layout的最后一个段落了。
blocks_group[i][0]["lines"] = [] #删除后一个layout第一个段落中的lines，因为他已经被合并到前一个layout的最后一个段落了
blocks_group[i][0]["lines"] = [] # 删除后一个layout第一个段落中的lines，因为他已经被合并到前一个layout的最后一个段落了
blocks_group[i][0][LINES_DELETED] = True
# if len(layout_paras[i]) == 0:
# layout_paras.pop(i)
......@@ -564,7 +565,6 @@ def __connect_para_inter_layoutbox(blocks_group, new_layout_bbox):
connected_layout_blocks.append(blocks_group[i])
else:
"""连接段落条件不成立，将前一个layout的段落加入到结果中。"""
#connected_layout_paras.append(layout_paras[i])
connected_layout_blocks.append(blocks_group[i])
return connected_layout_blocks
......@@ -622,7 +622,7 @@ def __connect_para_inter_page(pre_page_paras, next_page_paras, pre_page_layout_b
span[CROSS_PAGE] = True
pre_last_para.extend(next_first_para)
#next_page_paras[0].pop(0) # 删除后一个页面的第一个段落， 因为他已经被合并到前一个页面的最后一个段落了。
# next_page_paras[0].pop(0) # 删除后一个页面的第一个段落， 因为他已经被合并到前一个页面的最后一个段落了。
next_page_paras[0][0]["lines"] = []
next_page_paras[0][0][LINES_DELETED] = True
return True
......@@ -666,16 +666,15 @@ def __connect_middle_align_text(page_paras, new_layout_bbox, page_num, lang):
layout_box = new_layout_bbox[layout_i]
single_line_paras_tag = []
for i in range(len(layout_para)):
#single_line_paras_tag.append(len(layout_para[i]) == 1 and layout_para[i][0]['spans'][0]['type'] == TEXT)
# single_line_paras_tag.append(len(layout_para[i]) == 1 and layout_para[i][0]['spans'][0]['type'] == TEXT)
single_line_paras_tag.append(layout_para[i]['type'] == BlockType.Text and len(layout_para[i]["lines"]) == 1)
"""找出来连续的单行文本，如果连续行高度相同，那么合并为一个段落。"""
consecutive_single_line_indices = find_consecutive_true_regions(single_line_paras_tag)
if len(consecutive_single_line_indices) > 0:
#index_offset = 0
"""检查这些行是否是高度相同的，居中的"""
for start, end in consecutive_single_line_indices:
#start += index_offset
#end += index_offset
# start += index_offset
# end += index_offset
line_hi = np.array([block["lines"][0]['bbox'][3] - block["lines"][0]['bbox'][1] for block in
layout_para[start:end + 1]])
first_line_text = ''.join([__get_span_text(span) for span in layout_para[start]["lines"][0]['spans']])
......@@ -700,9 +699,9 @@ def __connect_middle_align_text(page_paras, new_layout_bbox, page_num, lang):
for i_para in range(start + 1, end + 1):
layout_para[i_para]["lines"] = []
layout_para[i_para][LINES_DELETED] = True
#layout_para[start:end + 1] = [merge_para]
# layout_para[start:end + 1] = [merge_para]
#index_offset -= end - start
# index_offset -= end - start
return
......@@ -742,7 +741,7 @@ def para_split(pdf_info_dict, debug_mode, lang="en"):
new_layout_of_pages = [] # 数组的数组，每个元素是一个页面的layoutS
all_page_list_info = [] # 保存每个页面开头和结尾是否是列表
for page_num, page in pdf_info_dict.items():
blocks = page['preproc_blocks']
blocks = copy.deepcopy(page['preproc_blocks'])
layout_bboxes = page['layout_bboxes']
new_layout_bbox = __common_pre_proc(blocks, layout_bboxes)
new_layout_of_pages.append(new_layout_bbox)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment