Skip to content

Commit

Permalink
refactor(para_split_v3): merge list and index block detection
Browse files Browse the repository at this point in the history
- Combine __is_list_block() and __is_index_block() into a single function __is_list_or_index_block()
- Simplify block type determination logic
- Remove redundant code and improve readability
- Optimize block merging process
  • Loading branch information
myhloli committed Oct 15, 2024
1 parent 0a9a6d3 commit fdcb49d
Showing 1 changed file with 43 additions and 71 deletions.
114 changes: 43 additions & 71 deletions magic_pdf/para/para_split_v3.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,9 @@ class ListLineTag:


def __process_blocks(blocks):
# 对所有block预处理
# 1.通过title和interline_equation将block分组
# 2.bbox边界根据line信息重置

result = []
current_group = []
Expand Down Expand Up @@ -47,11 +50,15 @@ def __process_blocks(blocks):
return result


def __is_list_block(block):
def __is_list_or_index_block(block):
# 一个block如果是list block 应该同时满足以下特征
# 1.block内有多个line 2.block 内有多个line左侧顶格写 3.block内有多个line 右侧不顶格(狗牙状)
# 1.block内有多个line 2.block 内有多个line左侧顶格写 3.多个line以endflag结尾
# 1.block内有多个line 2.block 内有多个line左侧顶格写 3.block内有多个line 左侧不顶格

# index block 是一种特殊的list block
# 一个block如果是index block 应该同时满足以下特征
# 1.block内有多个line 2.block 内有多个line两侧均顶格写 3.line的开头或者结尾均为数字
if len(block['lines']) >= 3:
first_line = block['lines'][0]
line_height = first_line['bbox'][3] - first_line['bbox'][1]
Expand All @@ -60,6 +67,7 @@ def __is_list_block(block):
left_close_num = 0
left_not_close_num = 0
right_not_close_num = 0
right_close_num = 0
lines_text_list = []
for line in block['lines']:

Expand All @@ -79,92 +87,55 @@ def __is_list_block(block):
# logger.info(f"{line_text}, {block['bbox_fs']}, {line['bbox']}")
left_not_close_num += 1

# 计算右侧是否不顶格,拍脑袋用0.3block宽度做阈值
closed_area = 0.3 * block_weight
# closed_area = 5 * line_height
if block['bbox_fs'][2] - line['bbox'][2] > closed_area:
right_not_close_num += 1

# 判断lines_text_list中的元素是否有超过80%都以LIST_END_FLAG结尾
line_end_flag = False
if len(lines_text_list) > 0:
num_end_count = 0
for line_text in lines_text_list:
if len(line_text) > 0:
if line_text[-1] in LIST_END_FLAG:
num_end_count += 1

if num_end_count / len(lines_text_list) >= 0.8:
line_end_flag = True

if left_close_num >= 2 and (right_not_close_num >= 2 or line_end_flag or left_not_close_num >= 2):
for line in block['lines']:
if abs(block['bbox_fs'][0] - line['bbox'][0]) < line_height / 2:
line[ListLineTag.IS_LIST_START_LINE] = True
if abs(block['bbox_fs'][2] - line['bbox'][2]) > line_height:
line[ListLineTag.IS_LIST_END_LINE] = True

return True
else:
return False
else:
return False


def __is_index_block(block):
# 一个block如果是index block 应该同时满足以下特征
# 1.block内有多个line 2.block 内有多个line两侧均顶格写 3.line的开头或者结尾均为数字
if len(block['lines']) >= 3:
first_line = block['lines'][0]
line_height = first_line['bbox'][3] - first_line['bbox'][1]

left_close_num = 0
right_close_num = 0

lines_text_list = []
for line in block['lines']:

# 计算line左侧顶格数量是否大于2,是否顶格用abs(block['bbox_fs'][0] - line['bbox'][0]) < line_height/2 来判断
if abs(block['bbox_fs'][0] - line['bbox'][0]) < line_height / 2:
left_close_num += 1

# 计算右侧是否不顶格
# 计算右侧是否顶格
if abs(block['bbox_fs'][2] - line['bbox'][2]) < line_height / 2:
right_close_num += 1
else:
# 右侧不顶格情况下是否有一段距离,拍脑袋用0.3block宽度做阈值
closed_area = 0.3 * block_weight
# closed_area = 5 * line_height
if block['bbox_fs'][2] - line['bbox'][2] > closed_area:
right_not_close_num += 1

line_text = ""

for span in line['spans']:
span_type = span['type']
if span_type == ContentType.Text:
line_text += span['content'].strip()

lines_text_list.append(line_text)

# 判断lines_text_list中的元素是否有超过80%都以LIST_END_FLAG结尾
line_end_flag = False
# 判断lines_text_list中的元素是否有超过80%都以数字开头或都以数字结尾
line_num_flag = False
if len(lines_text_list) > 0:
num_start_count = 0
num_end_count = 0
for line_text in lines_text_list:
if len(line_text) > 0:
if line_text[-1] in LIST_END_FLAG:
num_end_count += 1
if line_text[0].isdigit():
num_start_count += 1
if line_text[-1].isdigit():
num_end_count += 1

if num_end_count / len(lines_text_list) >= 0.8:
line_end_flag = True

if num_start_count / len(lines_text_list) >= 0.8 or num_end_count / len(lines_text_list) >= 0.8:
line_num_flag = True

if left_close_num >= 2 and right_close_num >= 2 and line_num_flag:
for line in block['lines']:
line[ListLineTag.IS_LIST_START_LINE] = True
return BlockType.Index

return True
elif left_close_num >= 2 and (right_not_close_num >= 2 or line_end_flag or left_not_close_num > 2):
for line in block['lines']:
if abs(block['bbox_fs'][0] - line['bbox'][0]) < line_height / 2:
line[ListLineTag.IS_LIST_START_LINE] = True
if abs(block['bbox_fs'][2] - line['bbox'][2]) > line_height:
line[ListLineTag.IS_LIST_END_LINE] = True

return BlockType.List
else:
return False
return BlockType.Text
else:
return False
return BlockType.Text


def __merge_2_text_blocks(block1, block2):
Expand Down Expand Up @@ -206,12 +177,11 @@ def __para_merge_page(blocks):
for text_blocks_group in page_text_blocks_groups:

if len(text_blocks_group) > 0:
# 需要先在合并前对所有block判断是否为list block
# 需要先在合并前对所有block判断是否为list or index block
for block in text_blocks_group:
if __is_list_block(block):
block['type'] = BlockType.List
elif __is_index_block(block):
block['type'] = BlockType.Index
block_type = __is_list_or_index_block(block)
block['type'] = block_type
# logger.info(f"{block['type']}:{block}")

if len(text_blocks_group) > 1:
# 倒序遍历
Expand All @@ -224,10 +194,12 @@ def __para_merge_page(blocks):

if current_block['type'] == 'text' and prev_block['type'] == 'text':
__merge_2_text_blocks(current_block, prev_block)
if current_block['type'] == BlockType.List and prev_block['type'] == BlockType.List:
__merge_2_list_blocks(current_block, prev_block)
if current_block['type'] == BlockType.Index and prev_block['type'] == BlockType.Index:
elif (
(current_block['type'] == BlockType.List and prev_block['type'] == BlockType.List) or
(current_block['type'] == BlockType.Index and prev_block['type'] == BlockType.Index)
):
__merge_2_list_blocks(current_block, prev_block)

else:
continue

Expand Down

0 comments on commit fdcb49d

Please sign in to comment.