Skip to content

GitLab

"lib/template.py" did not exist on "3ca8f7232722b0936d437daa1835be1d842b8e73"
Unverified Commit e6da37dd authored by Xiaomeng Zhao's avatar Xiaomeng Zhao Committed by GitHub
Browse files

Merge pull request #1099 from myhloli/dev 

refactor(magic_pdf): remove unused functions and simplify code
parents 79b58a1e 6a22b5ab
Hide whitespace changes
Inline Side-by-side
Showing with 0 additions and 8847 deletions
magic_pdf/para/para_split_v2.py.bak deleted 100644 → 0
View file @ 79b58a1e
import copy
import re
import numpy as np
from loguru import logger
from sklearn.cluster import DBSCAN
from magic_pdf.config.constants import * # noqa: F403
from magic_pdf.config.ocr_content_type import BlockType, ContentType
from magic_pdf.libs.boxbase import \
_is_in_or_part_overlap_with_area_ratio as is_in_layout
LINE_STOP_FLAG = ['.', '!', '?', '。', '！', '？', '：', ':', ')', '）', ';']
INLINE_EQUATION = ContentType.InlineEquation
INTERLINE_EQUATION = ContentType.InterlineEquation
TEXT = ContentType.Text
debug_able = False
def __get_span_text(span):
c = span.get('content', '')
if len(c) == 0:
c = span.get('image_path', '')
return c
def __detect_list_lines(lines, new_layout_bboxes, lang):
global debug_able
"""
探测是否包含了列表，并且把列表的行分开.
这样的段落特点是，顶格字母大写/数字，紧跟着几行缩进的。缩进的行首字母含小写的。
"""
def find_repeating_patterns2(lst):
indices = []
ones_indices = []
i = 0
while i < len(lst): # Loop through the entire list
if (
lst[i] == 1
): # If we encounter a '1', we might be at the start of a pattern
start = i
ones_in_this_interval = [i]
i += 1
# Traverse elements that are 1, 2 or 3, until we encounter something else
while i < len(lst) and lst[i] in [1, 2, 3]:
if lst[i] == 1:
ones_in_this_interval.append(i)
i += 1
if len(ones_in_this_interval) > 1 or (
start < len(lst) - 1
and ones_in_this_interval
and lst[start + 1] in [2, 3]
):
indices.append((start, i - 1))
ones_indices.append(ones_in_this_interval)
else:
i += 1
return indices, ones_indices
def find_repeating_patterns(lst):
indices = []
ones_indices = []
i = 0
while i < len(lst) - 1: # 确保余下元素至少有2个
if lst[i] == 1 and lst[i + 1] in [2, 3]: # 额外检查以防止连续出现的1
start = i
ones_in_this_interval = [i]
i += 1
while i < len(lst) and lst[i] in [2, 3]:
i += 1
# 验证下一个序列是否符合条件
if (
i < len(lst) - 1
and lst[i] == 1
and lst[i + 1] in [2, 3]
and lst[i - 1] in [2, 3]
):
while i < len(lst) and lst[i] in [1, 2, 3]:
if lst[i] == 1:
ones_in_this_interval.append(i)
i += 1
indices.append((start, i - 1))
ones_indices.append(ones_in_this_interval)
else:
i += 1
else:
i += 1
return indices, ones_indices
"""===================="""
def split_indices(slen, index_array):
result = []
last_end = 0
for start, end in sorted(index_array):
if start > last_end:
# 前一个区间结束到下一个区间开始之间的部分标记为"text"
result.append(('text', last_end, start - 1))
# 区间内标记为"list"
result.append(('list', start, end))
last_end = end + 1
if last_end < slen:
# 如果最后一个区间结束后还有剩余的字符串，将其标记为"text"
result.append(('text', last_end, slen - 1))
return result
"""===================="""
if lang != 'en':
return lines, None
total_lines = len(lines)
line_fea_encode = []
"""
对每一行进行特征编码，编码规则如下：
1. 如果行顶格，且大写字母开头或者数字开头，编码为1
2. 如果顶格，其他非大写开头编码为4
3. 如果非顶格，首字符大写，编码为2
4. 如果非顶格，首字符非大写编码为3
"""
if len(lines) > 0:
x_map_tag_dict, min_x_tag = cluster_line_x(lines)
for l in lines: # noqa: E741
span_text = __get_span_text(l['spans'][0])
if not span_text:
line_fea_encode.append(0)
continue
first_char = span_text[0]
layout = __find_layout_bbox_by_line(l['bbox'], new_layout_bboxes)
if not layout:
line_fea_encode.append(0)
else:
#
if x_map_tag_dict[round(l['bbox'][0])] == min_x_tag:
# if first_char.isupper() or first_char.isdigit() or not first_char.isalnum():
if not first_char.isalnum() or if_match_reference_list(span_text):
line_fea_encode.append(1)
else:
line_fea_encode.append(4)
else:
if first_char.isupper():
line_fea_encode.append(2)
else:
line_fea_encode.append(3)
# 然后根据编码进行分段, 选出来 1,2,3连续出现至少2次的行，认为是列表。
list_indice, list_start_idx = find_repeating_patterns2(line_fea_encode)
if len(list_indice) > 0:
if debug_able:
logger.info(f'发现了列表，列表行数：{list_indice}， {list_start_idx}')
# TODO check一下这个特列表里缩进的行左侧是不是对齐的。
for start, end in list_indice:
for i in range(start, end + 1):
if i > 0:
if line_fea_encode[i] == 4:
if debug_able:
logger.info(f'列表行的第{i}行不是顶格的')
break
else:
if debug_able:
logger.info(f'列表行的第{start}到第{end}行是列表')
return split_indices(total_lines, list_indice), list_start_idx
def cluster_line_x(lines: list) -> dict:
"""对一个block内所有lines的bbox的x0聚类."""
min_distance = 5
min_sample = 1
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])
x0_2_new_val = {} # 存储旧值对应的新值映射
min_x0 = round(lines[0]['bbox'][0])
for label in x0_uniq_label:
if label == -1:
continue
x0_index_of_label = np.where(x0_clusters.labels_ == label)
x0_raw_val = x0_lst[x0_index_of_label][:, 0]
x0_new_val = np.min(x0_lst[x0_index_of_label][:, 0])
x0_2_new_val.update(
{round(raw_val): round(x0_new_val) for raw_val in x0_raw_val}
)
if x0_new_val < min_x0:
min_x0 = x0_new_val
return x0_2_new_val, min_x0
def if_match_reference_list(text: str) -> bool:
pattern = re.compile(r'^\d+\..*')
if pattern.match(text):
return True
else:
return False
def __valign_lines(blocks, layout_bboxes):
"""在一个layoutbox内对齐行的左侧和右侧。 扫描行的左侧和右侧，如果x0,
x1差距不超过一个阈值，就强行对齐到所处layout的左右两侧（和layout有一段距离）。
3是个经验值，TODO，计算得来，可以设置为1.5个正文字符。"""
min_distance = 3
min_sample = 2
new_layout_bboxes = []
# add bbox_fs for para split calculation
for block in blocks:
block['bbox_fs'] = copy.deepcopy(block['bbox'])
for layout_box in layout_bboxes:
blocks_in_layoutbox = [
b
for b in blocks
if b['type'] == BlockType.Text
and is_in_layout(b['bbox'], layout_box['layout_bbox'])
]
if len(blocks_in_layoutbox) == 0 or len(blocks_in_layoutbox[0]['lines']) == 0:
new_layout_bboxes.append(layout_box['layout_bbox'])
continue
x0_lst = np.array(
[
[line['bbox'][0], 0]
for block in blocks_in_layoutbox
for line in block['lines']
]
)
x1_lst = np.array(
[
[line['bbox'][2], 0]
for block in blocks_in_layoutbox
for line in block['lines']
]
)
x0_clusters = DBSCAN(eps=min_distance, min_samples=min_sample).fit(x0_lst)
x1_clusters = DBSCAN(eps=min_distance, min_samples=min_sample).fit(x1_lst)
x0_uniq_label = np.unique(x0_clusters.labels_)
x1_uniq_label = np.unique(x1_clusters.labels_)
x0_2_new_val = {} # 存储旧值对应的新值映射
x1_2_new_val = {}
for label in x0_uniq_label:
if label == -1:
continue
x0_index_of_label = np.where(x0_clusters.labels_ == label)
x0_raw_val = x0_lst[x0_index_of_label][:, 0]
x0_new_val = np.min(x0_lst[x0_index_of_label][:, 0])
x0_2_new_val.update({idx: x0_new_val for idx in x0_raw_val})
for label in x1_uniq_label:
if label == -1:
continue
x1_index_of_label = np.where(x1_clusters.labels_ == label)
x1_raw_val = x1_lst[x1_index_of_label][:, 0]
x1_new_val = np.max(x1_lst[x1_index_of_label][:, 0])
x1_2_new_val.update({idx: x1_new_val for idx in x1_raw_val})
for block in blocks_in_layoutbox:
for line in block['lines']:
x0, x1 = line['bbox'][0], line['bbox'][2]
if x0 in x0_2_new_val:
line['bbox'][0] = int(x0_2_new_val[x0])
if x1 in x1_2_new_val:
line['bbox'][2] = int(x1_2_new_val[x1])
# 其余对不齐的保持不动
# 由于修改了block里的line长度，现在需要重新计算block的bbox
for block in blocks_in_layoutbox:
if len(block['lines']) > 0:
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']]),
]
"""新计算layout的bbox，因为block的bbox变了。"""
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
def __align_text_in_layout(blocks, layout_bboxes):
"""由于ocr出来的line，有时候会在前后有一段空白，这个时候需要对文本进行对齐，超出的部分被layout左右侧截断。"""
for layout in layout_bboxes:
lb = layout['layout_bbox']
blocks_in_layoutbox = [
block
for block in blocks
if block['type'] == BlockType.Text and is_in_layout(block['bbox'], lb)
]
if len(blocks_in_layoutbox) == 0:
continue
for block in blocks_in_layoutbox:
for line in block.get('lines', []):
x0, x1 = line['bbox'][0], line['bbox'][2]
if x0 < lb[0]:
line['bbox'][0] = lb[0]
if x1 > lb[2]:
line['bbox'][2] = lb[2]
def __common_pre_proc(blocks, layout_bboxes):
"""不分语言的，对文本进行预处理."""
# __add_line_period(blocks, layout_bboxes)
__align_text_in_layout(blocks, layout_bboxes)
aligned_layout_bboxes = __valign_lines(blocks, layout_bboxes)
return aligned_layout_bboxes
def __pre_proc_zh_blocks(blocks, layout_bboxes):
"""对中文文本进行分段预处理."""
pass
def __pre_proc_en_blocks(blocks, layout_bboxes):
"""对英文文本进行分段预处理."""
pass
def __group_line_by_layout(blocks, layout_bboxes):
"""每个layout内的行进行聚合."""
# 因为只是一个block一行目前, 一个block就是一个段落
blocks_group = []
for lyout in layout_bboxes:
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
def __split_para_in_layoutbox(blocks_group, new_layout_bbox, lang='en'):
"""
lines_group 进行行分段——layout内部进行分段。lines_group内每个元素是一个Layoutbox内的所有行。
1. 先计算每个group的左右边界。
2. 然后根据行末尾特征进行分段。
末尾特征：以句号等结束符结尾。并且距离右侧边界有一定距离。
且下一行开头不留空白。
"""
list_info = [] # 这个layout最后是不是列表,记录每一个layout里是不是列表开头，列表结尾
for blocks in blocks_group:
is_start_list = None
is_end_list = None
if len(blocks) == 0:
list_info.append([False, False])
continue
if blocks[0]['type'] != BlockType.Text and blocks[-1]['type'] != BlockType.Text:
list_info.append([False, False])
continue
if blocks[0]['type'] != BlockType.Text:
is_start_list = False
if blocks[-1]['type'] != BlockType.Text:
is_end_list = False
lines = [
line
for block in blocks
if block['type'] == BlockType.Text
for line in block['lines']
]
total_lines = len(lines)
if total_lines == 1 or total_lines == 0:
list_info.append([False, False])
continue
"""在进入到真正的分段之前，要对文字块从统计维度进行对齐方式的探测，
对齐方式分为以下：
1. 左对齐的文本块(特点是左侧顶格，或者左侧不顶格但是右侧顶格的行数大于非顶格的行数，顶格的首字母有大写也有小写)
1) 右侧对齐的行，单独成一段
2) 中间对齐的行，按照字体/行高聚合成一段
2. 左对齐的列表块（其特点是左侧顶格的行数小于等于非顶格的行数，非定格首字母会有小写，顶格90%是大写。并且左侧顶格行数大于1，大于1是为了这种模式连续出现才能称之为列表）
这样的文本块，顶格的为一个段落开头，紧随其后非顶格的行属于这个段落。
"""
text_segments, list_start_line = __detect_list_lines(
lines, new_layout_bbox, lang
)
"""根据list_range，把lines分成几个部分
"""
for list_start in list_start_line:
if len(list_start) > 1:
for i in range(0, len(list_start)):
index = list_start[i] - 1
if index >= 0:
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:
if content_type == 'list':
if start == 0 and is_start_list is None:
layout_list_info[0] = True
if end == total_lines - 1 and is_end_list is None:
layout_list_info[1] = True
list_info.append(layout_list_info)
return list_info
def __split_para_lines(lines: list, text_blocks: list) -> list:
text_paras = []
other_paras = []
text_lines = []
for line in lines:
spans_types = [span['type'] for span in line]
if ContentType.Table in spans_types:
other_paras.append([line])
continue
if ContentType.Image in spans_types:
other_paras.append([line])
continue
if ContentType.InterlineEquation in spans_types:
other_paras.append([line])
continue
text_lines.append(line)
for block in text_blocks:
block_bbox = block['bbox']
para = []
for line in text_lines:
bbox = line['bbox']
if is_in_layout(bbox, block_bbox):
para.append(line)
if len(para) > 0:
text_paras.append(para)
paras = other_paras.extend(text_paras)
paras_sorted = sorted(paras, key=lambda x: x[0]['bbox'][1])
return paras_sorted
def __connect_list_inter_layout(
blocks_group, new_layout_bbox, layout_list_info, page_num, lang
):
global debug_able
"""
如果上个layout的最后一个段落是列表，下一个layout的第一个段落也是列表，那么将他们连接起来。 TODO 因为没有区分列表和段落，所以这个方法暂时不实现。
根据layout_list_info判断是不是列表。，下个layout的第一个段如果不是列表，那么看他们是否有几行都有相同的缩进。
"""
if len(blocks_group) == 0 or len(blocks_group) == 0: # 0的时候最后的return 会出错
return blocks_group, [False, False]
for i in range(1, len(blocks_group)):
if len(blocks_group[i]) == 0 or len(blocks_group[i - 1]) == 0:
continue
pre_layout_list_info = layout_list_info[i - 1]
next_layout_list_info = layout_list_info[i]
pre_last_para = blocks_group[i - 1][-1].get('lines', [])
next_paras = blocks_group[i]
next_first_para = next_paras[0]
if (
pre_layout_list_info[1]
and not next_layout_list_info[0]
and next_first_para['type'] == BlockType.Text
): # 前一个是列表结尾，后一个是非列表开头，此时检测是否有相同的缩进
if debug_able:
logger.info(f'连接page {page_num} 内的list')
# 向layout_paras[i] 寻找开头具有相同缩进的连续的行
may_list_lines = []
lines = next_first_para.get('lines', [])
for line in lines:
if (
line['bbox'][0]
> __find_layout_bbox_by_line(line['bbox'], new_layout_bbox)[0]
):
may_list_lines.append(line)
else:
break
# 如果这些行的缩进是相等的，那么连到上一个layout的最后一个段落上。
if (
len(may_list_lines) > 0
and len(set([x['bbox'][0] for x in may_list_lines])) == 1
):
pre_last_para.extend(may_list_lines)
next_first_para['lines'] = next_first_para['lines'][
len(may_list_lines) :
]
return blocks_group, [
layout_list_info[0][0],
layout_list_info[-1][1],
] # 同时还返回了这个页面级别的开头、结尾是不是列表的信息
def __connect_list_inter_page(
pre_page_paras,
next_page_paras,
pre_page_layout_bbox,
next_page_layout_bbox,
pre_page_list_info,
next_page_list_info,
page_num,
lang,
):
"""如果上个layout的最后一个段落是列表，下一个layout的第一个段落也是列表，那么将他们连接起来。 TODO
因为没有区分列表和段落，所以这个方法暂时不实现。
根据layout_list_info判断是不是列表。，下个layout的第一个段如果不是列表，那么看他们是否有几行都有相同的缩进。"""
if (
len(pre_page_paras) == 0 or len(next_page_paras) == 0
): # 0的时候最后的return 会出错
return False
if len(pre_page_paras[-1]) == 0 or len(next_page_paras[0]) == 0:
return False
if (
pre_page_paras[-1][-1]['type'] != BlockType.Text
or next_page_paras[0][0]['type'] != BlockType.Text
):
return False
if (
pre_page_list_info[1] and not next_page_list_info[0]
): # 前一个是列表结尾，后一个是非列表开头，此时检测是否有相同的缩进
if debug_able:
logger.info(f'连接page {page_num} 内的list')
# 向layout_paras[i] 寻找开头具有相同缩进的连续的行
may_list_lines = []
next_page_first_para = next_page_paras[0][0]
if next_page_first_para['type'] == BlockType.Text:
lines = next_page_first_para['lines']
for line in lines:
if (
line['bbox'][0]
> __find_layout_bbox_by_line(line['bbox'], next_page_layout_bbox)[0]
):
may_list_lines.append(line)
else:
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)
# 下一页合并到上一页最后一段，打一个cross_page的标签
for line in may_list_lines:
for span in line['spans']:
span[CROSS_PAGE] = True # noqa: F405
pre_page_paras[-1][-1]['lines'].extend(may_list_lines)
next_page_first_para['lines'] = next_page_first_para['lines'][
len(may_list_lines) :
]
return True
return False
def __find_layout_bbox_by_line(line_bbox, layout_bboxes):
"""根据line找到所在的layout."""
for layout in layout_bboxes:
if is_in_layout(line_bbox, layout):
return layout
return None
def __connect_para_inter_layoutbox(blocks_group, new_layout_bbox):
"""
layout之间进行分段。
主要是计算前一个layOut的最后一行和后一个layout的第一行是否可以连接。
连接的条件需要同时满足：
1. 上一个layout的最后一行沾满整个行。并且没有结尾符号。
2. 下一行开头不留空白。
"""
connected_layout_blocks = []
if len(blocks_group) == 0:
return connected_layout_blocks
connected_layout_blocks.append(blocks_group[0])
for i in range(1, len(blocks_group)):
try:
if len(blocks_group[i]) == 0:
continue
if len(blocks_group[i - 1]) == 0: # TODO 考虑连接问题，
connected_layout_blocks.append(blocks_group[i])
continue
# text类型的段才需要考虑layout间的合并
if (
blocks_group[i - 1][-1]['type'] != BlockType.Text
or blocks_group[i][0]['type'] != BlockType.Text
):
connected_layout_blocks.append(blocks_group[i])
continue
if (
len(blocks_group[i - 1][-1]['lines']) == 0
or len(blocks_group[i][0]['lines']) == 0
):
connected_layout_blocks.append(blocks_group[i])
continue
pre_last_line = blocks_group[i - 1][-1]['lines'][-1]
next_first_line = blocks_group[i][0]['lines'][0]
except Exception:
logger.error(f'page layout {i} has no line')
continue
pre_last_line_text = ''.join(
[__get_span_text(span) for span in pre_last_line['spans']]
)
pre_last_line_type = pre_last_line['spans'][-1]['type']
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_blocks.append(blocks_group[i])
continue
pre_layout = __find_layout_bbox_by_line(pre_last_line['bbox'], new_layout_bbox)
next_layout = __find_layout_bbox_by_line(
next_first_line['bbox'], new_layout_bbox
)
pre_x2_max = pre_layout[2] if pre_layout else -1
next_x0_min = next_layout[0] if next_layout else -1
pre_last_line_text = pre_last_line_text.strip()
next_first_line_text = next_first_line_text.strip()
if (
pre_last_line['bbox'][2] == pre_x2_max
and pre_last_line_text
and pre_last_line_text[-1] not in LINE_STOP_FLAG
and next_first_line['bbox'][0] == next_x0_min
): # 前面一行沾满了整个行，并且没有结尾符号.下一行没有空白开头。
"""连接段落条件成立，将前一个layout的段落和后一个layout的段落连接。"""
connected_layout_blocks[-1][-1]['lines'].extend(blocks_group[i][0]['lines'])
blocks_group[i][0][
'lines'
] = [] # 删除后一个layout第一个段落中的lines，因为他已经被合并到前一个layout的最后一个段落了
blocks_group[i][0][LINES_DELETED] = True # noqa: F405
# if len(layout_paras[i]) == 0:
# layout_paras.pop(i)
# else:
# connected_layout_paras.append(layout_paras[i])
connected_layout_blocks.append(blocks_group[i])
else:
"""连接段落条件不成立，将前一个layout的段落加入到结果中。"""
connected_layout_blocks.append(blocks_group[i])
return connected_layout_blocks
def __connect_para_inter_page(
pre_page_paras,
next_page_paras,
pre_page_layout_bbox,
next_page_layout_bbox,
page_num,
lang,
):
"""
连接起来相邻两个页面的段落——前一个页面最后一个段落和后一个页面的第一个段落。
是否可以连接的条件：
1. 前一个页面的最后一个段落最后一行沾满整个行。并且没有结尾符号。
2. 后一个页面的第一个段落第一行没有空白开头。
"""
# 有的页面可能压根没有文字
if (
len(pre_page_paras) == 0
or len(next_page_paras) == 0
or len(pre_page_paras[0]) == 0
or len(next_page_paras[0]) == 0
): # TODO [[]]为什么出现在pre_page_paras里？
return False
pre_last_block = pre_page_paras[-1][-1]
next_first_block = next_page_paras[0][0]
if (
pre_last_block['type'] != BlockType.Text
or next_first_block['type'] != BlockType.Text
):
return False
if len(pre_last_block['lines']) == 0 or len(next_first_block['lines']) == 0:
return False
pre_last_para = pre_last_block['lines']
next_first_para = next_first_block['lines']
pre_last_line = pre_last_para[-1]
next_first_line = next_first_para[0]
pre_last_line_text = ''.join(
[__get_span_text(span) for span in pre_last_line['spans']]
)
pre_last_line_type = pre_last_line['spans'][-1]['type']
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,
]: # TODO，真的要做好，要考虑跨table, image, 行间的情况
# 不是文本，不连接
return False
pre_x2_max_bbox = __find_layout_bbox_by_line(
pre_last_line['bbox'], pre_page_layout_bbox
)
if not pre_x2_max_bbox:
return False
next_x0_min_bbox = __find_layout_bbox_by_line(
next_first_line['bbox'], next_page_layout_bbox
)
if not next_x0_min_bbox:
return False
pre_x2_max = pre_x2_max_bbox[2]
next_x0_min = next_x0_min_bbox[0]
pre_last_line_text = pre_last_line_text.strip()
next_first_line_text = next_first_line_text.strip()
if (
pre_last_line['bbox'][2] == pre_x2_max
and pre_last_line_text[-1] not in LINE_STOP_FLAG
and next_first_line['bbox'][0] == next_x0_min
): # 前面一行沾满了整个行，并且没有结尾符号.下一行没有空白开头。
"""连接段落条件成立，将前一个layout的段落和后一个layout的段落连接。"""
# 下一页合并到上一页最后一段，打一个cross_page的标签
for line in next_first_para:
for span in line['spans']:
span[CROSS_PAGE] = True # noqa: F405
pre_last_para.extend(next_first_para)
# next_page_paras[0].pop(0) # 删除后一个页面的第一个段落， 因为他已经被合并到前一个页面的最后一个段落了。
next_page_paras[0][0]['lines'] = []
next_page_paras[0][0][LINES_DELETED] = True # noqa: F405
return True
else:
return False
def find_consecutive_true_regions(input_array):
start_index = None # 连续True区域的起始索引
regions = [] # 用于保存所有连续True区域的起始和结束索引
for i in range(len(input_array)):
# 如果我们找到了一个True值，并且当前并没有在连续True区域中
if input_array[i] and start_index is None:
start_index = i # 记录连续True区域的起始索引
# 如果我们找到了一个False值，并且当前在连续True区域中
elif not input_array[i] and start_index is not None:
# 如果连续True区域长度大于1，那么将其添加到结果列表中
if i - start_index > 1:
regions.append((start_index, i - 1))
start_index = None # 重置起始索引
# 如果最后一个元素是True，那么需要将最后一个连续True区域加入到结果列表中
if start_index is not None and len(input_array) - start_index > 1:
regions.append((start_index, len(input_array) - 1))
return regions
def __connect_middle_align_text(page_paras, new_layout_bbox, page_num, lang):
global debug_able
"""
找出来中间对齐的连续单行文本，如果连续行高度相同，那么合并为一个段落。
一个line居中的条件是：
1. 水平中心点跨越layout的中心点。
2. 左右两侧都有空白
"""
for layout_i, layout_para in enumerate(page_paras):
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(
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:
"""检查这些行是否是高度相同的，居中的."""
for start, end in consecutive_single_line_indices:
# 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']
]
)
if 'Table' in first_line_text or 'Figure' in first_line_text:
pass
if debug_able:
logger.info(line_hi.std())
if line_hi.std() < 2:
"""行高度相同，那么判断是否居中."""
all_left_x0 = [
block['lines'][0]['bbox'][0]
for block in layout_para[start : end + 1]
]
all_right_x1 = [
block['lines'][0]['bbox'][2]
for block in layout_para[start : end + 1]
]
layout_center = (layout_box[0] + layout_box[2]) / 2
if (
all(
[
x0 < layout_center < x1
for x0, x1 in zip(all_left_x0, all_right_x1)
]
)
and not all([x0 == layout_box[0] for x0 in all_left_x0])
and not all([x1 == layout_box[2] for x1 in all_right_x1])
):
merge_para = [
block['lines'][0] for block in layout_para[start : end + 1]
]
para_text = ''.join(
[
__get_span_text(span)
for line in merge_para
for span in line['spans']
]
)
if debug_able:
logger.info(para_text)
layout_para[start]['lines'] = merge_para
for i_para in range(start + 1, end + 1):
layout_para[i_para]['lines'] = []
layout_para[i_para][LINES_DELETED] = True # noqa: F405
# layout_para[start:end + 1] = [merge_para]
# index_offset -= end - start
return
def __merge_signle_list_text(page_paras, new_layout_bbox, page_num, lang):
"""找出来连续的单行文本，如果首行顶格，接下来的几个单行段落缩进对齐，那么合并为一个段落。"""
pass
def __do_split_page(blocks, layout_bboxes, new_layout_bbox, page_num, lang):
"""根据line和layout情况进行分段 先实现一个根据行末尾特征分段的简单方法。"""
"""
算法思路：
1. 扫描layout里每一行，找出来行尾距离layout有边界有一定距离的行。
2. 从上述行中找到末尾是句号等可作为断行标志的行。
3. 参照上述行尾特征进行分段。
4. 图、表，目前独占一行，不考虑分段。
"""
blocks_group = __group_line_by_layout(blocks, layout_bboxes) # block内分段
layout_list_info = __split_para_in_layoutbox(
blocks_group, new_layout_bbox, lang
) # layout内分段
blocks_group, page_list_info = __connect_list_inter_layout(
blocks_group, new_layout_bbox, layout_list_info, page_num, lang
) # layout之间连接列表段落
connected_layout_blocks = __connect_para_inter_layoutbox(
blocks_group, new_layout_bbox
) # layout间链接段落
return connected_layout_blocks, page_list_info
def para_split(pdf_info_dict, debug_mode, lang='en'):
global debug_able
debug_able = debug_mode
new_layout_of_pages = [] # 数组的数组，每个元素是一个页面的layoutS
all_page_list_info = [] # 保存每个页面开头和结尾是否是列表
for page_num, page in pdf_info_dict.items():
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)
splited_blocks, page_list_info = __do_split_page(
blocks, layout_bboxes, new_layout_bbox, page_num, lang
)
all_page_list_info.append(page_list_info)
page['para_blocks'] = splited_blocks
"""连接页面与页面之间的可能合并的段落"""
pdf_infos = list(pdf_info_dict.values())
for page_num, page in enumerate(pdf_info_dict.values()):
if page_num == 0:
continue
pre_page_paras = pdf_infos[page_num - 1]['para_blocks']
next_page_paras = pdf_infos[page_num]['para_blocks']
pre_page_layout_bbox = new_layout_of_pages[page_num - 1]
next_page_layout_bbox = new_layout_of_pages[page_num]
is_conn = __connect_para_inter_page(
pre_page_paras,
next_page_paras,
pre_page_layout_bbox,
next_page_layout_bbox,
page_num,
lang,
)
if debug_able:
if is_conn:
logger.info(f'连接了第{page_num - 1}页和第{page_num}页的段落')
is_list_conn = __connect_list_inter_page(
pre_page_paras,
next_page_paras,
pre_page_layout_bbox,
next_page_layout_bbox,
all_page_list_info[page_num - 1],
all_page_list_info[page_num],
page_num,
lang,
)
if debug_able:
if is_list_conn:
logger.info(f'连接了第{page_num - 1}页和第{page_num}页的列表段落')
"""接下来可能会漏掉一些特别的一些可以合并的内容，对他们进行段落连接
1. 正文中有时出现一个行顶格，接下来几行缩进的情况。
2. 居中的一些连续单行，如果高度相同，那么可能是一个段落。
"""
for page_num, page in enumerate(pdf_info_dict.values()):
page_paras = page['para_blocks']
new_layout_bbox = new_layout_of_pages[page_num]
__connect_middle_align_text(page_paras, new_layout_bbox, page_num, lang)
__merge_signle_list_text(page_paras, new_layout_bbox, page_num, lang)
# layout展平
for page_num, page in enumerate(pdf_info_dict.values()):
page_paras = page['para_blocks']
page_blocks = [block for layout in page_paras for block in layout]
page['para_blocks'] = page_blocks
magic_pdf/para/raw_processor.py.bak deleted 100644 → 0
View file @ 79b58a1e
class RawBlockProcessor:
def __init__(self) -> None:
self.y_tolerance = 2
self.pdf_dic = {}
def __span_flags_decomposer(self, span_flags):
"""
Make font flags human readable.
Parameters
----------
self : object
The instance of the class.
span_flags : int
span flags
Returns
-------
l : dict
decomposed flags
"""
l = {
"is_superscript": False,
"is_italic": False,
"is_serifed": False,
"is_sans_serifed": False,
"is_monospaced": False,
"is_proportional": False,
"is_bold": False,
}
if span_flags & 2**0:
l["is_superscript"] = True # 表示上标
if span_flags & 2**1:
l["is_italic"] = True # 表示斜体
if span_flags & 2**2:
l["is_serifed"] = True # 表示衬线字体
else:
l["is_sans_serifed"] = True # 表示非衬线字体
if span_flags & 2**3:
l["is_monospaced"] = True # 表示等宽字体
else:
l["is_proportional"] = True # 表示比例字体
if span_flags & 2**4:
l["is_bold"] = True # 表示粗体
return l
def __make_new_lines(self, raw_lines):
"""
This function makes new lines.
Parameters
----------
self : object
The instance of the class.
raw_lines : list
raw lines
Returns
-------
new_lines : list
new lines
"""
new_lines = []
new_line = None
for raw_line in raw_lines:
raw_line_bbox = raw_line["bbox"]
raw_line_spans = raw_line["spans"]
raw_line_text = "".join([span["text"] for span in raw_line_spans])
raw_line_dir = raw_line.get("dir", None)
decomposed_line_spans = []
for span in raw_line_spans:
raw_flags = span["flags"]
decomposed_flags = self.__span_flags_decomposer(raw_flags)
span["decomposed_flags"] = decomposed_flags
decomposed_line_spans.append(span)
if new_line is None:
new_line = {
"bbox": raw_line_bbox,
"text": raw_line_text,
"dir": raw_line_dir if raw_line_dir else (0, 0),
"spans": decomposed_line_spans,
}
else:
if (
abs(raw_line_bbox[1] - new_line["bbox"][1]) <= self.y_tolerance
and abs(raw_line_bbox[3] - new_line["bbox"][3]) <= self.y_tolerance
):
new_line["bbox"] = (
min(new_line["bbox"][0], raw_line_bbox[0]), # left
new_line["bbox"][1], # top
max(new_line["bbox"][2], raw_line_bbox[2]), # right
raw_line_bbox[3], # bottom
)
new_line["text"] += " " + raw_line_text
new_line["spans"].extend(raw_line_spans)
new_line["dir"] = (
new_line["dir"][0] + raw_line_dir[0],
new_line["dir"][1] + raw_line_dir[1],
)
else:
new_lines.append(new_line)
new_line = {
"bbox": raw_line_bbox,
"text": raw_line_text,
"dir": raw_line_dir if raw_line_dir else (0, 0),
"spans": raw_line_spans,
}
if new_line:
new_lines.append(new_line)
return new_lines
def __make_new_block(self, raw_block):
"""
This function makes a new block.
Parameters
----------
self : object
The instance of the class.
----------
raw_block : dict
a raw block
Returns
-------
new_block : dict
Schema of new_block:
{
"block_id": "block_1",
"bbox": [0, 0, 100, 100],
"text": "This is a block.",
"lines": [
{
"bbox": [0, 0, 100, 100],
"text": "This is a line.",
"spans": [
{
"text": "This is a span.",
"font": "Times New Roman",
"size": 12,
"color": "#000000",
}
],
}
],
}
"""
new_block = {}
block_id = raw_block["number"]
block_bbox = raw_block["bbox"]
block_text = " ".join(span["text"] for line in raw_block["lines"] for span in line["spans"])
raw_lines = raw_block["lines"]
block_lines = self.__make_new_lines(raw_lines)
new_block["block_id"] = block_id
new_block["bbox"] = block_bbox
new_block["text"] = block_text
new_block["lines"] = block_lines
return new_block
def batch_process_blocks(self, pdf_dic):
"""
This function processes the blocks in batch.
Parameters
----------
self : object
The instance of the class.
----------
blocks : list
Input block is a list of raw blocks. Schema can refer to the value of key ""preproc_blocks", demo file is app/pdf_toolbox/tests/preproc_2_parasplit_example.json.
Returns
-------
result_dict : dict
result dictionary
"""
for page_id, blocks in pdf_dic.items():
if page_id.startswith("page_"):
para_blocks = []
if "preproc_blocks" in blocks.keys():
input_blocks = blocks["preproc_blocks"]
for raw_block in input_blocks:
new_block = self.__make_new_block(raw_block)
para_blocks.append(new_block)
blocks["para_blocks"] = para_blocks
return pdf_dic
magic_pdf/para/stats.py.bak deleted 100644 → 0
View file @ 79b58a1e
from collections import Counter
import numpy as np
from magic_pdf.para.commons import *
if sys.version_info[0] >= 3:
sys.stdout.reconfigure(encoding="utf-8") # type: ignore
class BlockStatisticsCalculator:
def __init__(self) -> None:
pass
def __calc_stats_of_new_lines(self, new_lines):
"""
This function calculates the paragraph metrics
Parameters
----------
combined_lines : list
combined lines
Returns
-------
X0 : float
Median of x0 values, which represents the left average boundary of the block
X1 : float
Median of x1 values, which represents the right average boundary of the block
avg_char_width : float
Average of char widths, which represents the average char width of the block
avg_char_height : float
Average of line heights, which represents the average line height of the block
"""
x0_values = []
x1_values = []
char_widths = []
char_heights = []
block_font_types = []
block_font_sizes = []
block_directions = []
if len(new_lines) > 0:
for i, line in enumerate(new_lines):
line_bbox = line["bbox"]
line_text = line["text"]
line_spans = line["spans"]
num_chars = len([ch for ch in line_text if not ch.isspace()])
x0_values.append(line_bbox[0])
x1_values.append(line_bbox[2])
if num_chars > 0:
char_width = (line_bbox[2] - line_bbox[0]) / num_chars
char_widths.append(char_width)
for span in line_spans:
block_font_types.append(span["font"])
block_font_sizes.append(span["size"])
if "dir" in line:
block_directions.append(line["dir"])
# line_font_types = [span["font"] for span in line_spans]
char_heights = [span["size"] for span in line_spans]
X0 = np.median(x0_values) if x0_values else 0
X1 = np.median(x1_values) if x1_values else 0
avg_char_width = sum(char_widths) / len(char_widths) if char_widths else 0
avg_char_height = sum(char_heights) / len(char_heights) if char_heights else 0
# max_freq_font_type = max(set(block_font_types), key=block_font_types.count) if block_font_types else None
max_span_length = 0
max_span_font_type = None
for line in new_lines:
line_spans = line["spans"]
for span in line_spans:
span_length = span["bbox"][2] - span["bbox"][0]
if span_length > max_span_length:
max_span_length = span_length
max_span_font_type = span["font"]
max_freq_font_type = max_span_font_type
avg_font_size = sum(block_font_sizes) / len(block_font_sizes) if block_font_sizes else None
avg_dir_horizontal = sum([dir[0] for dir in block_directions]) / len(block_directions) if block_directions else 0
avg_dir_vertical = sum([dir[1] for dir in block_directions]) / len(block_directions) if block_directions else 0
median_font_size = float(np.median(block_font_sizes)) if block_font_sizes else None
return (
X0,
X1,
avg_char_width,
avg_char_height,
max_freq_font_type,
avg_font_size,
(avg_dir_horizontal, avg_dir_vertical),
median_font_size,
)
def __make_new_block(self, input_block):
new_block = {}
raw_lines = input_block["lines"]
stats = self.__calc_stats_of_new_lines(raw_lines)
block_id = input_block["block_id"]
block_bbox = input_block["bbox"]
block_text = input_block["text"]
block_lines = raw_lines
block_avg_left_boundary = stats[0]
block_avg_right_boundary = stats[1]
block_avg_char_width = stats[2]
block_avg_char_height = stats[3]
block_font_type = stats[4]
block_font_size = stats[5]
block_direction = stats[6]
block_median_font_size = stats[7]
new_block["block_id"] = block_id
new_block["bbox"] = block_bbox
new_block["text"] = block_text
new_block["dir"] = block_direction
new_block["X0"] = block_avg_left_boundary
new_block["X1"] = block_avg_right_boundary
new_block["avg_char_width"] = block_avg_char_width
new_block["avg_char_height"] = block_avg_char_height
new_block["block_font_type"] = block_font_type
new_block["block_font_size"] = block_font_size
new_block["lines"] = block_lines
new_block["median_font_size"] = block_median_font_size
return new_block
def batch_process_blocks(self, pdf_dic):
"""
This function processes the blocks in batch.
Parameters
----------
self : object
The instance of the class.
----------
blocks : list
Input block is a list of raw blocks. Schema can refer to the value of key ""preproc_blocks", demo file is app/pdf_toolbox/tests/preproc_2_parasplit_example.json
Returns
-------
result_dict : dict
result dictionary
"""
for page_id, blocks in pdf_dic.items():
if page_id.startswith("page_"):
para_blocks = []
if "para_blocks" in blocks.keys():
input_blocks = blocks["para_blocks"]
for input_block in input_blocks:
new_block = self.__make_new_block(input_block)
para_blocks.append(new_block)
blocks["para_blocks"] = para_blocks
return pdf_dic
class DocStatisticsCalculator:
def __init__(self) -> None:
pass
def calc_stats_of_doc(self, pdf_dict):
"""
This function computes the statistics of the document
Parameters
----------
result_dict : dict
result dictionary
Returns
-------
statistics : dict
statistics of the document
"""
total_text_length = 0
total_num_blocks = 0
for page_id, blocks in pdf_dict.items():
if page_id.startswith("page_"):
if "para_blocks" in blocks.keys():
para_blocks = blocks["para_blocks"]
for para_block in para_blocks:
total_text_length += len(para_block["text"])
total_num_blocks += 1
avg_text_length = total_text_length / total_num_blocks if total_num_blocks else 0
font_list = []
for page_id, blocks in pdf_dict.items():
if page_id.startswith("page_"):
if "para_blocks" in blocks.keys():
input_blocks = blocks["para_blocks"]
for input_block in input_blocks:
block_text_length = len(input_block.get("text", ""))
if block_text_length < avg_text_length * 0.5:
continue
block_font_type = safe_get(input_block, "block_font_type", "")
block_font_size = safe_get(input_block, "block_font_size", 0)
font_list.append((block_font_type, block_font_size))
font_counter = Counter(font_list)
most_common_font = font_counter.most_common(1)[0] if font_list else (("", 0), 0)
second_most_common_font = font_counter.most_common(2)[1] if len(font_counter) > 1 else (("", 0), 0)
statistics = {
"num_pages": 0,
"num_blocks": 0,
"num_paras": 0,
"num_titles": 0,
"num_header_blocks": 0,
"num_footer_blocks": 0,
"num_watermark_blocks": 0,
"num_vertical_margin_note_blocks": 0,
"most_common_font_type": most_common_font[0][0],
"most_common_font_size": most_common_font[0][1],
"number_of_most_common_font": most_common_font[1],
"second_most_common_font_type": second_most_common_font[0][0],
"second_most_common_font_size": second_most_common_font[0][1],
"number_of_second_most_common_font": second_most_common_font[1],
"avg_text_length": avg_text_length,
}
for page_id, blocks in pdf_dict.items():
if page_id.startswith("page_"):
blocks = pdf_dict[page_id]["para_blocks"]
statistics["num_pages"] += 1
for block_id, block_data in enumerate(blocks):
statistics["num_blocks"] += 1
if "paras" in block_data.keys():
statistics["num_paras"] += len(block_data["paras"])
for line in block_data["lines"]:
if line.get("is_title", 0):
statistics["num_titles"] += 1
if block_data.get("is_header", 0):
statistics["num_header_blocks"] += 1
if block_data.get("is_footer", 0):
statistics["num_footer_blocks"] += 1
if block_data.get("is_watermark", 0):
statistics["num_watermark_blocks"] += 1
if block_data.get("is_vertical_margin_note", 0):
statistics["num_vertical_margin_note_blocks"] += 1
pdf_dict["statistics"] = statistics
return pdf_dict
magic_pdf/para/title_processor.py.bak deleted 100644 → 0
View file @ 79b58a1e
import os
import re
import numpy as np
from magic_pdf.libs.nlp_utils import NLPModels
from magic_pdf.para.commons import *
if sys.version_info[0] >= 3:
sys.stdout.reconfigure(encoding="utf-8") # type: ignore
class TitleProcessor:
def __init__(self, *doc_statistics) -> None:
if len(doc_statistics) > 0:
self.doc_statistics = doc_statistics[0]
self.nlp_model = NLPModels()
self.MAX_TITLE_LEVEL = 3
self.numbered_title_pattern = r"""
^ # 行首
( # 开始捕获组
[\(\（]\d+[\)\）] # 括号内数字，支持中文和英文括号，例如：(1) 或 （1）
|\d+[\)\）]\s # 数字后跟右括号和空格，支持中文和英文括号，例如：2) 或 2）
|[\(\（][A-Z][\)\）] # 括号内大写字母，支持中文和英文括号，例如：(A) 或 （A）
|[A-Z][\)\）]\s # 大写字母后跟右括号和空格，例如：A) 或 A）
|[\(\（][IVXLCDM]+[\)\）] # 括号内罗马数字，支持中文和英文括号，例如：(I) 或 （I）
|[IVXLCDM]+[\)\）]\s # 罗马数字后跟右括号和空格，例如：I) 或 I）
|\d+(\.\d+)*\s # 数字或复合数字编号后跟空格，例如：1. 或 3.2.1
|[一二三四五六七八九十百千]+[、\s] # 中文序号后跟顿号和空格，例如：一、
|[\（|\(][一二三四五六七八九十百千]+[\）|\)]\s* # 中文括号内中文序号后跟空格，例如：（一）
|[A-Z]\.\d+(\.\d+)?\s # 大写字母后跟点和数字，例如：A.1 或 A.1.1
|[\(\（][a-z][\)\）] # 括号内小写字母，支持中文和英文括号，例如：(a) 或 （a）
|[a-z]\)\s # 小写字母后跟右括号和空格，例如：a)
|[A-Z]-\s # 大写字母后跟短横线和空格，例如：A-
|\w+:\s # 英文序号词后跟冒号和空格，例如：First:
|第[一二三四五六七八九十百千]+[章节部分条款]\s # 以“第”开头的中文标题后跟空格
|[IVXLCDM]+\. # 罗马数字后跟点，例如：I.
|\d+\.\s # 单个数字后跟点和空格，例如：1.
) # 结束捕获组
.+ # 标题的其余部分
"""
def _is_potential_title(
self,
curr_line,
prev_line,
prev_line_is_title,
next_line,
avg_char_width,
avg_char_height,
median_font_size,
):
"""
This function checks if the line is a potential title.
Parameters
----------
curr_line : dict
current line
prev_line : dict
previous line
next_line : dict
next line
avg_char_width : float
average of char widths
avg_char_height : float
average of line heights
Returns
-------
bool
True if the line is a potential title, False otherwise.
"""
def __is_line_centered(line_bbox, page_bbox, avg_char_width):
"""
This function checks if the line is centered on the page
Parameters
----------
line_bbox : list
bbox of the line
page_bbox : list
bbox of the page
avg_char_width : float
average of char widths
Returns
-------
bool
True if the line is centered on the page, False otherwise.
"""
horizontal_ratio = 0.5
horizontal_thres = horizontal_ratio * avg_char_width
x0, _, x1, _ = line_bbox
_, _, page_x1, _ = page_bbox
return abs((x0 + x1) / 2 - page_x1 / 2) < horizontal_thres
def __is_bold_font_line(line):
"""
Check if a line contains any bold font style.
"""
def _is_bold_span(span):
# if span text is empty or only contains space, return False
if not span["text"].strip():
return False
return bool(span["flags"] & 2**4) # Check if the font is bold
for span in line["spans"]:
if not _is_bold_span(span):
return False
return True
def __is_italic_font_line(line):
"""
Check if a line contains any italic font style.
"""
def __is_italic_span(span):
return bool(span["flags"] & 2**1) # Check if the font is italic
for span in line["spans"]:
if not __is_italic_span(span):
return False
return True
def __is_punctuation_heavy(line_text):
"""
Check if the line contains a high ratio of punctuation marks, which may indicate
that the line is not a title.
Parameters:
line_text (str): Text of the line.
Returns:
bool: True if the line is heavy with punctuation, False otherwise.
"""
# Pattern for common title format like "X.Y. Title"
pattern = r"\b\d+\.\d+\..*\b"
# If the line matches the title format, return False
if re.match(pattern, line_text.strip()):
return False
# Find all punctuation marks in the line
punctuation_marks = re.findall(r"[^\w\s]", line_text)
number_of_punctuation_marks = len(punctuation_marks)
text_length = len(line_text)
if text_length == 0:
return False
punctuation_ratio = number_of_punctuation_marks / text_length
if punctuation_ratio >= 0.1:
return True
return False
def __has_mixed_font_styles(spans, strict_mode=False):
"""
This function checks if the line has mixed font styles, the strict mode will compare the font types
Parameters
----------
spans : list
spans of the line
strict_mode : bool
True for strict mode, the font types will be fully compared
False for non-strict mode, the font types will be compared by the most longest common prefix
Returns
-------
bool
True if the line has mixed font styles, False otherwise.
"""
if strict_mode:
font_styles = set()
for span in spans:
font_style = span["font"].lower()
font_styles.add(font_style)
return len(font_styles) > 1
else: # non-strict mode
font_styles = []
for span in spans:
font_style = span["font"].lower()
font_styles.append(font_style)
if len(font_styles) > 1:
longest_common_prefix = os.path.commonprefix(font_styles)
if len(longest_common_prefix) > 0:
return False
else:
return True
else:
return False
def __is_different_font_type_from_neighbors(curr_line_font_type, prev_line_font_type, next_line_font_type):
"""
This function checks if the current line has a different font type from the previous and next lines
Parameters
----------
curr_line_font_type : str
font type of the current line
prev_line_font_type : str
font type of the previous line
next_line_font_type : str
font type of the next line
Returns
-------
bool
True if the current line has a different font type from the previous and next lines, False otherwise.
"""
return all(
curr_line_font_type != other_font_type.lower()
for other_font_type in [prev_line_font_type, next_line_font_type]
if other_font_type is not None
)
def __is_larger_font_size_from_neighbors(curr_line_font_size, prev_line_font_size, next_line_font_size):
"""
This function checks if the current line has a larger font size than the previous and next lines
Parameters
----------
curr_line_font_size : float
font size of the current line
prev_line_font_size : float
font size of the previous line
next_line_font_size : float
font size of the next line
Returns
-------
bool
True if the current line has a larger font size than the previous and next lines, False otherwise.
"""
return all(
curr_line_font_size > other_font_size * 1.2
for other_font_size in [prev_line_font_size, next_line_font_size]
if other_font_size is not None
)
def __is_similar_to_pre_line(curr_line_font_type, prev_line_font_type, curr_line_font_size, prev_line_font_size):
"""
This function checks if the current line is similar to the previous line
Parameters
----------
curr_line : dict
current line
prev_line : dict
previous line
Returns
-------
bool
True if the current line is similar to the previous line, False otherwise.
"""
if curr_line_font_type == prev_line_font_type and curr_line_font_size == prev_line_font_size:
return True
else:
return False
def __is_same_font_type_of_docAvg(curr_line_font_type):
"""
This function checks if the current line has the same font type as the document average font type
Parameters
----------
curr_line_font_type : str
font type of the current line
Returns
-------
bool
True if the current line has the same font type as the document average font type, False otherwise.
"""
doc_most_common_font_type = safe_get(self.doc_statistics, "most_common_font_type", "").lower()
doc_second_most_common_font_type = safe_get(self.doc_statistics, "second_most_common_font_type", "").lower()
return curr_line_font_type.lower() in [doc_most_common_font_type, doc_second_most_common_font_type]
def __is_font_size_not_less_than_docAvg(curr_line_font_size, ratio: float = 1):
"""
This function checks if the current line has a large enough font size
Parameters
----------
curr_line_font_size : float
font size of the current line
ratio : float
ratio of the current line font size to the document average font size
Returns
-------
bool
True if the current line has a large enough font size, False otherwise.
"""
doc_most_common_font_size = safe_get(self.doc_statistics, "most_common_font_size", 0)
doc_second_most_common_font_size = safe_get(self.doc_statistics, "second_most_common_font_size", 0)
doc_avg_font_size = min(doc_most_common_font_size, doc_second_most_common_font_size)
return curr_line_font_size >= doc_avg_font_size * ratio
def __is_sufficient_spacing_above_and_below(
curr_line_bbox,
prev_line_bbox,
next_line_bbox,
avg_char_height,
median_font_size,
):
"""
This function checks if the current line has sufficient spacing above and below
Parameters
----------
curr_line_bbox : list
bbox of the current line
prev_line_bbox : list
bbox of the previous line
next_line_bbox : list
bbox of the next line
avg_char_width : float
average of char widths
avg_char_height : float
average of line heights
Returns
-------
bool
True if the current line has sufficient spacing above and below, False otherwise.
"""
vertical_ratio = 1.25
vertical_thres = vertical_ratio * median_font_size
_, y0, _, y1 = curr_line_bbox
sufficient_spacing_above = False
if prev_line_bbox:
vertical_spacing_above = min(y0 - prev_line_bbox[1], y1 - prev_line_bbox[3])
sufficient_spacing_above = vertical_spacing_above > vertical_thres
else:
sufficient_spacing_above = True
sufficient_spacing_below = False
if next_line_bbox:
vertical_spacing_below = min(next_line_bbox[1] - y0, next_line_bbox[3] - y1)
sufficient_spacing_below = vertical_spacing_below > vertical_thres
else:
sufficient_spacing_below = True
return (sufficient_spacing_above, sufficient_spacing_below)
def __is_word_list_line_by_rules(curr_line_text):
"""
This function checks if the current line is a word list
Parameters
----------
curr_line_text : str
text of the current line
Returns
-------
bool
True if the current line is a name list, False otherwise.
"""
# name_list_pattern = r"([a-zA-Z][a-zA-Z\s]{0,20}[a-zA-Z]|[\u4e00-\u9fa5·]{2,16})(?=[，,;；\s]|$)"
name_list_pattern = r"(?<![\u4e00-\u9fa5])([A-Z][a-z]{0,19}\s[A-Z][a-z]{0,19}|[\u4e00-\u9fa5]{2,6})(?=[，,;；\s]|$)"
compiled_pattern = re.compile(name_list_pattern)
if compiled_pattern.search(curr_line_text):
return True
else:
return False
# """
def __get_text_catgr_by_nlp(curr_line_text):
"""
This function checks if the current line is a name list using nlp model, such as spacy
Parameters
----------
curr_line_text : str
text of the current line
Returns
-------
bool
True if the current line is a name list, False otherwise.
"""
result = self.nlp_model.detect_entity_catgr_using_nlp(curr_line_text)
return result
# """
def __is_numbered_title(curr_line_text):
"""
This function checks if the current line is a numbered list
Parameters
----------
curr_line_text : str
text of the current line
Returns
-------
bool
True if the current line is a numbered list, False otherwise.
"""
compiled_pattern = re.compile(self.numbered_title_pattern, re.VERBOSE)
if compiled_pattern.search(curr_line_text):
return True
else:
return False
def __is_end_with_ending_puncs(line_text):
"""
This function checks if the current line ends with a ending punctuation mark
Parameters
----------
line_text : str
text of the current line
Returns
-------
bool
True if the current line ends with a punctuation mark, False otherwise.
"""
end_puncs = [".", "?", "!", "。", "？", "！", "…"]
line_text = line_text.rstrip()
if line_text[-1] in end_puncs:
return True
return False
def __contains_only_no_meaning_symbols(line_text):
"""
This function checks if the current line contains only symbols that have no meaning, if so, it is not a title.
Situation contains:
1. Only have punctuation marks
2. Only have other non-meaning symbols
Parameters
----------
line_text : str
text of the current line
Returns
-------
bool
True if the current line contains only symbols that have no meaning, False otherwise.
"""
punctuation_marks = re.findall(r"[^\w\s]", line_text) # find all punctuation marks
number_of_punctuation_marks = len(punctuation_marks)
text_length = len(line_text)
if text_length == 0:
return False
punctuation_ratio = number_of_punctuation_marks / text_length
if punctuation_ratio >= 0.9:
return True
return False
def __is_equation(line_text):
"""
This function checks if the current line is an equation.
Parameters
----------
line_text : str
Returns
-------
bool
True if the current line is an equation, False otherwise.
"""
equation_reg = r"\$.*?\\overline.*?\$" # to match interline equations
if re.search(equation_reg, line_text):
return True
else:
return False
def __is_title_by_len(text, max_length=200):
"""
This function checks if the current line is a title by length.
Parameters
----------
text : str
text of the current line
max_length : int
max length of the title
Returns
-------
bool
True if the current line is a title, False otherwise.
"""
text = text.strip()
return len(text) <= max_length
def __compute_line_font_type_and_size(curr_line):
"""
This function computes the font type and font size of the line.
Parameters
----------
line : dict
line
Returns
-------
font_type : str
font type of the line
font_size : float
font size of the line
"""
spans = curr_line["spans"]
max_accumulated_length = 0
max_span_font_size = curr_line["spans"][0]["size"] # default value, float type
max_span_font_type = curr_line["spans"][0]["font"].lower() # default value, string type
for span in spans:
if span["text"].isspace():
continue
span_length = span["bbox"][2] - span["bbox"][0]
if span_length > max_accumulated_length:
max_accumulated_length = span_length
max_span_font_size = span["size"]
max_span_font_type = span["font"].lower()
return max_span_font_type, max_span_font_size
"""
Title detecting main Process.
"""
"""
Basic features about the current line.
"""
curr_line_bbox = curr_line["bbox"]
curr_line_text = curr_line["text"]
curr_line_font_type, curr_line_font_size = __compute_line_font_type_and_size(curr_line)
if len(curr_line_text.strip()) == 0: # skip empty lines
return False
prev_line_bbox = prev_line["bbox"] if prev_line else None
if prev_line:
prev_line_font_type, prev_line_font_size = __compute_line_font_type_and_size(prev_line)
else:
prev_line_font_type, prev_line_font_size = None, None
next_line_bbox = next_line["bbox"] if next_line else None
if next_line:
next_line_font_type, next_line_font_size = __compute_line_font_type_and_size(next_line)
else:
next_line_font_type, next_line_font_size = None, None
"""
Aggregated features about the current line.
"""
is_italc_font = __is_italic_font_line(curr_line)
is_bold_font = __is_bold_font_line(curr_line)
is_font_size_little_less_than_doc_avg = __is_font_size_not_less_than_docAvg(curr_line_font_size, ratio=0.8)
is_font_size_not_less_than_doc_avg = __is_font_size_not_less_than_docAvg(curr_line_font_size, ratio=1)
is_much_larger_font_than_doc_avg = __is_font_size_not_less_than_docAvg(curr_line_font_size, ratio=1.6)
is_not_same_font_type_of_docAvg = not __is_same_font_type_of_docAvg(curr_line_font_type)
is_potential_title_font = is_bold_font or is_font_size_not_less_than_doc_avg or is_not_same_font_type_of_docAvg
is_mix_font_styles_strict = __has_mixed_font_styles(curr_line["spans"], strict_mode=True)
is_mix_font_styles_loose = __has_mixed_font_styles(curr_line["spans"], strict_mode=False)
is_punctuation_heavy = __is_punctuation_heavy(curr_line_text)
is_word_list_line_by_rules = __is_word_list_line_by_rules(curr_line_text)
is_person_or_org_list_line_by_nlp = __get_text_catgr_by_nlp(curr_line_text) in ["PERSON", "GPE", "ORG"]
is_font_size_larger_than_neighbors = __is_larger_font_size_from_neighbors(
curr_line_font_size, prev_line_font_size, next_line_font_size
)
is_font_type_diff_from_neighbors = __is_different_font_type_from_neighbors(
curr_line_font_type, prev_line_font_type, next_line_font_type
)
has_sufficient_spaces_above, has_sufficient_spaces_below = __is_sufficient_spacing_above_and_below(
curr_line_bbox, prev_line_bbox, next_line_bbox, avg_char_height, median_font_size
)
is_similar_to_pre_line = __is_similar_to_pre_line(
curr_line_font_type, prev_line_font_type, curr_line_font_size, prev_line_font_size
)
"""
Further aggregated features about the current line.
Attention:
Features that start with __ are for internal use.
"""
__is_line_left_aligned_from_neighbors = is_line_left_aligned_from_neighbors(
curr_line_bbox, prev_line_bbox, next_line_bbox, avg_char_width
)
__is_font_diff_from_neighbors = is_font_size_larger_than_neighbors or is_font_type_diff_from_neighbors
is_a_left_inline_title = (
is_mix_font_styles_strict and __is_line_left_aligned_from_neighbors and __is_font_diff_from_neighbors
)
is_title_by_check_prev_line = prev_line is None and has_sufficient_spaces_above and is_potential_title_font
is_title_by_check_next_line = next_line is None and has_sufficient_spaces_below and is_potential_title_font
is_title_by_check_pre_and_next_line = (
(prev_line is not None or next_line is not None)
and has_sufficient_spaces_above
and has_sufficient_spaces_below
and is_potential_title_font
)
is_numbered_title = __is_numbered_title(curr_line_text) and (
(has_sufficient_spaces_above or prev_line is None) and (has_sufficient_spaces_below or next_line is None)
)
is_not_end_with_ending_puncs = not __is_end_with_ending_puncs(curr_line_text)
is_not_only_no_meaning_symbols = not __contains_only_no_meaning_symbols(curr_line_text)
is_equation = __is_equation(curr_line_text)
is_title_by_len = __is_title_by_len(curr_line_text)
"""
Decide if the line is a title.
"""
# is_title = False
# if prev_line_is_title:
is_title = (
is_not_end_with_ending_puncs # not end with ending punctuation marks
and is_not_only_no_meaning_symbols # not only have no meaning symbols
and is_title_by_len # is a title by length, default max length is 200
and not is_equation # an interline equation should never be a title
and is_potential_title_font # is a potential title font, which is bold or larger than the document average font size or not the same font type as the document average font type
and (
(is_not_same_font_type_of_docAvg and is_font_size_not_less_than_doc_avg)
or (is_bold_font and is_much_larger_font_than_doc_avg and is_not_same_font_type_of_docAvg)
or (
is_much_larger_font_than_doc_avg
and (is_title_by_check_prev_line or is_title_by_check_next_line or is_title_by_check_pre_and_next_line)
)
or (
is_font_size_little_less_than_doc_avg
and is_bold_font
and (is_title_by_check_prev_line or is_title_by_check_next_line or is_title_by_check_pre_and_next_line)
)
) # not the same font type as the document average font type, which includes the most common font type and the second most common font type
and (
(
not is_person_or_org_list_line_by_nlp
and (
is_much_larger_font_than_doc_avg
or (is_not_same_font_type_of_docAvg and is_font_size_not_less_than_doc_avg)
)
)
or (
not (is_word_list_line_by_rules and is_person_or_org_list_line_by_nlp)
and not is_a_left_inline_title
and not is_punctuation_heavy
and (is_title_by_check_prev_line or is_title_by_check_next_line or is_title_by_check_pre_and_next_line)
)
or (
is_person_or_org_list_line_by_nlp
and (is_bold_font and is_much_larger_font_than_doc_avg and is_not_same_font_type_of_docAvg)
and (is_bold_font and is_much_larger_font_than_doc_avg and is_not_same_font_type_of_docAvg)
)
or (is_numbered_title and not is_a_left_inline_title)
)
)
# ) or (is_similar_to_pre_line and prev_line_is_title)
is_name_or_org_list_to_be_removed = (
(is_person_or_org_list_line_by_nlp)
and is_punctuation_heavy
and (is_title_by_check_prev_line or is_title_by_check_next_line or is_title_by_check_pre_and_next_line)
) and not is_title
if is_name_or_org_list_to_be_removed:
is_author_or_org_list = True
# print curr_line_text to check
# print_yellow(f"Text of is_author_or_org_list: {curr_line_text}")
else:
is_author_or_org_list = False
"""
# print reason why the line is a title
if is_title:
print_green("This line is a title.")
print_green("↓" * 10)
print()
print("curr_line_text: ", curr_line_text)
print()
# print reason why the line is not a title
line_text = curr_line_text.strip()
test_text = "Career/Personal Life"
text_content_condition = line_text == test_text
if not is_title and text_content_condition: # Print specific line
# if not is_title: # Print each line
print_red("This line is not a title.")
print_red("↓" * 10)
print()
print("curr_line_text: ", curr_line_text)
print()
if is_not_end_with_ending_puncs:
print_green(f"is_not_end_with_ending_puncs")
else:
print_red(f"is_end_with_ending_puncs")
if is_not_only_no_meaning_symbols:
print_green(f"is_not_only_no_meaning_symbols")
else:
print_red(f"is_only_no_meaning_symbols")
if is_title_by_len:
print_green(f"is_title_by_len: {is_title_by_len}")
else:
print_red(f"is_not_title_by_len: {is_title_by_len}")
if is_equation:
print_red(f"is_equation")
else:
print_green(f"is_not_equation")
if is_potential_title_font:
print_green(f"is_potential_title_font")
else:
print_red(f"is_not_potential_title_font")
if is_punctuation_heavy:
print_red("is_punctuation_heavy")
else:
print_green("is_not_punctuation_heavy")
if is_bold_font:
print_green(f"is_bold_font")
else:
print_red(f"is_not_bold_font")
if is_font_size_not_less_than_doc_avg:
print_green(f"is_larger_font_than_doc_avg")
else:
print_red(f"is_not_larger_font_than_doc_avg")
if is_much_larger_font_than_doc_avg:
print_green(f"is_much_larger_font_than_doc_avg")
else:
print_red(f"is_not_much_larger_font_than_doc_avg")
if is_not_same_font_type_of_docAvg:
print_green(f"is_not_same_font_type_of_docAvg")
else:
print_red(f"is_same_font_type_of_docAvg")
if is_word_list_line_by_rules:
print_red("is_word_list_line_by_rules")
else:
print_green("is_not_name_list_by_rules")
if is_person_or_org_list_line_by_nlp:
print_red("is_person_or_org_list_line_by_nlp")
else:
print_green("is_not_person_or_org_list_line_by_nlp")
if not is_numbered_title:
print_red("is_not_numbered_title")
else:
print_green("is_numbered_title")
if is_a_left_inline_title:
print_red("is_a_left_inline_title")
else:
print_green("is_not_a_left_inline_title")
if not is_title_by_check_prev_line:
print_red("is_not_title_by_check_prev_line")
else:
print_green("is_title_by_check_prev_line")
if not is_title_by_check_next_line:
print_red("is_not_title_by_check_next_line")
else:
print_green("is_title_by_check_next_line")
if not is_title_by_check_pre_and_next_line:
print_red("is_not_title_by_check_pre_and_next_line")
else:
print_green("is_title_by_check_pre_and_next_line")
# print_green("Common features:")
# print_green("↓" * 10)
# print(f" curr_line_font_type: {curr_line_font_type}")
# print(f" curr_line_font_size: {curr_line_font_size}")
# print()
"""
return is_title, is_author_or_org_list
def _detect_block_title(self, input_block):
"""
Use the functions 'is_potential_title' to detect titles of each paragraph block.
If a line is a title, then the value of key 'is_title' of the line will be set to True.
"""
raw_lines = input_block["lines"]
prev_line_is_title_flag = False
for i, curr_line in enumerate(raw_lines):
prev_line = raw_lines[i - 1] if i > 0 else None
next_line = raw_lines[i + 1] if i < len(raw_lines) - 1 else None
blk_avg_char_width = input_block["avg_char_width"]
blk_avg_char_height = input_block["avg_char_height"]
blk_media_font_size = input_block["median_font_size"]
is_title, is_author_or_org_list = self._is_potential_title(
curr_line,
prev_line,
prev_line_is_title_flag,
next_line,
blk_avg_char_width,
blk_avg_char_height,
blk_media_font_size,
)
if is_title:
curr_line["is_title"] = is_title
prev_line_is_title_flag = True
else:
curr_line["is_title"] = False
prev_line_is_title_flag = False
if is_author_or_org_list:
curr_line["is_author_or_org_list"] = is_author_or_org_list
else:
curr_line["is_author_or_org_list"] = False
return input_block
def batch_process_blocks_detect_titles(self, pdf_dic):
"""
This function batch process the blocks to detect titles.
Parameters
----------
pdf_dict : dict
result dictionary
Returns
-------
pdf_dict : dict
result dictionary
"""
num_titles = 0
for page_id, blocks in pdf_dic.items():
if page_id.startswith("page_"):
para_blocks = []
if "para_blocks" in blocks.keys():
para_blocks = blocks["para_blocks"]
all_single_line_blocks = []
for block in para_blocks:
if len(block["lines"]) == 1:
all_single_line_blocks.append(block)
new_para_blocks = []
if not len(all_single_line_blocks) == len(para_blocks): # Not all blocks are single line blocks.
for para_block in para_blocks:
new_block = self._detect_block_title(para_block)
new_para_blocks.append(new_block)
num_titles += sum([line.get("is_title", 0) for line in new_block["lines"]])
else: # All blocks are single line blocks.
for para_block in para_blocks:
new_para_blocks.append(para_block)
num_titles += sum([line.get("is_title", 0) for line in para_block["lines"]])
para_blocks = new_para_blocks
blocks["para_blocks"] = para_blocks
for para_block in para_blocks:
all_titles = all(safe_get(line, "is_title", False) for line in para_block["lines"])
para_text_len = sum([len(line["text"]) for line in para_block["lines"]])
if (
all_titles and para_text_len < 200
): # total length of the paragraph is less than 200, more than this should not be a title
para_block["is_block_title"] = 1
else:
para_block["is_block_title"] = 0
all_name_or_org_list_to_be_removed = all(
safe_get(line, "is_author_or_org_list", False) for line in para_block["lines"]
)
if all_name_or_org_list_to_be_removed and page_id == "page_0":
para_block["is_block_an_author_or_org_list"] = 1
else:
para_block["is_block_an_author_or_org_list"] = 0
pdf_dic["statistics"]["num_titles"] = num_titles
return pdf_dic
def __determine_size_based_level(self, title_blocks):
"""
This function determines the title level based on the font size of the title.
Parameters
----------
title_blocks : list
Returns
-------
title_blocks : list
"""
font_sizes = np.array([safe_get(tb["block"], "block_font_size", 0) for tb in title_blocks])
# Use the mean and std of font sizes to remove extreme values
mean_font_size = np.mean(font_sizes)
std_font_size = np.std(font_sizes)
min_extreme_font_size = mean_font_size - std_font_size # type: ignore
max_extreme_font_size = mean_font_size + std_font_size # type: ignore
# Compute the threshold for title level
middle_font_sizes = font_sizes[(font_sizes > min_extreme_font_size) & (font_sizes < max_extreme_font_size)]
if middle_font_sizes.size > 0:
middle_mean_font_size = np.mean(middle_font_sizes)
level_threshold = middle_mean_font_size
else:
level_threshold = mean_font_size
for tb in title_blocks:
title_block = tb["block"]
title_font_size = safe_get(title_block, "block_font_size", 0)
current_level = 1 # Initialize title level, the biggest level is 1
# print(f"Before adjustment by font size, {current_level}")
if title_font_size >= max_extreme_font_size:
current_level = 1
elif title_font_size <= min_extreme_font_size:
current_level = 3
elif float(title_font_size) >= float(level_threshold):
current_level = 2
else:
current_level = 3
# print(f"After adjustment by font size, {current_level}")
title_block["block_title_level"] = current_level
return title_blocks
def batch_process_blocks_recog_title_level(self, pdf_dic):
title_blocks = []
# Collect all titles
for page_id, blocks in pdf_dic.items():
if page_id.startswith("page_"):
para_blocks = blocks.get("para_blocks", [])
for block in para_blocks:
if block.get("is_block_title"):
title_obj = {"page_id": page_id, "block": block}
title_blocks.append(title_obj)
# Determine title level
if title_blocks:
# Determine title level based on font size
title_blocks = self.__determine_size_based_level(title_blocks)
return pdf_dic
magic_pdf/post_proc.bak/detect_para.py.bak deleted 100644 → 0
View file @ 79b58a1e
import os
import sys
import json
import re
import math
import unicodedata
from collections import Counter
import numpy as np
from termcolor import cprint
from magic_pdf.libs.commons import fitz
from magic_pdf.libs.nlp_utils import NLPModels
if sys.version_info[0] >= 3:
sys.stdout.reconfigure(encoding="utf-8") # type: ignore
def open_pdf(pdf_path):
try:
pdf_document = fitz.open(pdf_path) # type: ignore
return pdf_document
except Exception as e:
print(f"无法打开PDF文件：{pdf_path}。原因是：{e}")
raise e
def print_green_on_red(text):
cprint(text, "green", "on_red", attrs=["bold"], end="\n\n")
def print_green(text):
print()
cprint(text, "green", attrs=["bold"], end="\n\n")
def print_red(text):
print()
cprint(text, "red", attrs=["bold"], end="\n\n")
def print_yellow(text):
print()
cprint(text, "yellow", attrs=["bold"], end="\n\n")
def safe_get(dict_obj, key, default):
val = dict_obj.get(key)
if val is None:
return default
else:
return val
def is_bbox_overlap(bbox1, bbox2):
"""
This function checks if bbox1 and bbox2 overlap or not
Parameters
----------
bbox1 : list
bbox1
bbox2 : list
bbox2
Returns
-------
bool
True if bbox1 and bbox2 overlap, else False
"""
x0_1, y0_1, x1_1, y1_1 = bbox1
x0_2, y0_2, x1_2, y1_2 = bbox2
if x0_1 > x1_2 or x0_2 > x1_1:
return False
if y0_1 > y1_2 or y0_2 > y1_1:
return False
return True
def is_in_bbox(bbox1, bbox2):
"""
This function checks if bbox1 is in bbox2
Parameters
----------
bbox1 : list
bbox1
bbox2 : list
bbox2
Returns
-------
bool
True if bbox1 is in bbox2, else False
"""
x0_1, y0_1, x1_1, y1_1 = bbox1
x0_2, y0_2, x1_2, y1_2 = bbox2
if x0_1 >= x0_2 and y0_1 >= y0_2 and x1_1 <= x1_2 and y1_1 <= y1_2:
return True
else:
return False
def calculate_para_bbox(lines):
"""
This function calculates the minimum bbox of the paragraph
Parameters
----------
lines : list
lines
Returns
-------
para_bbox : list
bbox of the paragraph
"""
x0 = min(line["bbox"][0] for line in lines)
y0 = min(line["bbox"][1] for line in lines)
x1 = max(line["bbox"][2] for line in lines)
y1 = max(line["bbox"][3] for line in lines)
return [x0, y0, x1, y1]
def is_line_right_aligned_from_neighbors(curr_line_bbox, prev_line_bbox, next_line_bbox, avg_char_width, direction=2):
"""
This function checks if the line is right aligned from its neighbors
Parameters
----------
curr_line_bbox : list
bbox of the current line
prev_line_bbox : list
bbox of the previous line
next_line_bbox : list
bbox of the next line
avg_char_width : float
average of char widths
direction : int
0 for prev, 1 for next, 2 for both
Returns
-------
bool
True if the line is right aligned from its neighbors, False otherwise.
"""
horizontal_ratio = 0.5
horizontal_thres = horizontal_ratio * avg_char_width
_, _, x1, _ = curr_line_bbox
_, _, prev_x1, _ = prev_line_bbox if prev_line_bbox else (0, 0, 0, 0)
_, _, next_x1, _ = next_line_bbox if next_line_bbox else (0, 0, 0, 0)
if direction == 0:
return abs(x1 - prev_x1) < horizontal_thres
elif direction == 1:
return abs(x1 - next_x1) < horizontal_thres
elif direction == 2:
return abs(x1 - prev_x1) < horizontal_thres and abs(x1 - next_x1) < horizontal_thres
else:
return False
def is_line_left_aligned_from_neighbors(curr_line_bbox, prev_line_bbox, next_line_bbox, avg_char_width, direction=2):
"""
This function checks if the line is left aligned from its neighbors
Parameters
----------
curr_line_bbox : list
bbox of the current line
prev_line_bbox : list
bbox of the previous line
next_line_bbox : list
bbox of the next line
avg_char_width : float
average of char widths
direction : int
0 for prev, 1 for next, 2 for both
Returns
-------
bool
True if the line is left aligned from its neighbors, False otherwise.
"""
horizontal_ratio = 0.5
horizontal_thres = horizontal_ratio * avg_char_width
x0, _, _, _ = curr_line_bbox
prev_x0, _, _, _ = prev_line_bbox if prev_line_bbox else (0, 0, 0, 0)
next_x0, _, _, _ = next_line_bbox if next_line_bbox else (0, 0, 0, 0)
if direction == 0:
return abs(x0 - prev_x0) < horizontal_thres
elif direction == 1:
return abs(x0 - next_x0) < horizontal_thres
elif direction == 2:
return abs(x0 - prev_x0) < horizontal_thres and abs(x0 - next_x0) < horizontal_thres
else:
return False
def end_with_punctuation(line_text):
"""
This function checks if the line ends with punctuation marks
"""
english_end_puncs = [".", "?", "!"]
chinese_end_puncs = ["。", "？", "！"]
end_puncs = english_end_puncs + chinese_end_puncs
last_non_space_char = None
for ch in line_text[::-1]:
if not ch.isspace():
last_non_space_char = ch
break
if last_non_space_char is None:
return False
return last_non_space_char in end_puncs
def is_nested_list(lst):
if isinstance(lst, list):
return any(isinstance(sub, list) for sub in lst)
return False
class DenseSingleLineBlockException(Exception):
"""
This class defines the exception type for dense single line-block.
"""
def __init__(self, message="DenseSingleLineBlockException"):
self.message = message
super().__init__(self.message)
def __str__(self):
return f"{self.message}"
def __repr__(self):
return f"{self.message}"
class TitleDetectionException(Exception):
"""
This class defines the exception type for title detection.
"""
def __init__(self, message="TitleDetectionException"):
self.message = message
super().__init__(self.message)
def __str__(self):
return f"{self.message}"
def __repr__(self):
return f"{self.message}"
class TitleLevelException(Exception):
"""
This class defines the exception type for title level.
"""
def __init__(self, message="TitleLevelException"):
self.message = message
super().__init__(self.message)
def __str__(self):
return f"{self.message}"
def __repr__(self):
return f"{self.message}"
class ParaSplitException(Exception):
"""
This class defines the exception type for paragraph splitting.
"""
def __init__(self, message="ParaSplitException"):
self.message = message
super().__init__(self.message)
def __str__(self):
return f"{self.message}"
def __repr__(self):
return f"{self.message}"
class ParaMergeException(Exception):
"""
This class defines the exception type for paragraph merging.
"""
def __init__(self, message="ParaMergeException"):
self.message = message
super().__init__(self.message)
def __str__(self):
return f"{self.message}"
def __repr__(self):
return f"{self.message}"
class DiscardByException:
"""
This class discards pdf files by exception
"""
def __init__(self) -> None:
pass
def discard_by_single_line_block(self, pdf_dic, exception: DenseSingleLineBlockException):
"""
This function discards pdf files by single line block exception
Parameters
----------
pdf_dic : dict
pdf dictionary
exception : str
exception message
Returns
-------
error_message : str
"""
exception_page_nums = 0
page_num = 0
for page_id, page in pdf_dic.items():
if page_id.startswith("page_"):
page_num += 1
if "preproc_blocks" in page.keys():
preproc_blocks = page["preproc_blocks"]
all_single_line_blocks = []
for block in preproc_blocks:
if len(block["lines"]) == 1:
all_single_line_blocks.append(block)
if len(preproc_blocks) > 0 and len(all_single_line_blocks) / len(preproc_blocks) > 0.9:
exception_page_nums += 1
if page_num == 0:
return None
if exception_page_nums / page_num > 0.1: # Low ratio means basically, whenever this is the case, it is discarded
return exception.message
return None
def discard_by_title_detection(self, pdf_dic, exception: TitleDetectionException):
"""
This function discards pdf files by title detection exception
Parameters
----------
pdf_dic : dict
pdf dictionary
exception : str
exception message
Returns
-------
error_message : str
"""
# return exception.message
return None
def discard_by_title_level(self, pdf_dic, exception: TitleLevelException):
"""
This function discards pdf files by title level exception
Parameters
----------
pdf_dic : dict
pdf dictionary
exception : str
exception message
Returns
-------
error_message : str
"""
# return exception.message
return None
def discard_by_split_para(self, pdf_dic, exception: ParaSplitException):
"""
This function discards pdf files by split para exception
Parameters
----------
pdf_dic : dict
pdf dictionary
exception : str
exception message
Returns
-------
error_message : str
"""
# return exception.message
return None
def discard_by_merge_para(self, pdf_dic, exception: ParaMergeException):
"""
This function discards pdf files by merge para exception
Parameters
----------
pdf_dic : dict
pdf dictionary
exception : str
exception message
Returns
-------
error_message : str
"""
# return exception.message
return None
class LayoutFilterProcessor:
def __init__(self) -> None:
pass
def batch_process_blocks(self, pdf_dict):
"""
This function processes the blocks in batch.
Parameters
----------
self : object
The instance of the class.
pdf_dict : dict
pdf dictionary
Returns
-------
pdf_dict : dict
pdf dictionary
"""
for page_id, blocks in pdf_dict.items():
if page_id.startswith("page_"):
if "layout_bboxes" in blocks.keys() and "para_blocks" in blocks.keys():
layout_bbox_objs = blocks["layout_bboxes"]
if layout_bbox_objs is None:
continue
layout_bboxes = [bbox_obj["layout_bbox"] for bbox_obj in layout_bbox_objs]
# Enlarge each value of x0, y0, x1, y1 for each layout_bbox to prevent loss of text.
layout_bboxes = [
[math.ceil(x0), math.ceil(y0), math.ceil(x1), math.ceil(y1)] for x0, y0, x1, y1 in layout_bboxes
]
para_blocks = blocks["para_blocks"]
if para_blocks is None:
continue
for lb_bbox in layout_bboxes:
for i, para_block in enumerate(para_blocks):
para_bbox = para_block["bbox"]
para_blocks[i]["in_layout"] = 0
if is_in_bbox(para_bbox, lb_bbox):
para_blocks[i]["in_layout"] = 1
blocks["para_blocks"] = para_blocks
return pdf_dict
class RawBlockProcessor:
def __init__(self) -> None:
self.y_tolerance = 2
self.pdf_dic = {}
def __span_flags_decomposer(self, span_flags):
"""
Make font flags human readable.
Parameters
----------
self : object
The instance of the class.
span_flags : int
span flags
Returns
-------
l : dict
decomposed flags
"""
l = {
"is_superscript": False,
"is_italic": False,
"is_serifed": False,
"is_sans_serifed": False,
"is_monospaced": False,
"is_proportional": False,
"is_bold": False,
}
if span_flags & 2**0:
l["is_superscript"] = True # 表示上标
if span_flags & 2**1:
l["is_italic"] = True # 表示斜体
if span_flags & 2**2:
l["is_serifed"] = True # 表示衬线字体
else:
l["is_sans_serifed"] = True # 表示非衬线字体
if span_flags & 2**3:
l["is_monospaced"] = True # 表示等宽字体
else:
l["is_proportional"] = True # 表示比例字体
if span_flags & 2**4:
l["is_bold"] = True # 表示粗体
return l
def __make_new_lines(self, raw_lines):
"""
This function makes new lines.
Parameters
----------
self : object
The instance of the class.
raw_lines : list
raw lines
Returns
-------
new_lines : list
new lines
"""
new_lines = []
new_line = None
for raw_line in raw_lines:
raw_line_bbox = raw_line["bbox"]
raw_line_spans = raw_line["spans"]
raw_line_text = "".join([span["text"] for span in raw_line_spans])
raw_line_dir = raw_line.get("dir", None)
decomposed_line_spans = []
for span in raw_line_spans:
raw_flags = span["flags"]
decomposed_flags = self.__span_flags_decomposer(raw_flags)
span["decomposed_flags"] = decomposed_flags
decomposed_line_spans.append(span)
if new_line is None: # Handle the first line
new_line = {
"bbox": raw_line_bbox,
"text": raw_line_text,
"dir": raw_line_dir if raw_line_dir else (0, 0),
"spans": decomposed_line_spans,
}
else: # Handle the rest lines
if (
abs(raw_line_bbox[1] - new_line["bbox"][1]) <= self.y_tolerance
and abs(raw_line_bbox[3] - new_line["bbox"][3]) <= self.y_tolerance
):
new_line["bbox"] = (
min(new_line["bbox"][0], raw_line_bbox[0]), # left
new_line["bbox"][1], # top
max(new_line["bbox"][2], raw_line_bbox[2]), # right
raw_line_bbox[3], # bottom
)
new_line["text"] += raw_line_text
new_line["spans"].extend(raw_line_spans)
new_line["dir"] = (
new_line["dir"][0] + raw_line_dir[0],
new_line["dir"][1] + raw_line_dir[1],
)
else:
new_lines.append(new_line)
new_line = {
"bbox": raw_line_bbox,
"text": raw_line_text,
"dir": raw_line_dir if raw_line_dir else (0, 0),
"spans": raw_line_spans,
}
if new_line:
new_lines.append(new_line)
return new_lines
def __make_new_block(self, raw_block):
"""
This function makes a new block.
Parameters
----------
self : object
The instance of the class.
----------
raw_block : dict
a raw block
Returns
-------
new_block : dict
"""
new_block = {}
block_id = raw_block["number"]
block_bbox = raw_block["bbox"]
block_text = "".join(span["text"] for line in raw_block["lines"] for span in line["spans"])
raw_lines = raw_block["lines"]
block_lines = self.__make_new_lines(raw_lines)
new_block["block_id"] = block_id
new_block["bbox"] = block_bbox
new_block["text"] = block_text
new_block["lines"] = block_lines
return new_block
def batch_process_blocks(self, pdf_dic):
"""
This function processes the blocks in batch.
Parameters
----------
self : object
The instance of the class.
----------
blocks : list
Input block is a list of raw blocks.
Returns
-------
result_dict : dict
result dictionary
"""
for page_id, blocks in pdf_dic.items():
if page_id.startswith("page_"):
para_blocks = []
if "preproc_blocks" in blocks.keys():
input_blocks = blocks["preproc_blocks"]
for raw_block in input_blocks:
new_block = self.__make_new_block(raw_block)
para_blocks.append(new_block)
blocks["para_blocks"] = para_blocks
return pdf_dic
class BlockStatisticsCalculator:
"""
This class calculates the statistics of the block.
"""
def __init__(self) -> None:
pass
def __calc_stats_of_new_lines(self, new_lines):
"""
This function calculates the paragraph metrics
Parameters
----------
combined_lines : list
combined lines
Returns
-------
X0 : float
Median of x0 values, which represents the left average boundary of the block
X1 : float
Median of x1 values, which represents the right average boundary of the block
avg_char_width : float
Average of char widths, which represents the average char width of the block
avg_char_height : float
Average of line heights, which represents the average line height of the block
"""
x0_values = []
x1_values = []
char_widths = []
char_heights = []
block_font_types = []
block_font_sizes = []
block_directions = []
if len(new_lines) > 0:
for i, line in enumerate(new_lines):
line_bbox = line["bbox"]
line_text = line["text"]
line_spans = line["spans"]
num_chars = len([ch for ch in line_text if not ch.isspace()])
x0_values.append(line_bbox[0])
x1_values.append(line_bbox[2])
if num_chars > 0:
char_width = (line_bbox[2] - line_bbox[0]) / num_chars
char_widths.append(char_width)
for span in line_spans:
block_font_types.append(span["font"])
block_font_sizes.append(span["size"])
if "dir" in line:
block_directions.append(line["dir"])
# line_font_types = [span["font"] for span in line_spans]
char_heights = [span["size"] for span in line_spans]
X0 = np.median(x0_values) if x0_values else 0
X1 = np.median(x1_values) if x1_values else 0
avg_char_width = sum(char_widths) / len(char_widths) if char_widths else 0
avg_char_height = sum(char_heights) / len(char_heights) if char_heights else 0
# max_freq_font_type = max(set(block_font_types), key=block_font_types.count) if block_font_types else None
max_span_length = 0
max_span_font_type = None
for line in new_lines:
line_spans = line["spans"]
for span in line_spans:
span_length = span["bbox"][2] - span["bbox"][0]
if span_length > max_span_length:
max_span_length = span_length
max_span_font_type = span["font"]
max_freq_font_type = max_span_font_type
avg_font_size = sum(block_font_sizes) / len(block_font_sizes) if block_font_sizes else None
avg_dir_horizontal = sum([dir[0] for dir in block_directions]) / len(block_directions) if block_directions else 0
avg_dir_vertical = sum([dir[1] for dir in block_directions]) / len(block_directions) if block_directions else 0
median_font_size = float(np.median(block_font_sizes)) if block_font_sizes else None
return (
X0,
X1,
avg_char_width,
avg_char_height,
max_freq_font_type,
avg_font_size,
(avg_dir_horizontal, avg_dir_vertical),
median_font_size,
)
def __make_new_block(self, input_block):
new_block = {}
raw_lines = input_block["lines"]
stats = self.__calc_stats_of_new_lines(raw_lines)
block_id = input_block["block_id"]
block_bbox = input_block["bbox"]
block_text = input_block["text"]
block_lines = raw_lines
block_avg_left_boundary = stats[0]
block_avg_right_boundary = stats[1]
block_avg_char_width = stats[2]
block_avg_char_height = stats[3]
block_font_type = stats[4]
block_font_size = stats[5]
block_direction = stats[6]
block_median_font_size = stats[7]
new_block["block_id"] = block_id
new_block["bbox"] = block_bbox
new_block["text"] = block_text
new_block["dir"] = block_direction
new_block["X0"] = block_avg_left_boundary
new_block["X1"] = block_avg_right_boundary
new_block["avg_char_width"] = block_avg_char_width
new_block["avg_char_height"] = block_avg_char_height
new_block["block_font_type"] = block_font_type
new_block["block_font_size"] = block_font_size
new_block["lines"] = block_lines
new_block["median_font_size"] = block_median_font_size
return new_block
def batch_process_blocks(self, pdf_dic):
"""
This function processes the blocks in batch.
Parameters
----------
self : object
The instance of the class.
----------
blocks : list
Input block is a list of raw blocks.
Schema can refer to the value of key ""preproc_blocks".
Returns
-------
result_dict : dict
result dictionary
"""
for page_id, blocks in pdf_dic.items():
if page_id.startswith("page_"):
para_blocks = []
if "para_blocks" in blocks.keys():
input_blocks = blocks["para_blocks"]
for input_block in input_blocks:
new_block = self.__make_new_block(input_block)
para_blocks.append(new_block)
blocks["para_blocks"] = para_blocks
return pdf_dic
class DocStatisticsCalculator:
"""
This class calculates the statistics of the document.
"""
def __init__(self) -> None:
pass
def calc_stats_of_doc(self, pdf_dict):
"""
This function computes the statistics of the document
Parameters
----------
result_dict : dict
result dictionary
Returns
-------
statistics : dict
statistics of the document
"""
total_text_length = 0
total_num_blocks = 0
for page_id, blocks in pdf_dict.items():
if page_id.startswith("page_"):
if "para_blocks" in blocks.keys():
para_blocks = blocks["para_blocks"]
for para_block in para_blocks:
total_text_length += len(para_block["text"])
total_num_blocks += 1
avg_text_length = total_text_length / total_num_blocks if total_num_blocks else 0
font_list = []
for page_id, blocks in pdf_dict.items():
if page_id.startswith("page_"):
if "para_blocks" in blocks.keys():
input_blocks = blocks["para_blocks"]
for input_block in input_blocks:
block_text_length = len(input_block.get("text", ""))
if block_text_length < avg_text_length * 0.5:
continue
block_font_type = safe_get(input_block, "block_font_type", "")
block_font_size = safe_get(input_block, "block_font_size", 0)
font_list.append((block_font_type, block_font_size))
font_counter = Counter(font_list)
most_common_font = font_counter.most_common(1)[0] if font_list else (("", 0), 0)
second_most_common_font = font_counter.most_common(2)[1] if len(font_counter) > 1 else (("", 0), 0)
statistics = {
"num_pages": 0,
"num_blocks": 0,
"num_paras": 0,
"num_titles": 0,
"num_header_blocks": 0,
"num_footer_blocks": 0,
"num_watermark_blocks": 0,
"num_vertical_margin_note_blocks": 0,
"most_common_font_type": most_common_font[0][0],
"most_common_font_size": most_common_font[0][1],
"number_of_most_common_font": most_common_font[1],
"second_most_common_font_type": second_most_common_font[0][0],
"second_most_common_font_size": second_most_common_font[0][1],
"number_of_second_most_common_font": second_most_common_font[1],
"avg_text_length": avg_text_length,
}
for page_id, blocks in pdf_dict.items():
if page_id.startswith("page_"):
blocks = pdf_dict[page_id]["para_blocks"]
statistics["num_pages"] += 1
for block_id, block_data in enumerate(blocks):
statistics["num_blocks"] += 1
if "paras" in block_data.keys():
statistics["num_paras"] += len(block_data["paras"])
for line in block_data["lines"]:
if line.get("is_title", 0):
statistics["num_titles"] += 1
if block_data.get("is_header", 0):
statistics["num_header_blocks"] += 1
if block_data.get("is_footer", 0):
statistics["num_footer_blocks"] += 1
if block_data.get("is_watermark", 0):
statistics["num_watermark_blocks"] += 1
if block_data.get("is_vertical_margin_note", 0):
statistics["num_vertical_margin_note_blocks"] += 1
pdf_dict["statistics"] = statistics
return pdf_dict
class TitleProcessor:
"""
This class processes the title.
"""
def __init__(self, *doc_statistics) -> None:
if len(doc_statistics) > 0:
self.doc_statistics = doc_statistics[0]
self.nlp_model = NLPModels()
self.MAX_TITLE_LEVEL = 3
self.numbered_title_pattern = r"""
^ # 行首
( # 开始捕获组
[\(\（]\d+[\)\）] # 括号内数字，支持中文和英文括号，例如：(1) 或 （1）
|\d+[\)\）]\s # 数字后跟右括号和空格，支持中文和英文括号，例如：2) 或 2）
|[\(\（][A-Z][\)\）] # 括号内大写字母，支持中文和英文括号，例如：(A) 或 （A）
|[A-Z][\)\）]\s # 大写字母后跟右括号和空格，例如：A) 或 A）
|[\(\（][IVXLCDM]+[\)\）] # 括号内罗马数字，支持中文和英文括号，例如：(I) 或 （I）
|[IVXLCDM]+[\)\）]\s # 罗马数字后跟右括号和空格，例如：I) 或 I）
|\d+(\.\d+)*\s # 数字或复合数字编号后跟空格，例如：1. 或 3.2.1
|[一二三四五六七八九十百千]+[、\s] # 中文序号后跟顿号和空格，例如：一、
|[\（|\(][一二三四五六七八九十百千]+[\）|\)]\s* # 中文括号内中文序号后跟空格，例如：（一）
|[A-Z]\.\d+(\.\d+)?\s # 大写字母后跟点和数字，例如：A.1 或 A.1.1
|[\(\（][a-z][\)\）] # 括号内小写字母，支持中文和英文括号，例如：(a) 或 （a）
|[a-z]\)\s # 小写字母后跟右括号和空格，例如：a)
|[A-Z]-\s # 大写字母后跟短横线和空格，例如：A-
|\w+:\s # 英文序号词后跟冒号和空格，例如：First:
|第[一二三四五六七八九十百千]+[章节部分条款]\s # 以“第”开头的中文标题后跟空格
|[IVXLCDM]+\. # 罗马数字后跟点，例如：I.
|\d+\.\s # 单个数字后跟点和空格，例如：1.
) # 结束捕获组
.+ # 标题的其余部分
"""
def _is_potential_title(
self,
curr_line,
prev_line,
prev_line_is_title,
next_line,
avg_char_width,
avg_char_height,
median_font_size,
):
"""
This function checks if the line is a potential title.
Parameters
----------
curr_line : dict
current line
prev_line : dict
previous line
next_line : dict
next line
avg_char_width : float
average of char widths
avg_char_height : float
average of line heights
Returns
-------
bool
True if the line is a potential title, False otherwise.
"""
def __is_line_centered(line_bbox, page_bbox, avg_char_width):
"""
This function checks if the line is centered on the page
Parameters
----------
line_bbox : list
bbox of the line
page_bbox : list
bbox of the page
avg_char_width : float
average of char widths
Returns
-------
bool
True if the line is centered on the page, False otherwise.
"""
horizontal_ratio = 0.5
horizontal_thres = horizontal_ratio * avg_char_width
x0, _, x1, _ = line_bbox
_, _, page_x1, _ = page_bbox
return abs((x0 + x1) / 2 - page_x1 / 2) < horizontal_thres
def __is_bold_font_line(line):
"""
Check if a line contains any bold font style.
"""
def _is_bold_span(span):
# if span text is empty or only contains space, return False
if not span["text"].strip():
return False
return bool(span["flags"] & 2**4) # Check if the font is bold
for span in line["spans"]:
if not _is_bold_span(span):
return False
return True
def __is_italic_font_line(line):
"""
Check if a line contains any italic font style.
"""
def __is_italic_span(span):
return bool(span["flags"] & 2**1) # Check if the font is italic
for span in line["spans"]:
if not __is_italic_span(span):
return False
return True
def __is_punctuation_heavy(line_text):
"""
Check if the line contains a high ratio of punctuation marks, which may indicate
that the line is not a title.
Parameters:
line_text (str): Text of the line.
Returns:
bool: True if the line is heavy with punctuation, False otherwise.
"""
# Pattern for common title format like "X.Y. Title"
pattern = r"\b\d+\.\d+\..*\b"
# If the line matches the title format, return False
if re.match(pattern, line_text.strip()):
return False
# Find all punctuation marks in the line
punctuation_marks = re.findall(r"[^\w\s]", line_text)
number_of_punctuation_marks = len(punctuation_marks)
text_length = len(line_text)
if text_length == 0:
return False
punctuation_ratio = number_of_punctuation_marks / text_length
if punctuation_ratio >= 0.1:
return True
return False
def __has_mixed_font_styles(spans, strict_mode=False):
"""
This function checks if the line has mixed font styles, the strict mode will compare the font types
Parameters
----------
spans : list
spans of the line
strict_mode : bool
True for strict mode, the font types will be fully compared
False for non-strict mode, the font types will be compared by the most longest common prefix
Returns
-------
bool
True if the line has mixed font styles, False otherwise.
"""
if strict_mode:
font_styles = set()
for span in spans:
font_style = span["font"].lower()
font_styles.add(font_style)
return len(font_styles) > 1
else: # non-strict mode
font_styles = []
for span in spans:
font_style = span["font"].lower()
font_styles.append(font_style)
if len(font_styles) > 1:
longest_common_prefix = os.path.commonprefix(font_styles)
if len(longest_common_prefix) > 0:
return False
else:
return True
else:
return False
def __is_different_font_type_from_neighbors(curr_line_font_type, prev_line_font_type, next_line_font_type):
"""
This function checks if the current line has a different font type from the previous and next lines
Parameters
----------
curr_line_font_type : str
font type of the current line
prev_line_font_type : str
font type of the previous line
next_line_font_type : str
font type of the next line
Returns
-------
bool
True if the current line has a different font type from the previous and next lines, False otherwise.
"""
return all(
curr_line_font_type != other_font_type.lower()
for other_font_type in [prev_line_font_type, next_line_font_type]
if other_font_type is not None
)
def __is_larger_font_size_from_neighbors(curr_line_font_size, prev_line_font_size, next_line_font_size):
"""
This function checks if the current line has a larger font size than the previous and next lines
Parameters
----------
curr_line_font_size : float
font size of the current line
prev_line_font_size : float
font size of the previous line
next_line_font_size : float
font size of the next line
Returns
-------
bool
True if the current line has a larger font size than the previous and next lines, False otherwise.
"""
return all(
curr_line_font_size > other_font_size * 1.2
for other_font_size in [prev_line_font_size, next_line_font_size]
if other_font_size is not None
)
def __is_similar_to_pre_line(curr_line_font_type, prev_line_font_type, curr_line_font_size, prev_line_font_size):
"""
This function checks if the current line is similar to the previous line
Parameters
----------
curr_line : dict
current line
prev_line : dict
previous line
Returns
-------
bool
True if the current line is similar to the previous line, False otherwise.
"""
if curr_line_font_type == prev_line_font_type and curr_line_font_size == prev_line_font_size:
return True
else:
return False
def __is_same_font_type_of_docAvg(curr_line_font_type):
"""
This function checks if the current line has the same font type as the document average font type
Parameters
----------
curr_line_font_type : str
font type of the current line
Returns
-------
bool
True if the current line has the same font type as the document average font type, False otherwise.
"""
doc_most_common_font_type = safe_get(self.doc_statistics, "most_common_font_type", "").lower()
doc_second_most_common_font_type = safe_get(self.doc_statistics, "second_most_common_font_type", "").lower()
return curr_line_font_type.lower() in [doc_most_common_font_type, doc_second_most_common_font_type]
def __is_font_size_not_less_than_docAvg(curr_line_font_size, ratio: float = 1):
"""
This function checks if the current line has a large enough font size
Parameters
----------
curr_line_font_size : float
font size of the current line
ratio : float
ratio of the current line font size to the document average font size
Returns
-------
bool
True if the current line has a large enough font size, False otherwise.
"""
doc_most_common_font_size = safe_get(self.doc_statistics, "most_common_font_size", 0)
doc_second_most_common_font_size = safe_get(self.doc_statistics, "second_most_common_font_size", 0)
doc_avg_font_size = min(doc_most_common_font_size, doc_second_most_common_font_size)
return curr_line_font_size >= doc_avg_font_size * ratio
def __is_sufficient_spacing_above_and_below(
curr_line_bbox,
prev_line_bbox,
next_line_bbox,
avg_char_height,
median_font_size,
):
"""
This function checks if the current line has sufficient spacing above and below
Parameters
----------
curr_line_bbox : list
bbox of the current line
prev_line_bbox : list
bbox of the previous line
next_line_bbox : list
bbox of the next line
avg_char_width : float
average of char widths
avg_char_height : float
average of line heights
Returns
-------
bool
True if the current line has sufficient spacing above and below, False otherwise.
"""
vertical_ratio = 1.25
vertical_thres = vertical_ratio * median_font_size
_, y0, _, y1 = curr_line_bbox
sufficient_spacing_above = False
if prev_line_bbox:
vertical_spacing_above = min(y0 - prev_line_bbox[1], y1 - prev_line_bbox[3])
sufficient_spacing_above = vertical_spacing_above > vertical_thres
else:
sufficient_spacing_above = True
sufficient_spacing_below = False
if next_line_bbox:
vertical_spacing_below = min(next_line_bbox[1] - y0, next_line_bbox[3] - y1)
sufficient_spacing_below = vertical_spacing_below > vertical_thres
else:
sufficient_spacing_below = True
return (sufficient_spacing_above, sufficient_spacing_below)
def __is_word_list_line_by_rules(curr_line_text):
"""
This function checks if the current line is a word list
Parameters
----------
curr_line_text : str
text of the current line
Returns
-------
bool
True if the current line is a name list, False otherwise.
"""
# name_list_pattern = r"([a-zA-Z][a-zA-Z\s]{0,20}[a-zA-Z]|[\u4e00-\u9fa5·]{2,16})(?=[，,;；\s]|$)"
name_list_pattern = r"(?<![\u4e00-\u9fa5])([A-Z][a-z]{0,19}\s[A-Z][a-z]{0,19}|[\u4e00-\u9fa5]{2,6})(?=[，,;；\s]|$)"
compiled_pattern = re.compile(name_list_pattern)
if compiled_pattern.search(curr_line_text):
return True
else:
return False
def __get_text_catgr_by_nlp(curr_line_text):
"""
This function checks if the current line is a name list using nlp model, such as spacy
Parameters
----------
curr_line_text : str
text of the current line
Returns
-------
bool
True if the current line is a name list, False otherwise.
"""
result = self.nlp_model.detect_entity_catgr_using_nlp(curr_line_text)
return result
def __is_numbered_title(curr_line_text):
"""
This function checks if the current line is a numbered list
Parameters
----------
curr_line_text : str
text of the current line
Returns
-------
bool
True if the current line is a numbered list, False otherwise.
"""
compiled_pattern = re.compile(self.numbered_title_pattern, re.VERBOSE)
if compiled_pattern.search(curr_line_text):
return True
else:
return False
def __is_end_with_ending_puncs(line_text):
"""
This function checks if the current line ends with a ending punctuation mark
Parameters
----------
line_text : str
text of the current line
Returns
-------
bool
True if the current line ends with a punctuation mark, False otherwise.
"""
end_puncs = [".", "?", "!", "。", "？", "！", "…"]
line_text = line_text.rstrip()
if line_text[-1] in end_puncs:
return True
return False
def __contains_only_no_meaning_symbols(line_text):
"""
This function checks if the current line contains only symbols that have no meaning, if so, it is not a title.
Situation contains:
1. Only have punctuation marks
2. Only have other non-meaning symbols
Parameters
----------
line_text : str
text of the current line
Returns
-------
bool
True if the current line contains only symbols that have no meaning, False otherwise.
"""
punctuation_marks = re.findall(r"[^\w\s]", line_text) # find all punctuation marks
number_of_punctuation_marks = len(punctuation_marks)
text_length = len(line_text)
if text_length == 0:
return False
punctuation_ratio = number_of_punctuation_marks / text_length
if punctuation_ratio >= 0.9:
return True
return False
def __is_equation(line_text):
"""
This function checks if the current line is an equation.
Parameters
----------
line_text : str
Returns
-------
bool
True if the current line is an equation, False otherwise.
"""
equation_reg = r"\$.*?\\overline.*?\$" # to match interline equations
if re.search(equation_reg, line_text):
return True
else:
return False
def __is_title_by_len(text, max_length=200):
"""
This function checks if the current line is a title by length.
Parameters
----------
text : str
text of the current line
max_length : int
max length of the title
Returns
-------
bool
True if the current line is a title, False otherwise.
"""
text = text.strip()
return len(text) <= max_length
def __compute_line_font_type_and_size(curr_line):
"""
This function computes the font type and font size of the line.
Parameters
----------
line : dict
line
Returns
-------
font_type : str
font type of the line
font_size : float
font size of the line
"""
spans = curr_line["spans"]
max_accumulated_length = 0
max_span_font_size = curr_line["spans"][0]["size"] # default value, float type
max_span_font_type = curr_line["spans"][0]["font"].lower() # default value, string type
for span in spans:
if span["text"].isspace():
continue
span_length = span["bbox"][2] - span["bbox"][0]
if span_length > max_accumulated_length:
max_accumulated_length = span_length
max_span_font_size = span["size"]
max_span_font_type = span["font"].lower()
return max_span_font_type, max_span_font_size
def __is_a_consistent_sub_title(pre_line, curr_line):
"""
This function checks if the current line is a consistent sub title.
Parameters
----------
pre_line : dict
previous line
curr_line : dict
current line
Returns
-------
bool
True if the current line is a consistent sub title, False otherwise.
"""
if pre_line is None:
return False
start_letter_of_pre_line = pre_line["text"][0]
start_letter_of_curr_line = curr_line["text"][0]
has_same_prefix_digit = (
start_letter_of_pre_line.isdigit()
and start_letter_of_curr_line.isdigit()
and start_letter_of_pre_line == start_letter_of_curr_line
)
# prefix text of curr_line satisfies the following title format: x.x
prefix_text_pattern = r"^\d+\.\d+"
has_subtitle_format = re.match(prefix_text_pattern, curr_line["text"])
if has_same_prefix_digit or has_subtitle_format:
return True
"""
Title detecting main Process.
"""
"""
Basic features about the current line.
"""
curr_line_bbox = curr_line["bbox"]
curr_line_text = curr_line["text"]
curr_line_font_type, curr_line_font_size = __compute_line_font_type_and_size(curr_line)
if len(curr_line_text.strip()) == 0: # skip empty lines
return False, False
prev_line_bbox = prev_line["bbox"] if prev_line else None
if prev_line:
prev_line_font_type, prev_line_font_size = __compute_line_font_type_and_size(prev_line)
else:
prev_line_font_type, prev_line_font_size = None, None
next_line_bbox = next_line["bbox"] if next_line else None
if next_line:
next_line_font_type, next_line_font_size = __compute_line_font_type_and_size(next_line)
else:
next_line_font_type, next_line_font_size = None, None
"""
Aggregated features about the current line.
"""
is_italc_font = __is_italic_font_line(curr_line)
is_bold_font = __is_bold_font_line(curr_line)
is_font_size_little_less_than_doc_avg = __is_font_size_not_less_than_docAvg(curr_line_font_size, ratio=0.8)
is_font_size_not_less_than_doc_avg = __is_font_size_not_less_than_docAvg(curr_line_font_size, ratio=1)
is_much_larger_font_than_doc_avg = __is_font_size_not_less_than_docAvg(curr_line_font_size, ratio=1.6)
is_not_same_font_type_of_docAvg = not __is_same_font_type_of_docAvg(curr_line_font_type)
is_potential_title_font = is_bold_font or is_font_size_not_less_than_doc_avg or is_not_same_font_type_of_docAvg
is_mix_font_styles_strict = __has_mixed_font_styles(curr_line["spans"], strict_mode=True)
is_mix_font_styles_loose = __has_mixed_font_styles(curr_line["spans"], strict_mode=False)
is_punctuation_heavy = __is_punctuation_heavy(curr_line_text)
is_word_list_line_by_rules = __is_word_list_line_by_rules(curr_line_text)
is_person_or_org_list_line_by_nlp = __get_text_catgr_by_nlp(curr_line_text) in ["PERSON", "GPE", "ORG"]
is_font_size_larger_than_neighbors = __is_larger_font_size_from_neighbors(
curr_line_font_size, prev_line_font_size, next_line_font_size
)
is_font_type_diff_from_neighbors = __is_different_font_type_from_neighbors(
curr_line_font_type, prev_line_font_type, next_line_font_type
)
has_sufficient_spaces_above, has_sufficient_spaces_below = __is_sufficient_spacing_above_and_below(
curr_line_bbox, prev_line_bbox, next_line_bbox, avg_char_height, median_font_size
)
is_similar_to_pre_line = __is_similar_to_pre_line(
curr_line_font_type, prev_line_font_type, curr_line_font_size, prev_line_font_size
)
is_consis_sub_title = __is_a_consistent_sub_title(prev_line, curr_line)
"""
Further aggregated features about the current line.
Attention:
Features that start with __ are for internal use.
"""
__is_line_left_aligned_from_neighbors = is_line_left_aligned_from_neighbors(
curr_line_bbox, prev_line_bbox, next_line_bbox, avg_char_width
)
__is_font_diff_from_neighbors = is_font_size_larger_than_neighbors or is_font_type_diff_from_neighbors
is_a_left_inline_title = (
is_mix_font_styles_strict and __is_line_left_aligned_from_neighbors and __is_font_diff_from_neighbors
)
is_title_by_check_prev_line = prev_line is None and has_sufficient_spaces_above and is_potential_title_font
is_title_by_check_next_line = next_line is None and has_sufficient_spaces_below and is_potential_title_font
is_title_by_check_pre_and_next_line = (
(prev_line is not None or next_line is not None)
and has_sufficient_spaces_above
and has_sufficient_spaces_below
and is_potential_title_font
)
is_numbered_title = __is_numbered_title(curr_line_text) and (
(has_sufficient_spaces_above or prev_line is None) and (has_sufficient_spaces_below or next_line is None)
)
is_not_end_with_ending_puncs = not __is_end_with_ending_puncs(curr_line_text)
is_not_only_no_meaning_symbols = not __contains_only_no_meaning_symbols(curr_line_text)
is_equation = __is_equation(curr_line_text)
is_title_by_len = __is_title_by_len(curr_line_text)
"""
Decide if the line is a title.
"""
is_title = (
is_not_end_with_ending_puncs # not end with ending punctuation marks
and is_not_only_no_meaning_symbols # not only have no meaning symbols
and is_title_by_len # is a title by length, default max length is 200
and not is_equation # an interline equation should never be a title
and is_potential_title_font # is a potential title font, which is bold or larger than the document average font size or not the same font type as the document average font type
and (
(is_not_same_font_type_of_docAvg and is_font_size_not_less_than_doc_avg)
or (is_bold_font and is_much_larger_font_than_doc_avg and is_not_same_font_type_of_docAvg)
or (
is_much_larger_font_than_doc_avg
and (is_title_by_check_prev_line or is_title_by_check_next_line or is_title_by_check_pre_and_next_line)
)
or (
is_font_size_little_less_than_doc_avg
and is_bold_font
and (is_title_by_check_prev_line or is_title_by_check_next_line or is_title_by_check_pre_and_next_line)
)
) # Consider the following situations: bold font, much larger font than doc avg, not same font type as doc avg, sufficient spacing above and below
and (
(
not is_person_or_org_list_line_by_nlp
and (
is_much_larger_font_than_doc_avg
or (is_not_same_font_type_of_docAvg and is_font_size_not_less_than_doc_avg)
)
)
or (
not (is_word_list_line_by_rules and is_person_or_org_list_line_by_nlp)
and not is_a_left_inline_title
and not is_punctuation_heavy
and (is_title_by_check_prev_line or is_title_by_check_next_line or is_title_by_check_pre_and_next_line)
)
or (
is_person_or_org_list_line_by_nlp
and (is_bold_font and is_much_larger_font_than_doc_avg and is_not_same_font_type_of_docAvg)
and (is_bold_font and is_much_larger_font_than_doc_avg and is_not_same_font_type_of_docAvg)
)
or (is_numbered_title and not is_a_left_inline_title)
) # Exclude the following situations: person/org list
)
# ) or (prev_line_is_title and is_consis_sub_title)
is_name_or_org_list_to_be_removed = (
(is_person_or_org_list_line_by_nlp)
and is_punctuation_heavy
and (is_title_by_check_prev_line or is_title_by_check_next_line or is_title_by_check_pre_and_next_line)
) and not is_title
if is_name_or_org_list_to_be_removed:
is_author_or_org_list = True
else:
is_author_or_org_list = False
# return is_title, is_author_or_org_list
"""
# print reason why the line is a title
if is_title:
print_green("This line is a title.")
print_green("↓" * 10)
print()
print("curr_line_text: ", curr_line_text)
print()
# print reason why the line is not a title
line_text = curr_line_text.strip()
test_text = "Career/Personal Life"
text_content_condition = line_text == test_text
if not is_title and text_content_condition: # Print specific line
# if not is_title: # Print each line
print_red("This line is not a title.")
print_red("↓" * 10)
print()
print("curr_line_text: ", curr_line_text)
print()
if is_not_end_with_ending_puncs:
print_green(f"is_not_end_with_ending_puncs")
else:
print_red(f"is_end_with_ending_puncs")
if is_not_only_no_meaning_symbols:
print_green(f"is_not_only_no_meaning_symbols")
else:
print_red(f"is_only_no_meaning_symbols")
if is_title_by_len:
print_green(f"is_title_by_len: {is_title_by_len}")
else:
print_red(f"is_not_title_by_len: {is_title_by_len}")
if is_equation:
print_red(f"is_equation")
else:
print_green(f"is_not_equation")
if is_potential_title_font:
print_green(f"is_potential_title_font")
else:
print_red(f"is_not_potential_title_font")
if is_punctuation_heavy:
print_red("is_punctuation_heavy")
else:
print_green("is_not_punctuation_heavy")
if is_bold_font:
print_green(f"is_bold_font")
else:
print_red(f"is_not_bold_font")
if is_font_size_not_less_than_doc_avg:
print_green(f"is_larger_font_than_doc_avg")
else:
print_red(f"is_not_larger_font_than_doc_avg")
if is_much_larger_font_than_doc_avg:
print_green(f"is_much_larger_font_than_doc_avg")
else:
print_red(f"is_not_much_larger_font_than_doc_avg")
if is_not_same_font_type_of_docAvg:
print_green(f"is_not_same_font_type_of_docAvg")
else:
print_red(f"is_same_font_type_of_docAvg")
if is_word_list_line_by_rules:
print_red("is_word_list_line_by_rules")
else:
print_green("is_not_name_list_by_rules")
if is_person_or_org_list_line_by_nlp:
print_red("is_person_or_org_list_line_by_nlp")
else:
print_green("is_not_person_or_org_list_line_by_nlp")
if not is_numbered_title:
print_red("is_not_numbered_title")
else:
print_green("is_numbered_title")
if is_a_left_inline_title:
print_red("is_a_left_inline_title")
else:
print_green("is_not_a_left_inline_title")
if not is_title_by_check_prev_line:
print_red("is_not_title_by_check_prev_line")
else:
print_green("is_title_by_check_prev_line")
if not is_title_by_check_next_line:
print_red("is_not_title_by_check_next_line")
else:
print_green("is_title_by_check_next_line")
if not is_title_by_check_pre_and_next_line:
print_red("is_not_title_by_check_pre_and_next_line")
else:
print_green("is_title_by_check_pre_and_next_line")
# print_green("Common features:")
# print_green("↓" * 10)
# print(f" curr_line_font_type: {curr_line_font_type}")
# print(f" curr_line_font_size: {curr_line_font_size}")
# print()
"""
return is_title, is_author_or_org_list
def _detect_title(self, input_block):
"""
Use the functions 'is_potential_title' to detect titles of each paragraph block.
If a line is a title, then the value of key 'is_title' of the line will be set to True.
"""
raw_lines = input_block["lines"]
prev_line_is_title_flag = False
for i, curr_line in enumerate(raw_lines):
prev_line = raw_lines[i - 1] if i > 0 else None
next_line = raw_lines[i + 1] if i < len(raw_lines) - 1 else None
blk_avg_char_width = input_block["avg_char_width"]
blk_avg_char_height = input_block["avg_char_height"]
blk_media_font_size = input_block["median_font_size"]
is_title, is_author_or_org_list = self._is_potential_title(
curr_line,
prev_line,
prev_line_is_title_flag,
next_line,
blk_avg_char_width,
blk_avg_char_height,
blk_media_font_size,
)
if is_title:
curr_line["is_title"] = is_title
prev_line_is_title_flag = True
else:
curr_line["is_title"] = False
prev_line_is_title_flag = False
# print(f"curr_line['text']: {curr_line['text']}")
# print(f"curr_line['is_title']: {curr_line['is_title']}")
# print(f"prev_line['text']: {prev_line['text'] if prev_line else None}")
# print(f"prev_line_is_title_flag: {prev_line_is_title_flag}")
# print()
if is_author_or_org_list:
curr_line["is_author_or_org_list"] = is_author_or_org_list
else:
curr_line["is_author_or_org_list"] = False
return input_block
def batch_detect_titles(self, pdf_dic):
"""
This function batch process the blocks to detect titles.
Parameters
----------
pdf_dict : dict
result dictionary
Returns
-------
pdf_dict : dict
result dictionary
"""
num_titles = 0
for page_id, blocks in pdf_dic.items():
if page_id.startswith("page_"):
para_blocks = []
if "para_blocks" in blocks.keys():
para_blocks = blocks["para_blocks"]
all_single_line_blocks = []
for block in para_blocks:
if len(block["lines"]) == 1:
all_single_line_blocks.append(block)
new_para_blocks = []
if not len(all_single_line_blocks) == len(para_blocks): # Not all blocks are single line blocks.
for para_block in para_blocks:
new_block = self._detect_title(para_block)
new_para_blocks.append(new_block)
num_titles += sum([line.get("is_title", 0) for line in new_block["lines"]])
else: # All blocks are single line blocks.
for para_block in para_blocks:
new_para_blocks.append(para_block)
num_titles += sum([line.get("is_title", 0) for line in para_block["lines"]])
para_blocks = new_para_blocks
blocks["para_blocks"] = para_blocks
for para_block in para_blocks:
all_titles = all(safe_get(line, "is_title", False) for line in para_block["lines"])
para_text_len = sum([len(line["text"]) for line in para_block["lines"]])
if (
all_titles and para_text_len < 200
): # total length of the paragraph is less than 200, more than this should not be a title
para_block["is_block_title"] = 1
else:
para_block["is_block_title"] = 0
all_name_or_org_list_to_be_removed = all(
safe_get(line, "is_author_or_org_list", False) for line in para_block["lines"]
)
if all_name_or_org_list_to_be_removed and page_id == "page_0":
para_block["is_block_an_author_or_org_list"] = 1
else:
para_block["is_block_an_author_or_org_list"] = 0
pdf_dic["statistics"]["num_titles"] = num_titles
return pdf_dic
def _recog_title_level(self, title_blocks):
"""
This function determines the title level based on the font size of the title.
Parameters
----------
title_blocks : list
Returns
-------
title_blocks : list
"""
font_sizes = np.array([safe_get(tb["block"], "block_font_size", 0) for tb in title_blocks])
# Use the mean and std of font sizes to remove extreme values
mean_font_size = np.mean(font_sizes)
std_font_size = np.std(font_sizes)
min_extreme_font_size = mean_font_size - std_font_size # type: ignore
max_extreme_font_size = mean_font_size + std_font_size # type: ignore
# Compute the threshold for title level
middle_font_sizes = font_sizes[(font_sizes > min_extreme_font_size) & (font_sizes < max_extreme_font_size)]
if middle_font_sizes.size > 0:
middle_mean_font_size = np.mean(middle_font_sizes)
level_threshold = middle_mean_font_size
else:
level_threshold = mean_font_size
for tb in title_blocks:
title_block = tb["block"]
title_font_size = safe_get(title_block, "block_font_size", 0)
current_level = 1 # Initialize title level, the biggest level is 1
# print(f"Before adjustment by font size, {current_level}")
if title_font_size >= max_extreme_font_size:
current_level = 1
elif title_font_size <= min_extreme_font_size:
current_level = 3
elif float(title_font_size) >= float(level_threshold):
current_level = 2
else:
current_level = 3
# print(f"After adjustment by font size, {current_level}")
title_block["block_title_level"] = current_level
return title_blocks
def batch_recog_title_level(self, pdf_dic):
"""
This function batch process the blocks to recognize title level.
Parameters
----------
pdf_dict : dict
result dictionary
Returns
-------
pdf_dict : dict
result dictionary
"""
title_blocks = []
# Collect all titles
for page_id, blocks in pdf_dic.items():
if page_id.startswith("page_"):
para_blocks = blocks.get("para_blocks", [])
for block in para_blocks:
if block.get("is_block_title"):
title_obj = {"page_id": page_id, "block": block}
title_blocks.append(title_obj)
# Determine title level
if title_blocks:
# Determine title level based on font size
title_blocks = self._recog_title_level(title_blocks)
return pdf_dic
class BlockTerminationProcessor:
"""
This class is used to process the block termination.
"""
def __init__(self) -> None:
pass
def _is_consistent_lines(
self,
curr_line,
prev_line,
next_line,
consistent_direction, # 0 for prev, 1 for next, 2 for both
):
"""
This function checks if the line is consistent with its neighbors
Parameters
----------
curr_line : dict
current line
prev_line : dict
previous line
next_line : dict
next line
consistent_direction : int
0 for prev, 1 for next, 2 for both
Returns
-------
bool
True if the line is consistent with its neighbors, False otherwise.
"""
curr_line_font_size = curr_line["spans"][0]["size"]
curr_line_font_type = curr_line["spans"][0]["font"].lower()
if consistent_direction == 0:
if prev_line:
prev_line_font_size = prev_line["spans"][0]["size"]
prev_line_font_type = prev_line["spans"][0]["font"].lower()
return curr_line_font_size == prev_line_font_size and curr_line_font_type == prev_line_font_type
else:
return False
elif consistent_direction == 1:
if next_line:
next_line_font_size = next_line["spans"][0]["size"]
next_line_font_type = next_line["spans"][0]["font"].lower()
return curr_line_font_size == next_line_font_size and curr_line_font_type == next_line_font_type
else:
return False
elif consistent_direction == 2:
if prev_line and next_line:
prev_line_font_size = prev_line["spans"][0]["size"]
prev_line_font_type = prev_line["spans"][0]["font"].lower()
next_line_font_size = next_line["spans"][0]["size"]
next_line_font_type = next_line["spans"][0]["font"].lower()
return (curr_line_font_size == prev_line_font_size and curr_line_font_type == prev_line_font_type) and (
curr_line_font_size == next_line_font_size and curr_line_font_type == next_line_font_type
)
else:
return False
else:
return False
def _is_regular_line(self, curr_line_bbox, prev_line_bbox, next_line_bbox, avg_char_width, X0, X1, avg_line_height):
"""
This function checks if the line is a regular line
Parameters
----------
curr_line_bbox : list
bbox of the current line
prev_line_bbox : list
bbox of the previous line
next_line_bbox : list
bbox of the next line
avg_char_width : float
average of char widths
X0 : float
median of x0 values, which represents the left average boundary of the page
X1 : float
median of x1 values, which represents the right average boundary of the page
avg_line_height : float
average of line heights
Returns
-------
bool
True if the line is a regular line, False otherwise.
"""
horizontal_ratio = 0.5
vertical_ratio = 0.5
horizontal_thres = horizontal_ratio * avg_char_width
vertical_thres = vertical_ratio * avg_line_height
x0, y0, x1, y1 = curr_line_bbox
x0_near_X0 = abs(x0 - X0) < horizontal_thres
x1_near_X1 = abs(x1 - X1) < horizontal_thres
prev_line_is_end_of_para = prev_line_bbox and (abs(prev_line_bbox[2] - X1) > avg_char_width)
sufficient_spacing_above = False
if prev_line_bbox:
vertical_spacing_above = y1 - prev_line_bbox[3]
sufficient_spacing_above = vertical_spacing_above > vertical_thres
sufficient_spacing_below = False
if next_line_bbox:
vertical_spacing_below = next_line_bbox[1] - y0
sufficient_spacing_below = vertical_spacing_below > vertical_thres
return (
(sufficient_spacing_above or sufficient_spacing_below)
or (not x0_near_X0 and not x1_near_X1)
or prev_line_is_end_of_para
)
def _is_possible_start_of_para(self, curr_line, prev_line, next_line, X0, X1, avg_char_width, avg_font_size):
"""
This function checks if the line is a possible start of a paragraph
Parameters
----------
curr_line : dict
current line
prev_line : dict
previous line
next_line : dict
next line
X0 : float
median of x0 values, which represents the left average boundary of the page
X1 : float
median of x1 values, which represents the right average boundary of the page
avg_char_width : float
average of char widths
avg_line_height : float
average of line heights
Returns
-------
bool
True if the line is a possible start of a paragraph, False otherwise.
"""
start_confidence = 0.5 # Initial confidence of the line being a start of a paragraph
decision_path = [] # Record the decision path
curr_line_bbox = curr_line["bbox"]
prev_line_bbox = prev_line["bbox"] if prev_line else None
next_line_bbox = next_line["bbox"] if next_line else None
indent_ratio = 1
vertical_ratio = 1.5
vertical_thres = vertical_ratio * avg_font_size
left_horizontal_ratio = 0.5
left_horizontal_thres = left_horizontal_ratio * avg_char_width
right_horizontal_ratio = 2.5
right_horizontal_thres = right_horizontal_ratio * avg_char_width
x0, y0, x1, y1 = curr_line_bbox
indent_condition = x0 > X0 + indent_ratio * avg_char_width
if indent_condition:
start_confidence += 0.2
decision_path.append("indent_condition_met")
x0_near_X0 = abs(x0 - X0) < left_horizontal_thres
if x0_near_X0:
start_confidence += 0.1
decision_path.append("x0_near_X0")
x1_near_X1 = abs(x1 - X1) < right_horizontal_thres
if x1_near_X1:
start_confidence += 0.1
decision_path.append("x1_near_X1")
if prev_line is None:
prev_line_is_end_of_para = True
start_confidence += 0.2
decision_path.append("no_prev_line")
else:
prev_line_is_end_of_para, _, _ = self._is_possible_end_of_para(prev_line, next_line, X0, X1, avg_char_width)
if prev_line_is_end_of_para:
start_confidence += 0.1
decision_path.append("prev_line_is_end_of_para")
sufficient_spacing_above = False
if prev_line_bbox:
vertical_spacing_above = y1 - prev_line_bbox[3]
sufficient_spacing_above = vertical_spacing_above > vertical_thres
if sufficient_spacing_above:
start_confidence += 0.2
decision_path.append("sufficient_spacing_above")
sufficient_spacing_below = False
if next_line_bbox:
vertical_spacing_below = next_line_bbox[1] - y0
sufficient_spacing_below = vertical_spacing_below > vertical_thres
if sufficient_spacing_below:
start_confidence += 0.2
decision_path.append("sufficient_spacing_below")
is_regular_line = self._is_regular_line(
curr_line_bbox, prev_line_bbox, next_line_bbox, avg_char_width, X0, X1, avg_font_size
)
if is_regular_line:
start_confidence += 0.1
decision_path.append("is_regular_line")
is_start_of_para = (
(sufficient_spacing_above or sufficient_spacing_below)
or (indent_condition)
or (not indent_condition and x0_near_X0 and x1_near_X1 and not is_regular_line)
or prev_line_is_end_of_para
)
return (is_start_of_para, start_confidence, decision_path)
def _is_possible_end_of_para(self, curr_line, next_line, X0, X1, avg_char_width):
"""
This function checks if the line is a possible end of a paragraph
Parameters
----------
curr_line : dict
current line
next_line : dict
next line
X0 : float
median of x0 values, which represents the left average boundary of the page
X1 : float
median of x1 values, which represents the right average boundary of the page
avg_char_width : float
average of char widths
Returns
-------
bool
True if the line is a possible end of a paragraph, False otherwise.
"""
end_confidence = 0.5 # Initial confidence of the line being a end of a paragraph
decision_path = [] # Record the decision path
curr_line_bbox = curr_line["bbox"]
next_line_bbox = next_line["bbox"] if next_line else None
left_horizontal_ratio = 0.5
right_horizontal_ratio = 0.5
x0, _, x1, y1 = curr_line_bbox
next_x0, next_y0, _, _ = next_line_bbox if next_line_bbox else (0, 0, 0, 0)
x0_near_X0 = abs(x0 - X0) < left_horizontal_ratio * avg_char_width
if x0_near_X0:
end_confidence += 0.1
decision_path.append("x0_near_X0")
x1_smaller_than_X1 = x1 < X1 - right_horizontal_ratio * avg_char_width
if x1_smaller_than_X1:
end_confidence += 0.1
decision_path.append("x1_smaller_than_X1")
next_line_is_start_of_para = (
next_line_bbox
and (next_x0 > X0 + left_horizontal_ratio * avg_char_width)
and (not is_line_left_aligned_from_neighbors(curr_line_bbox, None, next_line_bbox, avg_char_width, direction=1))
)
if next_line_is_start_of_para:
end_confidence += 0.2
decision_path.append("next_line_is_start_of_para")
is_line_left_aligned_from_neighbors_bool = is_line_left_aligned_from_neighbors(
curr_line_bbox, None, next_line_bbox, avg_char_width
)
if is_line_left_aligned_from_neighbors_bool:
end_confidence += 0.1
decision_path.append("line_is_left_aligned_from_neighbors")
is_line_right_aligned_from_neighbors_bool = is_line_right_aligned_from_neighbors(
curr_line_bbox, None, next_line_bbox, avg_char_width
)
if not is_line_right_aligned_from_neighbors_bool:
end_confidence += 0.1
decision_path.append("line_is_not_right_aligned_from_neighbors")
is_end_of_para = end_with_punctuation(curr_line["text"]) and (
(x0_near_X0 and x1_smaller_than_X1)
or (is_line_left_aligned_from_neighbors_bool and not is_line_right_aligned_from_neighbors_bool)
)
return (is_end_of_para, end_confidence, decision_path)
def _cut_paras_per_block(
self,
block,
):
"""
Processes a raw block from PyMuPDF and returns the processed block.
Parameters
----------
raw_block : dict
A raw block from pymupdf.
Returns
-------
processed_block : dict
"""
def _construct_para(lines, is_block_title, para_title_level):
"""
Construct a paragraph from given lines.
"""
font_sizes = [span["size"] for line in lines for span in line["spans"]]
avg_font_size = sum(font_sizes) / len(font_sizes) if font_sizes else 0
font_colors = [span["color"] for line in lines for span in line["spans"]]
most_common_font_color = max(set(font_colors), key=font_colors.count) if font_colors else None
font_type_lengths = {}
for line in lines:
for span in line["spans"]:
font_type = span["font"]
bbox_width = span["bbox"][2] - span["bbox"][0]
if font_type in font_type_lengths:
font_type_lengths[font_type] += bbox_width
else:
font_type_lengths[font_type] = bbox_width
# get the font type with the longest bbox width
most_common_font_type = max(font_type_lengths, key=font_type_lengths.get) if font_type_lengths else None # type: ignore
para_bbox = calculate_para_bbox(lines)
para_text = " ".join(line["text"] for line in lines)
return {
"para_bbox": para_bbox,
"para_text": para_text,
"para_font_type": most_common_font_type,
"para_font_size": avg_font_size,
"para_font_color": most_common_font_color,
"is_para_title": is_block_title,
"para_title_level": para_title_level,
}
block_bbox = block["bbox"]
block_text = block["text"]
block_lines = block["lines"]
X0 = safe_get(block, "X0", 0)
X1 = safe_get(block, "X1", 0)
avg_char_width = safe_get(block, "avg_char_width", 0)
avg_char_height = safe_get(block, "avg_char_height", 0)
avg_font_size = safe_get(block, "avg_font_size", 0)
is_block_title = safe_get(block, "is_block_title", False)
para_title_level = safe_get(block, "block_title_level", 0)
# Segment into paragraphs
para_ranges = []
in_paragraph = False
start_idx_of_para = None
# Create the processed paragraphs
processed_paras = {}
para_bboxes = []
end_idx_of_para = 0
for line_index, line in enumerate(block_lines):
curr_line = line
prev_line = block_lines[line_index - 1] if line_index > 0 else None
next_line = block_lines[line_index + 1] if line_index < len(block_lines) - 1 else None
"""
Start processing paragraphs.
"""
# Check if the line is the start of a paragraph
is_start_of_para, start_confidence, decision_path = self._is_possible_start_of_para(
curr_line, prev_line, next_line, X0, X1, avg_char_width, avg_font_size
)
if not in_paragraph and is_start_of_para:
in_paragraph = True
start_idx_of_para = line_index
# print_green(">>> Start of a paragraph")
# print(" curr_line_text: ", curr_line["text"])
# print(" start_confidence: ", start_confidence)
# print(" decision_path: ", decision_path)
# Check if the line is the end of a paragraph
is_end_of_para, end_confidence, decision_path = self._is_possible_end_of_para(
curr_line, next_line, X0, X1, avg_char_width
)
if in_paragraph and (is_end_of_para or not next_line):
para_ranges.append((start_idx_of_para, line_index))
start_idx_of_para = None
in_paragraph = False
# print_red(">>> End of a paragraph")
# print(" curr_line_text: ", curr_line["text"])
# print(" end_confidence: ", end_confidence)
# print(" decision_path: ", decision_path)
# Add the last paragraph if it is not added
if in_paragraph and start_idx_of_para is not None:
para_ranges.append((start_idx_of_para, len(block_lines) - 1))
# Process the matched paragraphs
for para_index, (start_idx, end_idx) in enumerate(para_ranges):
matched_lines = block_lines[start_idx : end_idx + 1]
para_properties = _construct_para(matched_lines, is_block_title, para_title_level)
para_key = f"para_{len(processed_paras)}"
processed_paras[para_key] = para_properties
para_bboxes.append(para_properties["para_bbox"])
end_idx_of_para = end_idx + 1
# Deal with the remaining lines
if end_idx_of_para < len(block_lines):
unmatched_lines = block_lines[end_idx_of_para:]
unmatched_properties = _construct_para(unmatched_lines, is_block_title, para_title_level)
unmatched_key = f"para_{len(processed_paras)}"
processed_paras[unmatched_key] = unmatched_properties
para_bboxes.append(unmatched_properties["para_bbox"])
block["paras"] = processed_paras
return block
def batch_process_blocks(self, pdf_dict):
"""
Parses the blocks of all pages.
Parameters
----------
pdf_dict : dict
PDF dictionary.
filter_blocks : list
List of bounding boxes to filter.
Returns
-------
result_dict : dict
Result dictionary.
"""
num_paras = 0
for page_id, page in pdf_dict.items():
if page_id.startswith("page_"):
para_blocks = []
if "para_blocks" in page.keys():
input_blocks = page["para_blocks"]
for input_block in input_blocks:
new_block = self._cut_paras_per_block(input_block)
para_blocks.append(new_block)
num_paras += len(new_block["paras"])
page["para_blocks"] = para_blocks
pdf_dict["statistics"]["num_paras"] = num_paras
return pdf_dict
class BlockContinuationProcessor:
"""
This class is used to process the blocks to detect block continuations.
"""
def __init__(self) -> None:
pass
def __is_similar_font_type(self, font_type_1, font_type_2, prefix_length_ratio=0.3):
"""
This function checks if the two font types are similar.
Definition of similar font types: the two font types have a common prefix,
and the length of the common prefix is at least a certain ratio of the length of the shorter font type.
Parameters
----------
font_type1 : str
font type 1
font_type2 : str
font type 2
prefix_length_ratio : float
minimum ratio of the common prefix length to the length of the shorter font type
Returns
-------
bool
True if the two font types are similar, False otherwise.
"""
if isinstance(font_type_1, list):
font_type_1 = font_type_1[0] if font_type_1 else ""
if isinstance(font_type_2, list):
font_type_2 = font_type_2[0] if font_type_2 else ""
if font_type_1 == font_type_2:
return True
# Find the length of the common prefix
common_prefix_length = len(os.path.commonprefix([font_type_1, font_type_2]))
# Calculate the minimum prefix length based on the ratio
min_prefix_length = int(min(len(font_type_1), len(font_type_2)) * prefix_length_ratio)
return common_prefix_length >= min_prefix_length
def __is_same_block_font(self, block_1, block_2):
"""
This function compares the font of block1 and block2
Parameters
----------
block1 : dict
block1
block2 : dict
block2
Returns
-------
is_same : bool
True if block1 and block2 have the same font, else False
"""
block_1_font_type = safe_get(block_1, "block_font_type", "")
block_1_font_size = safe_get(block_1, "block_font_size", 0)
block_1_avg_char_width = safe_get(block_1, "avg_char_width", 0)
block_2_font_type = safe_get(block_2, "block_font_type", "")
block_2_font_size = safe_get(block_2, "block_font_size", 0)
block_2_avg_char_width = safe_get(block_2, "avg_char_width", 0)
if isinstance(block_1_font_size, list):
block_1_font_size = block_1_font_size[0] if block_1_font_size else 0
if isinstance(block_2_font_size, list):
block_2_font_size = block_2_font_size[0] if block_2_font_size else 0
block_1_text = safe_get(block_1, "text", "")
block_2_text = safe_get(block_2, "text", "")
if block_1_avg_char_width == 0 or block_2_avg_char_width == 0:
return False
if not block_1_text or not block_2_text:
return False
else:
text_len_ratio = len(block_2_text) / len(block_1_text)
if text_len_ratio < 0.2:
avg_char_width_condition = (
abs(block_1_avg_char_width - block_2_avg_char_width) / min(block_1_avg_char_width, block_2_avg_char_width)
< 0.5
)
else:
avg_char_width_condition = (
abs(block_1_avg_char_width - block_2_avg_char_width) / min(block_1_avg_char_width, block_2_avg_char_width)
< 0.2
)
block_font_size_condition = abs(block_1_font_size - block_2_font_size) < 1
return (
self.__is_similar_font_type(block_1_font_type, block_2_font_type)
and avg_char_width_condition
and block_font_size_condition
)
def _is_alphabet_char(self, char):
if (char >= "\u0041" and char <= "\u005a") or (char >= "\u0061" and char <= "\u007a"):
return True
else:
return False
def _is_chinese_char(self, char):
if char >= "\u4e00" and char <= "\u9fa5":
return True
else:
return False
def _is_other_letter_char(self, char):
try:
cat = unicodedata.category(char)
if cat == "Lu" or cat == "Ll":
return not self._is_alphabet_char(char) and not self._is_chinese_char(char)
except TypeError:
print("The input to the function must be a single character.")
return False
def _is_year(self, s: str):
try:
number = int(s)
return 1900 <= number <= 2099
except ValueError:
return False
def _match_brackets(self, text):
# pattern = r"^[\(\)\[\]（）【】{}｛｝<>＜＞〔〕〘〙\"\'“”‘’]"
pattern = r"^[\(\)\]（）】{}｛｝>＞〕〙\"\'“”‘’]"
return bool(re.match(pattern, text))
def _is_para_font_consistent(self, para_1, para_2):
"""
This function compares the font of para1 and para2
Parameters
----------
para1 : dict
para1
para2 : dict
para2
Returns
-------
is_same : bool
True if para1 and para2 have the same font, else False
"""
if para_1 is None or para_2 is None:
return False
para_1_font_type = safe_get(para_1, "para_font_type", "")
para_1_font_size = safe_get(para_1, "para_font_size", 0)
para_1_font_color = safe_get(para_1, "para_font_color", "")
para_2_font_type = safe_get(para_2, "para_font_type", "")
para_2_font_size = safe_get(para_2, "para_font_size", 0)
para_2_font_color = safe_get(para_2, "para_font_color", "")
if isinstance(para_1_font_type, list): # get the most common font type
para_1_font_type = max(set(para_1_font_type), key=para_1_font_type.count)
if isinstance(para_2_font_type, list):
para_2_font_type = max(set(para_2_font_type), key=para_2_font_type.count)
if isinstance(para_1_font_size, list): # compute average font type
para_1_font_size = sum(para_1_font_size) / len(para_1_font_size)
if isinstance(para_2_font_size, list): # compute average font type
para_2_font_size = sum(para_2_font_size) / len(para_2_font_size)
return (
self.__is_similar_font_type(para_1_font_type, para_2_font_type)
and abs(para_1_font_size - para_2_font_size) < 1.5
# and para_font_color1 == para_font_color2
)
def _is_para_puncs_consistent(self, para_1, para_2):
"""
This function determines whether para1 and para2 are originally from the same paragraph by checking the puncs of para1(former) and para2(latter)
Parameters
----------
para1 : dict
para1
para2 : dict
para2
Returns
-------
is_same : bool
True if para1 and para2 are from the same paragraph by using the puncs, else False
"""
para_1_text = safe_get(para_1, "para_text", "").strip()
para_2_text = safe_get(para_2, "para_text", "").strip()
para_1_bboxes = safe_get(para_1, "para_bbox", [])
para_1_font_sizes = safe_get(para_1, "para_font_size", 0)
para_2_bboxes = safe_get(para_2, "para_bbox", [])
para_2_font_sizes = safe_get(para_2, "para_font_size", 0)
# print_yellow(" Features of determine puncs_consistent:")
# print(f" para_1_text: {para_1_text}")
# print(f" para_2_text: {para_2_text}")
# print(f" para_1_bboxes: {para_1_bboxes}")
# print(f" para_2_bboxes: {para_2_bboxes}")
# print(f" para_1_font_sizes: {para_1_font_sizes}")
# print(f" para_2_font_sizes: {para_2_font_sizes}")
if is_nested_list(para_1_bboxes):
x0_1, y0_1, x1_1, y1_1 = para_1_bboxes[-1]
else:
x0_1, y0_1, x1_1, y1_1 = para_1_bboxes
if is_nested_list(para_2_bboxes):
x0_2, y0_2, x1_2, y1_2 = para_2_bboxes[0]
para_2_font_sizes = para_2_font_sizes[0] # type: ignore
else:
x0_2, y0_2, x1_2, y1_2 = para_2_bboxes
right_align_threshold = 0.5 * (para_1_font_sizes + para_2_font_sizes) * 0.8
are_two_paras_right_aligned = abs(x1_1 - x1_2) < right_align_threshold
left_indent_threshold = 0.5 * (para_1_font_sizes + para_2_font_sizes) * 0.8
is_para1_left_indent_than_papa2 = x0_1 - x0_2 > left_indent_threshold
is_para2_left_indent_than_papa1 = x0_2 - x0_1 > left_indent_threshold
# Check if either para_text1 or para_text2 is empty
if not para_1_text or not para_2_text:
return False
# Define the end puncs for a sentence to end and hyphen
end_puncs = [".", "?", "!", "。", "？", "！", "…"]
hyphen = ["-", "—"]
# Check if para_text1 ends with either hyphen or non-end punctuation or spaces
para_1_end_with_hyphen = para_1_text and para_1_text[-1] in hyphen
para_1_end_with_end_punc = para_1_text and para_1_text[-1] in end_puncs
para_1_end_with_space = para_1_text and para_1_text[-1] == " "
para_1_not_end_with_end_punc = para_1_text and para_1_text[-1] not in end_puncs
# print_yellow(f" para_1_end_with_hyphen: {para_1_end_with_hyphen}")
# print_yellow(f" para_1_end_with_end_punc: {para_1_end_with_end_punc}")
# print_yellow(f" para_1_not_end_with_end_punc: {para_1_not_end_with_end_punc}")
# print_yellow(f" para_1_end_with_space: {para_1_end_with_space}")
if para_1_end_with_hyphen: # If para_text1 ends with hyphen
# print_red(f"para_1 is end with hyphen.")
para_2_is_consistent = para_2_text and (
para_2_text[0] in hyphen
or (self._is_alphabet_char(para_2_text[0]) and para_2_text[0].islower())
or (self._is_chinese_char(para_2_text[0]))
or (self._is_other_letter_char(para_2_text[0]))
)
if para_2_is_consistent:
# print(f"para_2 is consistent.\n")
return True
else:
# print(f"para_2 is not consistent.\n")
pass
elif para_1_end_with_end_punc: # If para_text1 ends with ending punctuations
# print_red(f"para_1 is end with end_punc.")
para_2_is_consistent = (
para_2_text
and (
para_2_text[0]
== " "
# or (self._is_alphabet_char(para_2_text[0]) and para_2_text[0].isupper())
# or (self._is_chinese_char(para_2_text[0]))
# or (self._is_other_letter_char(para_2_text[0]))
)
and not is_para2_left_indent_than_papa1
)
if para_2_is_consistent:
# print(f"para_2 is consistent.\n")
return True
else:
# print(f"para_2 is not consistent.\n")
pass
elif para_1_not_end_with_end_punc: # If para_text1 is not end with ending punctuations
# print_red(f"para_1 is NOT end with end_punc.")
para_2_is_consistent = para_2_text and (
para_2_text[0] == " "
or (self._is_alphabet_char(para_2_text[0]) and para_2_text[0].islower())
or (self._is_alphabet_char(para_2_text[0]))
or (self._is_year(para_2_text[0:4]))
or (are_two_paras_right_aligned or is_para1_left_indent_than_papa2)
or (self._is_chinese_char(para_2_text[0]))
or (self._is_other_letter_char(para_2_text[0]))
or (self._match_brackets(para_2_text[0]))
)
if para_2_is_consistent:
# print(f"para_2 is consistent.\n")
return True
else:
# print(f"para_2 is not consistent.\n")
pass
elif para_1_end_with_space: # If para_text1 ends with space
# print_red(f"para_1 is end with space.")
para_2_is_consistent = para_2_text and (
para_2_text[0] == " "
or (self._is_alphabet_char(para_2_text[0]) and para_2_text[0].islower())
or (self._is_chinese_char(para_2_text[0]))
or (self._is_other_letter_char(para_2_text[0]))
)
if para_2_is_consistent:
# print(f"para_2 is consistent.\n")
return True
else:
pass
# print(f"para_2 is not consistent.\n")
return False
def _is_block_consistent(self, block_1, block_2):
"""
This function determines whether block1 and block2 are originally from the same block
Parameters
----------
block1 : dict
block1s
block2 : dict
block2
Returns
-------
is_same : bool
True if block1 and block2 are from the same block, else False
"""
return self.__is_same_block_font(block_1, block_2)
def _is_para_continued(self, para_1, para_2):
"""
This function determines whether para1 and para2 are originally from the same paragraph
Parameters
----------
para1 : dict
para1
para2 : dict
para2
Returns
-------
is_same : bool
True if para1 and para2 are from the same paragraph, else False
"""
is_para_font_consistent = self._is_para_font_consistent(para_1, para_2)
is_para_puncs_consistent = self._is_para_puncs_consistent(para_1, para_2)
return is_para_font_consistent and is_para_puncs_consistent
def _are_boundaries_of_block_consistent(self, block_1, block_2):
"""
This function checks if the boundaries of block1 and block2 are consistent
Parameters
----------
block1 : dict
block1
block2 : dict
block2
Returns
-------
is_consistent : bool
True if the boundaries of block1 and block2 are consistent, else False
"""
last_line_of_block_1 = block_1["lines"][-1]
first_line_of_block_2 = block_2["lines"][0]
spans_of_last_line_of_block_1 = last_line_of_block_1["spans"]
spans_of_first_line_of_block_2 = first_line_of_block_2["spans"]
font_type_of_last_line_of_block_1 = spans_of_last_line_of_block_1[0]["font"].lower()
font_size_of_last_line_of_block_1 = spans_of_last_line_of_block_1[0]["size"]
font_color_of_last_line_of_block_1 = spans_of_last_line_of_block_1[0]["color"]
font_flags_of_last_line_of_block_1 = spans_of_last_line_of_block_1[0]["flags"]
font_type_of_first_line_of_block_2 = spans_of_first_line_of_block_2[0]["font"].lower()
font_size_of_first_line_of_block_2 = spans_of_first_line_of_block_2[0]["size"]
font_color_of_first_line_of_block_2 = spans_of_first_line_of_block_2[0]["color"]
font_flags_of_first_line_of_block_2 = spans_of_first_line_of_block_2[0]["flags"]
return (
self.__is_similar_font_type(font_type_of_last_line_of_block_1, font_type_of_first_line_of_block_2)
and abs(font_size_of_last_line_of_block_1 - font_size_of_first_line_of_block_2) < 1
# and font_color_of_last_line_of_block1 == font_color_of_first_line_of_block2
and font_flags_of_last_line_of_block_1 == font_flags_of_first_line_of_block_2
)
def should_merge_next_para(self, curr_para, next_para):
"""
This function checks if the next_para should be merged into the curr_para.
Parameters
----------
curr_para : dict
The current paragraph.
next_para : dict
The next paragraph.
Returns
-------
bool
True if the next_para should be merged into the curr_para, False otherwise.
"""
if self._is_para_continued(curr_para, next_para):
return True
else:
return False
def batch_tag_paras(self, pdf_dict):
"""
This function tags the paragraphs in the pdf_dict.
Parameters
----------
pdf_dict : dict
PDF dictionary.
Returns
-------
pdf_dict : dict
PDF dictionary with tagged paragraphs.
"""
the_last_page_id = len(pdf_dict) - 1
for curr_page_idx, (curr_page_id, curr_page_content) in enumerate(pdf_dict.items()):
if curr_page_id.startswith("page_") and curr_page_content.get("para_blocks", []):
para_blocks_of_curr_page = curr_page_content["para_blocks"]
next_page_idx = curr_page_idx + 1
next_page_id = f"page_{next_page_idx}"
next_page_content = pdf_dict.get(next_page_id, {})
for i, current_block in enumerate(para_blocks_of_curr_page):
for para_id, curr_para in current_block["paras"].items():
curr_para["curr_para_location"] = [
curr_page_idx,
current_block["block_id"],
int(para_id.split("_")[-1]),
]
curr_para["next_para_location"] = None # 默认设置为None
curr_para["merge_next_para"] = False # 默认设置为False
next_block = para_blocks_of_curr_page[i + 1] if i < len(para_blocks_of_curr_page) - 1 else None
if next_block:
curr_block_last_para_key = list(current_block["paras"].keys())[-1]
curr_blk_last_para = current_block["paras"][curr_block_last_para_key]
next_block_first_para_key = list(next_block["paras"].keys())[0]
next_blk_first_para = next_block["paras"][next_block_first_para_key]
if self.should_merge_next_para(curr_blk_last_para, next_blk_first_para):
curr_blk_last_para["next_para_location"] = [
curr_page_idx,
next_block["block_id"],
int(next_block_first_para_key.split("_")[-1]),
]
curr_blk_last_para["merge_next_para"] = True
else:
# Handle the case where the next block is in a different page
curr_block_last_para_key = list(current_block["paras"].keys())[-1]
curr_blk_last_para = current_block["paras"][curr_block_last_para_key]
while not next_page_content.get("para_blocks", []) and next_page_idx <= the_last_page_id:
next_page_idx += 1
next_page_id = f"page_{next_page_idx}"
next_page_content = pdf_dict.get(next_page_id, {})
if next_page_content.get("para_blocks", []):
next_blk_first_para_key = list(next_page_content["para_blocks"][0]["paras"].keys())[0]
next_blk_first_para = next_page_content["para_blocks"][0]["paras"][next_blk_first_para_key]
if self.should_merge_next_para(curr_blk_last_para, next_blk_first_para):
curr_blk_last_para["next_para_location"] = [
next_page_idx,
next_page_content["para_blocks"][0]["block_id"],
int(next_blk_first_para_key.split("_")[-1]),
]
curr_blk_last_para["merge_next_para"] = True
return pdf_dict
def find_block_by_id(self, para_blocks, block_id):
"""
This function finds a block by its id.
Parameters
----------
para_blocks : list
List of blocks.
block_id : int
Id of the block to find.
Returns
-------
block : dict
The block with the given id.
"""
for blk_idx, block in enumerate(para_blocks):
if block.get("block_id") == block_id:
return block
return None
def batch_merge_paras(self, pdf_dict):
"""
This function merges the paragraphs in the pdf_dict.
Parameters
----------
pdf_dict : dict
PDF dictionary.
Returns
-------
pdf_dict : dict
PDF dictionary with merged paragraphs.
"""
for page_id, page_content in pdf_dict.items():
if page_id.startswith("page_") and page_content.get("para_blocks", []):
para_blocks_of_page = page_content["para_blocks"]
for i in range(len(para_blocks_of_page)):
current_block = para_blocks_of_page[i]
paras = current_block["paras"]
for para_id, curr_para in list(paras.items()):
# print(f"current para_id: {para_id}")
# 跳过标题段落
if curr_para.get("is_para_title"):
continue
while curr_para.get("merge_next_para"):
curr_para_location = curr_para.get("curr_para_location")
next_para_location = curr_para.get("next_para_location")
# print(f"curr_para_location: {curr_para_location}, next_para_location: {next_para_location}")
if not next_para_location:
break
if curr_para_location == next_para_location:
# print_red("The next para is in the same block as the current para.")
curr_para["merge_next_para"] = False
break
next_page_idx, next_block_id, next_para_id = next_para_location
next_page_id = f"page_{next_page_idx}"
next_page_content = pdf_dict.get(next_page_id)
if not next_page_content:
break
next_block = self.find_block_by_id(next_page_content.get("para_blocks", []), next_block_id)
if not next_block:
break
next_para = next_block["paras"].get(f"para_{next_para_id}")
if not next_para or next_para.get("is_para_title"):
break
# 合并段落文本
curr_para_text = curr_para.get("para_text", "")
next_para_text = next_para.get("para_text", "")
curr_para["para_text"] = curr_para_text + " " + next_para_text
# 更新 next_para_location
curr_para["next_para_location"] = next_para.get("next_para_location")
# 将下一个段落文本置为空，表示已被合并
next_para["para_text"] = ""
# 更新 merge_next_para 标记
curr_para["merge_next_para"] = next_para.get("merge_next_para", False)
return pdf_dict
class DrawAnnos:
"""
This class draws annotations on the pdf file
----------------------------------------
Color Code
----------------------------------------
Red: (1, 0, 0)
Green: (0, 1, 0)
Blue: (0, 0, 1)
Yellow: (1, 1, 0) - mix of red and green
Cyan: (0, 1, 1) - mix of green and blue
Magenta: (1, 0, 1) - mix of red and blue
White: (1, 1, 1) - red, green and blue full intensity
Black: (0, 0, 0) - no color component whatsoever
Gray: (0.5, 0.5, 0.5) - equal and medium intensity of red, green and blue color components
Orange: (1, 0.65, 0) - maximum intensity of red, medium intensity of green, no blue component
"""
def __init__(self) -> None:
pass
def __is_nested_list(self, lst):
"""
This function returns True if the given list is a nested list of any degree.
"""
if isinstance(lst, list):
return any(self.__is_nested_list(i) for i in lst) or any(isinstance(i, list) for i in lst)
return False
def __valid_rect(self, bbox):
# Ensure that the rectangle is not empty or invalid
if isinstance(bbox[0], list):
return False # It's a nested list, hence it can't be valid rect
else:
return bbox[0] < bbox[2] and bbox[1] < bbox[3]
def __draw_nested_boxes(self, page, nested_bbox, color=(0, 1, 1)):
"""
This function draws the nested boxes
Parameters
----------
page : fitz.Page
page
nested_bbox : list
nested bbox
color : tuple
color, by default (0, 1, 1) # draw with cyan color for combined paragraph
"""
if self.__is_nested_list(nested_bbox): # If it's a nested list
for bbox in nested_bbox:
self.__draw_nested_boxes(page, bbox, color) # Recursively call the function
elif self.__valid_rect(nested_bbox): # If valid rectangle
para_rect = fitz.Rect(nested_bbox)
para_anno = page.add_rect_annot(para_rect)
para_anno.set_colors(stroke=color) # draw with cyan color for combined paragraph
para_anno.set_border(width=1)
para_anno.update()
def draw_annos(self, input_pdf_path, pdf_dic, output_pdf_path):
"""
This function draws annotations on the pdf file.
Parameters
----------
input_pdf_path : str
path to the input pdf file
pdf_dic : dict
pdf dictionary
output_pdf_path : str
path to the output pdf file
pdf_dic : dict
pdf dictionary
"""
pdf_doc = open_pdf(input_pdf_path)
if pdf_dic is None:
pdf_dic = {}
if output_pdf_path is None:
output_pdf_path = input_pdf_path.replace(".pdf", "_anno.pdf")
for page_id, page in enumerate(pdf_doc): # type: ignore
page_key = f"page_{page_id}"
for ele_key, ele_data in pdf_dic[page_key].items():
if ele_key == "para_blocks":
para_blocks = ele_data
for para_block in para_blocks:
if "paras" in para_block.keys():
paras = para_block["paras"]
for para_key, para_content in paras.items():
para_bbox = para_content["para_bbox"]
# print(f"para_bbox: {para_bbox}")
# print(f"is a nested list: {self.__is_nested_list(para_bbox)}")
if self.__is_nested_list(para_bbox) and len(para_bbox) > 1:
color = (0, 1, 1)
self.__draw_nested_boxes(
page, para_bbox, color
) # draw with cyan color for combined paragraph
else:
if self.__valid_rect(para_bbox):
para_rect = fitz.Rect(para_bbox)
para_anno = page.add_rect_annot(para_rect)
para_anno.set_colors(stroke=(0, 1, 0)) # draw with green color for normal paragraph
para_anno.set_border(width=0.5)
para_anno.update()
is_para_title = para_content["is_para_title"]
if is_para_title:
if self.__is_nested_list(para_content["para_bbox"]) and len(para_content["para_bbox"]) > 1:
color = (0, 0, 1)
self.__draw_nested_boxes(
page, para_content["para_bbox"], color
) # draw with cyan color for combined title
else:
if self.__valid_rect(para_content["para_bbox"]):
para_rect = fitz.Rect(para_content["para_bbox"])
if self.__valid_rect(para_content["para_bbox"]):
para_anno = page.add_rect_annot(para_rect)
para_anno.set_colors(stroke=(0, 0, 1)) # draw with blue color for normal title
para_anno.set_border(width=0.5)
para_anno.update()
pdf_doc.save(output_pdf_path)
pdf_doc.close()
class ParaProcessPipeline:
def __init__(self) -> None:
pass
def para_process_pipeline(self, pdf_info_dict, para_debug_mode=None, input_pdf_path=None, output_pdf_path=None):
"""
This function processes the paragraphs, including:
1. Read raw input json file into pdf_dic
2. Detect and replace equations
3. Combine spans into a natural line
4. Check if the paragraphs are inside bboxes passed from "layout_bboxes" key
5. Compute statistics for each block
6. Detect titles in the document
7. Detect paragraphs inside each block
8. Divide the level of the titles
9. Detect and combine paragraphs from different blocks into one paragraph
10. Check whether the final results after checking headings, dividing paragraphs within blocks, and merging paragraphs between blocks are plausible and reasonable.
11. Draw annotations on the pdf file
Parameters
----------
pdf_dic_json_fpath : str
path to the pdf dictionary json file.
Notice: data noises, including overlap blocks, header, footer, watermark, vertical margin note have been removed already.
input_pdf_doc : str
path to the input pdf file
output_pdf_path : str
path to the output pdf file
Returns
-------
pdf_dict : dict
result dictionary
"""
error_info = None
output_json_file = ""
output_dir = ""
if input_pdf_path is not None:
input_pdf_path = os.path.abspath(input_pdf_path)
# print_green_on_red(f">>>>>>>>>>>>>>>>>>> Process the paragraphs of {input_pdf_path}")
if output_pdf_path is not None:
output_dir = os.path.dirname(output_pdf_path)
output_json_file = f"{output_dir}/pdf_dic.json"
def __save_pdf_dic(pdf_dic, output_pdf_path, stage="0", para_debug_mode=para_debug_mode):
"""
Save the pdf_dic to a json file
"""
output_pdf_file_name = os.path.basename(output_pdf_path)
# output_dir = os.path.dirname(output_pdf_path)
output_dir = "\\tmp\\pdf_parse"
output_pdf_file_name = output_pdf_file_name.replace(".pdf", f"_stage_{stage}.json")
pdf_dic_json_fpath = os.path.join(output_dir, output_pdf_file_name)
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if para_debug_mode == "full":
with open(pdf_dic_json_fpath, "w", encoding="utf-8") as f:
json.dump(pdf_dic, f, indent=2, ensure_ascii=False)
# Validate the output already exists
if not os.path.exists(pdf_dic_json_fpath):
print_red(f"Failed to save the pdf_dic to {pdf_dic_json_fpath}")
return None
else:
print_green(f"Succeed to save the pdf_dic to {pdf_dic_json_fpath}")
return pdf_dic_json_fpath
"""
Preprocess the lines of block
"""
# Combine spans into a natural line
rawBlockProcessor = RawBlockProcessor()
pdf_dic = rawBlockProcessor.batch_process_blocks(pdf_info_dict)
# print(f"pdf_dic['page_0']['para_blocks'][0]: {pdf_dic['page_0']['para_blocks'][0]}", end="\n\n")
# Check if the paragraphs are inside bboxes passed from "layout_bboxes" key
layoutFilter = LayoutFilterProcessor()
pdf_dic = layoutFilter.batch_process_blocks(pdf_dic)
# Compute statistics for each block
blockStatisticsCalculator = BlockStatisticsCalculator()
pdf_dic = blockStatisticsCalculator.batch_process_blocks(pdf_dic)
# print(f"pdf_dic['page_0']['para_blocks'][0]: {pdf_dic['page_0']['para_blocks'][0]}", end="\n\n")
# Compute statistics for all blocks(namely this pdf document)
docStatisticsCalculator = DocStatisticsCalculator()
pdf_dic = docStatisticsCalculator.calc_stats_of_doc(pdf_dic)
# print(f"pdf_dic['statistics']: {pdf_dic['statistics']}", end="\n\n")
# Dump the first three stages of pdf_dic to a json file
if para_debug_mode == "full":
pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="0", para_debug_mode=para_debug_mode)
"""
Detect titles in the document
"""
doc_statistics = pdf_dic["statistics"]
titleProcessor = TitleProcessor(doc_statistics)
pdf_dic = titleProcessor.batch_detect_titles(pdf_dic)
if para_debug_mode == "full":
pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="1", para_debug_mode=para_debug_mode)
"""
Detect and divide the level of the titles
"""
titleProcessor = TitleProcessor()
pdf_dic = titleProcessor.batch_recog_title_level(pdf_dic)
if para_debug_mode == "full":
pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="2", para_debug_mode=para_debug_mode)
"""
Detect and split paragraphs inside each block
"""
blockInnerParasProcessor = BlockTerminationProcessor()
pdf_dic = blockInnerParasProcessor.batch_process_blocks(pdf_dic)
if para_debug_mode == "full":
pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="3", para_debug_mode=para_debug_mode)
# pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="3", para_debug_mode="full")
# print_green(f"pdf_dic_json_fpath: {pdf_dic_json_fpath}")
"""
Detect and combine paragraphs from different blocks into one paragraph
"""
blockContinuationProcessor = BlockContinuationProcessor()
pdf_dic = blockContinuationProcessor.batch_tag_paras(pdf_dic)
pdf_dic = blockContinuationProcessor.batch_merge_paras(pdf_dic)
if para_debug_mode == "full":
pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="4", para_debug_mode=para_debug_mode)
# pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="4", para_debug_mode="full")
# print_green(f"pdf_dic_json_fpath: {pdf_dic_json_fpath}")
"""
Discard pdf files by checking exceptions and return the error info to the caller
"""
discardByException = DiscardByException()
is_discard_by_single_line_block = discardByException.discard_by_single_line_block(
pdf_dic, exception=DenseSingleLineBlockException()
)
is_discard_by_title_detection = discardByException.discard_by_title_detection(
pdf_dic, exception=TitleDetectionException()
)
is_discard_by_title_level = discardByException.discard_by_title_level(pdf_dic, exception=TitleLevelException())
is_discard_by_split_para = discardByException.discard_by_split_para(pdf_dic, exception=ParaSplitException())
is_discard_by_merge_para = discardByException.discard_by_merge_para(pdf_dic, exception=ParaMergeException())
if is_discard_by_single_line_block is not None:
error_info = is_discard_by_single_line_block
elif is_discard_by_title_detection is not None:
error_info = is_discard_by_title_detection
elif is_discard_by_title_level is not None:
error_info = is_discard_by_title_level
elif is_discard_by_split_para is not None:
error_info = is_discard_by_split_para
elif is_discard_by_merge_para is not None:
error_info = is_discard_by_merge_para
if error_info is not None:
return pdf_dic, error_info
"""
Dump the final pdf_dic to a json file
"""
if para_debug_mode is not None:
with open(output_json_file, "w", encoding="utf-8") as f:
json.dump(pdf_info_dict, f, ensure_ascii=False, indent=4)
"""
Draw the annotations
"""
if para_debug_mode is not None:
drawAnnos = DrawAnnos()
drawAnnos.draw_annos(input_pdf_path, pdf_dic, output_pdf_path)
"""
Remove the intermediate files which are generated in the process of paragraph processing if debug_mode is simple
"""
if para_debug_mode is not None:
for fpath in os.listdir(output_dir):
if fpath.endswith(".json") and "stage" in fpath:
os.remove(os.path.join(output_dir, fpath))
return pdf_dic, error_info
"""
Run this script to test the function with Command:
python detect_para.py [pdf_path] [output_pdf_path]
Params:
- pdf_path: the path of the pdf file
- output_pdf_path: the path of the output pdf file
"""
if __name__ == "__main__":
DEFAULT_PDF_PATH = (
"app/pdf_toolbox/tests/assets/paper/paper.pdf" if os.name != "nt" else "app\\pdf_toolbox\\tests\\assets\\paper\\paper.pdf"
)
input_pdf_path = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_PDF_PATH
output_pdf_path = sys.argv[2] if len(sys.argv) > 2 else input_pdf_path.split(".")[0] + "_recogPara.pdf"
output_json_path = sys.argv[3] if len(sys.argv) > 3 else input_pdf_path.split(".")[0] + "_recogPara.json"
import stat
# Remove existing output file if it exists
if os.path.exists(output_pdf_path):
os.chmod(output_pdf_path, stat.S_IWRITE)
os.remove(output_pdf_path)
input_pdf_doc = open_pdf(input_pdf_path)
# postprocess the paragraphs
paraProcessPipeline = ParaProcessPipeline()
# parse paragraph and save to json file
pdf_dic = {}
blockInnerParasProcessor = BlockTerminationProcessor()
"""
Construct the pdf dictionary.
"""
for page_id, page in enumerate(input_pdf_doc): # type: ignore
# print(f"Processing page {page_id}")
# print(f"page: {page}")
raw_blocks = page.get_text("dict")["blocks"]
# Save text blocks to "preproc_blocks"
preproc_blocks = []
for block in raw_blocks:
if block["type"] == 0:
preproc_blocks.append(block)
layout_bboxes = []
# Construct the pdf dictionary as schema above
page_dict = {
"para_blocks": None,
"preproc_blocks": preproc_blocks,
"images": None,
"tables": None,
"interline_equations": None,
"inline_equations": None,
"layout_bboxes": None,
"pymu_raw_blocks": None,
"global_statistic": None,
"droped_text_block": None,
"droped_image_block": None,
"droped_table_block": None,
"image_backup": None,
"table_backup": None,
}
pdf_dic[f"page_{page_id}"] = page_dict
# print(f"pdf_dic: {pdf_dic}")
with open(output_json_path, "w", encoding="utf-8") as f:
json.dump(pdf_dic, f, ensure_ascii=False, indent=4)
pdf_dic = paraProcessPipeline.para_process_pipeline(output_json_path, input_pdf_doc, output_pdf_path)
magic_pdf/post_proc.bak/pdf_post_filter.py.bak deleted 100644 → 0
View file @ 79b58a1e
from loguru import logger
from magic_pdf.config.drop_reason import DropReason
from magic_pdf.layout.layout_sort import get_columns_cnt_of_layout
def __is_pseudo_single_column(page_info) -> bool:
"""判断一个页面是否伪单列。
Args:
page_info (dict): 页面信息字典，包括'_layout_tree'和'preproc_blocks'。
Returns:
Tuple[bool, Optional[str]]: 如果页面伪单列返回(True, extra_info)，否则返回(False, None)。
"""
layout_tree = page_info['_layout_tree']
layout_column_width = get_columns_cnt_of_layout(layout_tree)
if layout_column_width == 1:
text_blocks = page_info['preproc_blocks']
# 遍历每一个text_block
for text_block in text_blocks:
lines = text_block['lines']
num_lines = len(lines)
num_satisfying_lines = 0
for i in range(num_lines - 1):
current_line = lines[i]
next_line = lines[i + 1]
# 获取当前line和下一个line的bbox属性
current_bbox = current_line['bbox']
next_bbox = next_line['bbox']
# 检查是否满足条件
if next_bbox[0] > current_bbox[2] or next_bbox[2] < current_bbox[0]:
num_satisfying_lines += 1
# 如果有一半以上的line满足条件，就drop
# print("num_satisfying_lines:", num_satisfying_lines, "num_lines:", num_lines)
if num_lines > 20:
radio = num_satisfying_lines / num_lines
if radio >= 0.5:
extra_info = f'{{num_lines: {num_lines}, num_satisfying_lines: {num_satisfying_lines}}}'
block_text = []
for line in lines:
if line['spans']:
for span in line['spans']:
block_text.append(span['text'])
logger.warning(f'pseudo_single_column block_text: {block_text}')
return True, extra_info
return False, None
def pdf_post_filter(page_info) -> tuple:
"""return:(True|False, err_msg) True, 如果pdf符合要求 False, 如果pdf不符合要求."""
bool_is_pseudo_single_column, extra_info = __is_pseudo_single_column(page_info)
if bool_is_pseudo_single_column:
return False, {'_need_drop': True, '_drop_reason': DropReason.PSEUDO_SINGLE_COLUMN, 'extra_info': extra_info}
return True, None
magic_pdf/post_proc.bak/remove_footnote.py.bak deleted 100644 → 0
View file @ 79b58a1e
from magic_pdf.libs.boxbase import _is_in, _is_in_or_part_overlap
import collections # 统计库
def is_below(bbox1, bbox2):
# 如果block1的上边y坐标大于block2的下边y坐标，那么block1在block2下面
return bbox1[1] > bbox2[3]
def merge_bboxes(bboxes):
# 找出所有blocks的最小x0，最大y1，最大x1，最小y0，这就是合并后的bbox
x0 = min(bbox[0] for bbox in bboxes)
y0 = min(bbox[1] for bbox in bboxes)
x1 = max(bbox[2] for bbox in bboxes)
y1 = max(bbox[3] for bbox in bboxes)
return [x0, y0, x1, y1]
def merge_footnote_blocks(page_info, main_text_font):
page_info['merged_bboxes'] = []
for layout in page_info['layout_bboxes']:
# 找出layout中的所有footnote blocks和preproc_blocks
footnote_bboxes = [block for block in page_info['footnote_bboxes_tmp'] if _is_in(block, layout['layout_bbox'])]
# 如果没有footnote_blocks，就跳过这个layout
if not footnote_bboxes:
continue
preproc_blocks = [block for block in page_info['preproc_blocks'] if _is_in(block['bbox'], layout['layout_bbox'])]
# preproc_bboxes = [block['bbox'] for block in preproc_blocks]
font_names = collections.Counter()
if len(preproc_blocks) > 0:
# 存储每一行的文本块大小的列表
line_sizes = []
# 存储每个文本块的平均行大小
block_sizes = []
for block in preproc_blocks:
block_line_sizes = []
block_fonts = collections.Counter()
for line in block['lines']:
# 提取每个span的size属性，并计算行大小
span_sizes = [span['size'] for span in line['spans'] if 'size' in span]
if span_sizes:
line_size = sum(span_sizes) / len(span_sizes)
line_sizes.append(line_size)
block_line_sizes.append(line_size)
span_font = [(span['font'], len(span['text'])) for span in line['spans'] if
'font' in span and len(span['text']) > 0]
if span_font:
# # todo main_text_font应该用基于字数最多的字体而不是span级别的统计
# font_names.append(font_name for font_name in span_font)
# block_fonts.append(font_name for font_name in span_font)
for font, count in span_font:
# font_names.extend([font] * count)
# block_fonts.extend([font] * count)
font_names[font] += count
block_fonts[font] += count
if block_line_sizes:
# 计算文本块的平均行大小
block_size = sum(block_line_sizes) / len(block_line_sizes)
block_font = block_fonts.most_common(1)[0][0]
block_sizes.append((block, block_size, block_font))
# 计算main_text_size
# main_text_font = font_names.most_common(1)[0][0]
main_text_size = collections.Counter(line_sizes).most_common(1)[0][0]
else:
continue
need_merge_bboxes = []
# 任何一个下面有正文block的footnote bbox都是假footnote
for footnote_bbox in footnote_bboxes:
# 检测footnote下面是否有正文block(正文block需满足，block平均size大于等于main_text_size，且block行数大于等于5)
main_text_bboxes_below = [block['bbox'] for block, size, block_font in block_sizes if
is_below(block['bbox'], footnote_bbox) and
sum([size >= main_text_size,
len(block['lines']) >= 5,
block_font == main_text_font])
>= 2]
# 如果main_text_bboxes_below不为空，说明footnote下面有正文block，这个footnote不成立，跳过
if len(main_text_bboxes_below) > 0:
continue
else:
# 否则，说明footnote下面没有正文block，这个footnote成立，添加到待merge的footnote_bboxes中
need_merge_bboxes.append(footnote_bbox)
if len(need_merge_bboxes) == 0:
continue
# 找出最靠上的footnote block
top_footnote_bbox = min(need_merge_bboxes, key=lambda bbox: bbox[1])
# 找出所有在top_footnote_block下面的preproc_blocks，并确保这些preproc_blocks的平均行大小小于main_text_size
bboxes_below = [block['bbox'] for block, size, block_font in block_sizes if is_below(block['bbox'], top_footnote_bbox)]
# # 找出所有在top_footnote_block下面的preproc_blocks
# bboxes_below = [bbox for bbox in preproc_bboxes if is_below(bbox, top_footnote_bbox)]
# 合并top_footnote_block和blocks_below
merged_bbox = merge_bboxes([top_footnote_bbox] + bboxes_below)
# 添加到新的footnote_bboxes_tmp中
page_info['merged_bboxes'].append(merged_bbox)
return page_info
def remove_footnote_blocks(page_info):
if page_info.get('merged_bboxes'):
# 从文字中去掉footnote
remain_text_blocks, removed_footnote_text_blocks = remove_footnote_text(page_info['preproc_blocks'], page_info['merged_bboxes'])
# 从图片中去掉footnote
image_blocks, removed_footnote_imgs_blocks = remove_footnote_image(page_info['images'], page_info['merged_bboxes'])
# 更新page_info
page_info['preproc_blocks'] = remain_text_blocks
page_info['images'] = image_blocks
page_info['droped_text_block'].extend(removed_footnote_text_blocks)
page_info['droped_image_block'].extend(removed_footnote_imgs_blocks)
# 删除footnote_bboxes_tmp和merged_bboxes
del page_info['merged_bboxes']
del page_info['footnote_bboxes_tmp']
return page_info
def remove_footnote_text(raw_text_block, footnote_bboxes):
"""
:param raw_text_block: str类型，是当前页的文本内容
:param footnoteBboxes: list类型，是当前页的脚注bbox
"""
footnote_text_blocks = []
for block in raw_text_block:
text_bbox = block['bbox']
# TODO 更严谨点在line级别做
if any([_is_in_or_part_overlap(text_bbox, footnote_bbox) for footnote_bbox in footnote_bboxes]):
# if any([text_bbox[3]>=footnote_bbox[1] for footnote_bbox in footnote_bboxes]):
block['tag'] = 'footnote'
footnote_text_blocks.append(block)
# raw_text_block.remove(block)
# 移除，不能再内部移除，否则会出错
for block in footnote_text_blocks:
raw_text_block.remove(block)
return raw_text_block, footnote_text_blocks
def remove_footnote_image(image_blocks, footnote_bboxes):
"""
:param image_bboxes: list类型，是当前页的图片bbox(结构体)
:param footnoteBboxes: list类型，是当前页的脚注bbox
"""
footnote_imgs_blocks = []
for image_block in image_blocks:
if any([_is_in(image_block['bbox'], footnote_bbox) for footnote_bbox in footnote_bboxes]):
footnote_imgs_blocks.append(image_block)
for footnote_imgs_block in footnote_imgs_blocks:
image_blocks.remove(footnote_imgs_block)
return image_blocks, footnote_imgs_blocks
\ No newline at end of file
magic_pdf/pre_proc/citationmarker_remove.py.bak deleted 100644 → 0
View file @ 79b58a1e
"""
去掉正文的引文引用marker
https://aicarrier.feishu.cn/wiki/YLOPwo1PGiwFRdkwmyhcZmr0n3d
"""
import re
# from magic_pdf.libs.nlp_utils import NLPModels
# __NLP_MODEL = NLPModels()
def check_1(spans, cur_span_i):
"""寻找前一个char,如果是句号，逗号，那么就是角标"""
if cur_span_i==0:
return False # 不是角标
pre_span = spans[cur_span_i-1]
pre_char = pre_span['chars'][-1]['c']
if pre_char in ['。', '，', '.', ',']:
return True
return False
# def check_2(spans, cur_span_i):
# """检查前面一个span的最后一个单词，如果长度大于5，全都是字母，并且不含大写，就是角标"""
# pattern = r'\b[A-Z]\.\s[A-Z][a-z]*\b' # 形如A. Bcde, L. Bcde, 人名的缩写
#
# if cur_span_i==0 and len(spans)>1:
# next_span = spans[cur_span_i+1]
# next_txt = "".join([c['c'] for c in next_span['chars']])
# result = __NLP_MODEL.detect_entity_catgr_using_nlp(next_txt)
# if result in ["PERSON", "GPE", "ORG"]:
# return True
#
# if re.findall(pattern, next_txt):
# return True
#
# return False # 不是角标
# elif cur_span_i==0 and len(spans)==1: # 角标占用了整行？谨慎删除
# return False
#
# # 如果这个span是最后一个span,
# if cur_span_i==len(spans)-1:
# pre_span = spans[cur_span_i-1]
# pre_txt = "".join([c['c'] for c in pre_span['chars']])
# pre_word = pre_txt.split(' ')[-1]
# result = __NLP_MODEL.detect_entity_catgr_using_nlp(pre_txt)
# if result in ["PERSON", "GPE", "ORG"]:
# return True
#
# if re.findall(pattern, pre_txt):
# return True
#
# return len(pre_word) > 5 and pre_word.isalpha() and pre_word.islower()
# else: # 既不是第一个span，也不是最后一个span，那么此时检查一下这个角标距离前后哪个单词更近就属于谁的角标
# pre_span = spans[cur_span_i-1]
# next_span = spans[cur_span_i+1]
# cur_span = spans[cur_span_i]
# # 找到前一个和后一个span里的距离最近的单词
# pre_distance = 10000 # 一个很大的数
# next_distance = 10000 # 一个很大的数
# for c in pre_span['chars'][::-1]:
# if c['c'].isalpha():
# pre_distance = cur_span['bbox'][0] - c['bbox'][2]
# break
# for c in next_span['chars']:
# if c['c'].isalpha():
# next_distance = c['bbox'][0] - cur_span['bbox'][2]
# break
#
# if pre_distance<next_distance:
# belong_to_span = pre_span
# else:
# belong_to_span = next_span
#
# txt = "".join([c['c'] for c in belong_to_span['chars']])
# pre_word = txt.split(' ')[-1]
# result = __NLP_MODEL.detect_entity_catgr_using_nlp(txt)
# if result in ["PERSON", "GPE", "ORG"]:
# return True
#
# if re.findall(pattern, txt):
# return True
#
# return len(pre_word) > 5 and pre_word.isalpha() and pre_word.islower()
def check_3(spans, cur_span_i):
"""上标里有[], 有*， 有-， 有逗号"""
# 如[2-3],[22]
# 如 2,3,4
cur_span_txt = ''.join(c['c'] for c in spans[cur_span_i]['chars']).strip()
bad_char = ['[', ']', '*', ',']
if any([c in cur_span_txt for c in bad_char]) and any(character.isdigit() for character in cur_span_txt):
return True
# 如2-3, a-b
patterns = [r'\d+-\d+', r'[a-zA-Z]-[a-zA-Z]', r'[a-zA-Z],[a-zA-Z]']
for pattern in patterns:
match = re.match(pattern, cur_span_txt)
if match is not None:
return True
return False
def remove_citation_marker(with_char_text_blcoks):
for blk in with_char_text_blcoks:
for line in blk['lines']:
# 如果span里的个数少于2个，那只能忽略，角标不可能自己独占一行
if len(line['spans'])<=1:
continue
# 找到高度最高的span作为位置比较的基准
max_hi_span = line['spans'][0]['bbox']
min_font_sz = 10000 # line里最小的字体
max_font_sz = 0 # line里最大的字体
for s in line['spans']:
if max_hi_span[3]-max_hi_span[1]<s['bbox'][3]-s['bbox'][1]:
max_hi_span = s['bbox']
if min_font_sz>s['size']:
min_font_sz = s['size']
if max_font_sz<s['size']:
max_font_sz = s['size']
base_span_mid_y = (max_hi_span[3]+max_hi_span[1])/2
span_to_del = []
for i, span in enumerate(line['spans']):
span_hi = span['bbox'][3]-span['bbox'][1]
span_mid_y = (span['bbox'][3]+span['bbox'][1])/2
span_font_sz = span['size']
if max_font_sz-span_font_sz<1: # 先以字体过滤正文，如果是正文就不再继续判断了
continue
# 对被除数为0的情况进行过滤
if span_hi==0 or min_font_sz==0:
continue
if (base_span_mid_y-span_mid_y)/span_hi>0.2 or (base_span_mid_y-span_mid_y>0 and abs(span_font_sz-min_font_sz)/min_font_sz<0.1):
"""
1. 它的前一个char如果是句号或者逗号的话，那么肯定是角标而不是公式
2. 如果这个角标的前面是一个单词（长度大于5）而不是任何大写或小写的短字母的话 应该也是角标
3. 上标里有数字和逗号或者数字+星号的组合，方括号，一般肯定就是角标了
4. 这个角标属于前文还是后文要根据距离来判断，如果距离前面的文本太近，那么就是前面的角标，否则就是后面的角标
"""
if (check_1(line['spans'], i) or
# check_2(line['spans'], i) or
check_3(line['spans'], i)
):
"""删除掉这个角标：删除这个span, 同时还要更新line的text"""
span_to_del.append(span)
if len(span_to_del)>0:
for span in span_to_del:
line['spans'].remove(span)
line['text'] = ''.join([c['c'] for s in line['spans'] for c in s['chars']])
return with_char_text_blcoks
magic_pdf/pre_proc/detect_equation.py.bak deleted 100644 → 0
View file @ 79b58a1e
from magic_pdf.libs.boxbase import _is_in, calculate_overlap_area_2_minbox_area_ratio # 正则
from magic_pdf.libs.commons import fitz # pyMuPDF库
def __solve_contain_bboxs(all_bbox_list: list):
"""将两个公式的bbox做判断是否有包含关系，若有的话则删掉较小的bbox"""
dump_list = []
for i in range(len(all_bbox_list)):
for j in range(i + 1, len(all_bbox_list)):
# 获取当前两个值
bbox1 = all_bbox_list[i][:4]
bbox2 = all_bbox_list[j][:4]
# 删掉较小的框
if _is_in(bbox1, bbox2):
dump_list.append(all_bbox_list[i])
elif _is_in(bbox2, bbox1):
dump_list.append(all_bbox_list[j])
else:
ratio = calculate_overlap_area_2_minbox_area_ratio(bbox1, bbox2)
if ratio > 0.7:
s1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
s2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])
if s2 > s1:
dump_list.append(all_bbox_list[i])
else:
dump_list.append(all_bbox_list[i])
# 遍历需要删除的列表中的每个元素
for item in dump_list:
while item in all_bbox_list:
all_bbox_list.remove(item)
return all_bbox_list
def parse_equations(page_ID: int, page: fitz.Page, json_from_DocXchain_obj: dict):
"""
:param page_ID: int类型，当前page在当前pdf文档中是第page_D页。
:param page :fitz读取的当前页的内容
:param res_dir_path: str类型，是每一个pdf文档，在当前.py文件的目录下生成一个与pdf文档同名的文件夹，res_dir_path就是文件夹的dir
:param json_from_DocXchain_obj: dict类型，把pdf文档送入DocXChain模型中后，提取bbox，结果保存到pdf文档同名文件夹下的 page_ID.json文件中了。json_from_DocXchain_obj就是打开后的dict
"""
DPI = 72 # use this resolution
pix = page.get_pixmap(dpi=DPI)
pageL = 0
pageR = int(pix.w)
pageU = 0
pageD = int(pix.h)
#--------- 通过json_from_DocXchain来获取 table ---------#
equationEmbedding_from_DocXChain_bboxs = []
equationIsolated_from_DocXChain_bboxs = []
xf_json = json_from_DocXchain_obj
width_from_json = xf_json['page_info']['width']
height_from_json = xf_json['page_info']['height']
LR_scaleRatio = width_from_json / (pageR - pageL)
UD_scaleRatio = height_from_json / (pageD - pageU)
for xf in xf_json['layout_dets']:
# {0: 'title', 1: 'figure', 2: 'plain text', 3: 'header', 4: 'page number', 5: 'footnote', 6: 'footer', 7: 'table', 8: 'table caption', 9: 'figure caption', 10: 'equation', 11: 'full column', 12: 'sub column'}
L = xf['poly'][0] / LR_scaleRatio
U = xf['poly'][1] / UD_scaleRatio
R = xf['poly'][2] / LR_scaleRatio
D = xf['poly'][5] / UD_scaleRatio
# L += pageL # 有的页面，artBox偏移了。不在（0,0）
# R += pageL
# U += pageU
# D += pageU
L, R = min(L, R), max(L, R)
U, D = min(U, D), max(U, D)
# equation
img_suffix = f"{page_ID}_{int(L)}_{int(U)}_{int(R)}_{int(D)}"
if xf['category_id'] == 13 and xf['score'] >= 0.3:
latex_text = xf.get("latex", "EmptyInlineEquationResult")
debugable_latex_text = f"{latex_text}|{img_suffix}"
equationEmbedding_from_DocXChain_bboxs.append((L, U, R, D, latex_text))
if xf['category_id'] == 14 and xf['score'] >= 0.3:
latex_text = xf.get("latex", "EmptyInterlineEquationResult")
debugable_latex_text = f"{latex_text}|{img_suffix}"
equationIsolated_from_DocXChain_bboxs.append((L, U, R, D, latex_text))
#---------------------------------------- 排序，编号，保存 -----------------------------------------#
equationIsolated_from_DocXChain_bboxs.sort(key = lambda LURD: (LURD[1], LURD[0]))
equationIsolated_from_DocXChain_bboxs.sort(key = lambda LURD: (LURD[1], LURD[0]))
equationEmbedding_from_DocXChain_names = []
equationEmbedding_ID = 0
equationIsolated_from_DocXChain_names = []
equationIsolated_ID = 0
for L, U, R, D, _ in equationEmbedding_from_DocXChain_bboxs:
if not(L < R and U < D):
continue
try:
# cur_equation = page.get_pixmap(clip=(L,U,R,D))
new_equation_name = "equationEmbedding_{}_{}.png".format(page_ID, equationEmbedding_ID) # 公式name
# cur_equation.save(res_dir_path + '/' + new_equation_name) # 把公式存出在新建的文件夹，并命名
equationEmbedding_from_DocXChain_names.append(new_equation_name) # 把公式的名字存在list中，方便在md中插入引用
equationEmbedding_ID += 1
except:
pass
for L, U, R, D, _ in equationIsolated_from_DocXChain_bboxs:
if not(L < R and U < D):
continue
try:
# cur_equation = page.get_pixmap(clip=(L,U,R,D))
new_equation_name = "equationEmbedding_{}_{}.png".format(page_ID, equationIsolated_ID) # 公式name
# cur_equation.save(res_dir_path + '/' + new_equation_name) # 把公式存出在新建的文件夹，并命名
equationIsolated_from_DocXChain_names.append(new_equation_name) # 把公式的名字存在list中，方便在md中插入引用
equationIsolated_ID += 1
except:
pass
equationEmbedding_from_DocXChain_bboxs.sort(key = lambda LURD: (LURD[1], LURD[0]))
equationIsolated_from_DocXChain_bboxs.sort(key = lambda LURD: (LURD[1], LURD[0]))
"""根据pdf可视区域，调整bbox的坐标"""
cropbox = page.cropbox
if cropbox[0]!=page.rect[0] or cropbox[1]!=page.rect[1]:
for eq_box in equationEmbedding_from_DocXChain_bboxs:
eq_box = [eq_box[0]+cropbox[0], eq_box[1]+cropbox[1], eq_box[2]+cropbox[0], eq_box[3]+cropbox[1], eq_box[4]]
for eq_box in equationIsolated_from_DocXChain_bboxs:
eq_box = [eq_box[0]+cropbox[0], eq_box[1]+cropbox[1], eq_box[2]+cropbox[0], eq_box[3]+cropbox[1], eq_box[4]]
deduped_embedding_eq_bboxes = __solve_contain_bboxs(equationEmbedding_from_DocXChain_bboxs)
return deduped_embedding_eq_bboxes, equationIsolated_from_DocXChain_bboxs
magic_pdf/pre_proc/detect_footer_by_model.py.bak deleted 100644 → 0
View file @ 79b58a1e
from magic_pdf.libs.commons import fitz # pyMuPDF库
from magic_pdf.libs.coordinate_transform import get_scale_ratio
def parse_footers(page_ID: int, page: fitz.Page, json_from_DocXchain_obj: dict):
"""
:param page_ID: int类型，当前page在当前pdf文档中是第page_D页。
:param page :fitz读取的当前页的内容
:param res_dir_path: str类型，是每一个pdf文档，在当前.py文件的目录下生成一个与pdf文档同名的文件夹，res_dir_path就是文件夹的dir
:param json_from_DocXchain_obj: dict类型，把pdf文档送入DocXChain模型中后，提取bbox，结果保存到pdf文档同名文件夹下的 page_ID.json文件中了。json_from_DocXchain_obj就是打开后的dict
"""
#--------- 通过json_from_DocXchain来获取 footer ---------#
footer_bbox_from_DocXChain = []
xf_json = json_from_DocXchain_obj
horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(xf_json, page)
# {0: 'title', # 标题
# 1: 'figure', # 图片
# 2: 'plain text', # 文本
# 3: 'header', # 页眉
# 4: 'page number', # 页码
# 5: 'footnote', # 脚注
# 6: 'footer', # 页脚
# 7: 'table', # 表格
# 8: 'table caption', # 表格描述
# 9: 'figure caption', # 图片描述
# 10: 'equation', # 公式
# 11: 'full column', # 单栏
# 12: 'sub column', # 多栏
# 13: 'embedding', # 嵌入公式
# 14: 'isolated'} # 单行公式
for xf in xf_json['layout_dets']:
L = xf['poly'][0] / horizontal_scale_ratio
U = xf['poly'][1] / vertical_scale_ratio
R = xf['poly'][2] / horizontal_scale_ratio
D = xf['poly'][5] / vertical_scale_ratio
# L += pageL # 有的页面，artBox偏移了。不在（0,0）
# R += pageL
# U += pageU
# D += pageU
L, R = min(L, R), max(L, R)
U, D = min(U, D), max(U, D)
if xf['category_id'] == 6 and xf['score'] >= 0.3:
footer_bbox_from_DocXChain.append((L, U, R, D))
footer_final_names = []
footer_final_bboxs = []
footer_ID = 0
for L, U, R, D in footer_bbox_from_DocXChain:
# cur_footer = page.get_pixmap(clip=(L,U,R,D))
new_footer_name = "footer_{}_{}.png".format(page_ID, footer_ID) # 脚注name
# cur_footer.save(res_dir_path + '/' + new_footer_name) # 把页脚存储在新建的文件夹，并命名
footer_final_names.append(new_footer_name) # 把脚注的名字存在list中
footer_final_bboxs.append((L, U, R, D))
footer_ID += 1
footer_final_bboxs.sort(key = lambda LURD: (LURD[1], LURD[0]))
curPage_all_footer_bboxs = footer_final_bboxs
return curPage_all_footer_bboxs
magic_pdf/pre_proc/detect_footer_header_by_statistics.py.bak deleted 100644 → 0
View file @ 79b58a1e
from collections import defaultdict
from magic_pdf.libs.boxbase import calculate_iou
def compare_bbox_with_list(bbox, bbox_list, tolerance=1):
return any(all(abs(a - b) < tolerance for a, b in zip(bbox, common_bbox)) for common_bbox in bbox_list)
def is_single_line_block(block):
# Determine based on the width and height of the block
block_width = block["X1"] - block["X0"]
block_height = block["bbox"][3] - block["bbox"][1]
# If the height of the block is close to the average character height and the width is large, it is considered a single line
return block_height <= block["avg_char_height"] * 3 and block_width > block["avg_char_width"] * 3
def get_most_common_bboxes(bboxes, page_height, position="top", threshold=0.25, num_bboxes=3, min_frequency=2):
"""
This function gets the most common bboxes from the bboxes
Parameters
----------
bboxes : list
bboxes
page_height : float
height of the page
position : str, optional
"top" or "bottom", by default "top"
threshold : float, optional
threshold, by default 0.25
num_bboxes : int, optional
number of bboxes to return, by default 3
min_frequency : int, optional
minimum frequency of the bbox, by default 2
Returns
-------
common_bboxes : list
common bboxes
"""
# Filter bbox by position
if position == "top":
filtered_bboxes = [bbox for bbox in bboxes if bbox[1] < page_height * threshold]
else:
filtered_bboxes = [bbox for bbox in bboxes if bbox[3] > page_height * (1 - threshold)]
# Find the most common bbox
bbox_count = defaultdict(int)
for bbox in filtered_bboxes:
bbox_count[tuple(bbox)] += 1
# Get the most frequently occurring bbox, but only consider it when the frequency exceeds min_frequency
common_bboxes = [
bbox for bbox, count in sorted(bbox_count.items(), key=lambda item: item[1], reverse=True) if count >= min_frequency
][:num_bboxes]
return common_bboxes
def detect_footer_header2(result_dict, similarity_threshold=0.5):
"""
This function detects the header and footer of the document.
Parameters
----------
result_dict : dict
result dictionary
Returns
-------
result_dict : dict
result dictionary
"""
# Traverse all blocks in the document
single_line_blocks = 0
total_blocks = 0
single_line_blocks = 0
for page_id, blocks in result_dict.items():
if page_id.startswith("page_"):
for block_key, block in blocks.items():
if block_key.startswith("block_"):
total_blocks += 1
if is_single_line_block(block):
single_line_blocks += 1
# If there are no blocks, skip the header and footer detection
if total_blocks == 0:
print("No blocks found. Skipping header/footer detection.")
return result_dict
# If most of the blocks are single-line, skip the header and footer detection
if single_line_blocks / total_blocks > 0.5: # 50% of the blocks are single-line
# print("Skipping header/footer detection for text-dense document.")
return result_dict
# Collect the bounding boxes of all blocks
all_bboxes = []
all_texts = []
for page_id, blocks in result_dict.items():
if page_id.startswith("page_"):
for block_key, block in blocks.items():
if block_key.startswith("block_"):
all_bboxes.append(block["bbox"])
# Get the height of the page
page_height = max(bbox[3] for bbox in all_bboxes)
# Get the most common bbox lists for headers and footers
common_header_bboxes = get_most_common_bboxes(all_bboxes, page_height, position="top") if all_bboxes else []
common_footer_bboxes = get_most_common_bboxes(all_bboxes, page_height, position="bottom") if all_bboxes else []
# Detect and mark headers and footers
for page_id, blocks in result_dict.items():
if page_id.startswith("page_"):
for block_key, block in blocks.items():
if block_key.startswith("block_"):
bbox = block["bbox"]
text = block["text"]
is_header = compare_bbox_with_list(bbox, common_header_bboxes)
is_footer = compare_bbox_with_list(bbox, common_footer_bboxes)
block["is_header"] = int(is_header)
block["is_footer"] = int(is_footer)
return result_dict
def __get_page_size(page_sizes:list):
"""
页面大小可能不一样
"""
w = sum([w for w,h in page_sizes])/len(page_sizes)
h = sum([h for w,h in page_sizes])/len(page_sizes)
return w, h
def __calculate_iou(bbox1, bbox2):
iou = calculate_iou(bbox1, bbox2)
return iou
def __is_same_pos(box1, box2, iou_threshold):
iou = __calculate_iou(box1, box2)
return iou >= iou_threshold
def get_most_common_bbox(bboxes:list, page_size:list, page_cnt:int, page_range_threshold=0.2, iou_threshold=0.9):
"""
common bbox必须大于page_cnt的1/3
"""
min_occurance_cnt = max(3, page_cnt//4)
header_det_bbox = []
footer_det_bbox = []
hdr_same_pos_group = []
btn_same_pos_group = []
page_w, page_h = __get_page_size(page_size)
top_y, bottom_y = page_w*page_range_threshold, page_h*(1-page_range_threshold)
top_bbox = [b for b in bboxes if b[3]<top_y]
bottom_bbox = [b for b in bboxes if b[1]>bottom_y]
# 然后开始排序，寻找最经常出现的bbox, 寻找的时候如果IOU>iou_threshold就算是一个
for i in range(0, len(top_bbox)):
hdr_same_pos_group.append([top_bbox[i]])
for j in range(i+1, len(top_bbox)):
if __is_same_pos(top_bbox[i], top_bbox[j], iou_threshold):
#header_det_bbox = [min(top_bbox[i][0], top_bbox[j][0]), min(top_bbox[i][1], top_bbox[j][1]), max(top_bbox[i][2], top_bbox[j][2]), max(top_bbox[i][3],top_bbox[j][3])]
hdr_same_pos_group[i].append(top_bbox[j])
for i in range(0, len(bottom_bbox)):
btn_same_pos_group.append([bottom_bbox[i]])
for j in range(i+1, len(bottom_bbox)):
if __is_same_pos(bottom_bbox[i], bottom_bbox[j], iou_threshold):
#footer_det_bbox = [min(bottom_bbox[i][0], bottom_bbox[j][0]), min(bottom_bbox[i][1], bottom_bbox[j][1]), max(bottom_bbox[i][2], bottom_bbox[j][2]), max(bottom_bbox[i][3],bottom_bbox[j][3])]
btn_same_pos_group[i].append(bottom_bbox[j])
# 然后看下每一组的bbox，是否符合大于page_cnt一定比例
hdr_same_pos_group = [g for g in hdr_same_pos_group if len(g)>=min_occurance_cnt]
btn_same_pos_group = [g for g in btn_same_pos_group if len(g)>=min_occurance_cnt]
# 平铺2个list[list]
hdr_same_pos_group = [bbox for g in hdr_same_pos_group for bbox in g]
btn_same_pos_group = [bbox for g in btn_same_pos_group for bbox in g]
# 寻找hdr_same_pos_group中的box[3]最大值，btn_same_pos_group中的box[1]最小值
hdr_same_pos_group.sort(key=lambda b:b[3])
btn_same_pos_group.sort(key=lambda b:b[1])
hdr_y = hdr_same_pos_group[-1][3] if hdr_same_pos_group else 0
btn_y = btn_same_pos_group[0][1] if btn_same_pos_group else page_h
header_det_bbox = [0, 0, page_w, hdr_y]
footer_det_bbox = [0, btn_y, page_w, page_h]
# logger.warning(f"header: {header_det_bbox}, footer: {footer_det_bbox}")
return header_det_bbox, footer_det_bbox, page_w, page_h
def drop_footer_header(pdf_info_dict:dict):
"""
启用规则探测,在全局的视角上通过统计的方法。
"""
header = []
footer = []
all_text_bboxes = [blk['bbox'] for _, val in pdf_info_dict.items() for blk in val['preproc_blocks']]
image_bboxes = [img['bbox'] for _, val in pdf_info_dict.items() for img in val['images']] + [img['bbox'] for _, val in pdf_info_dict.items() for img in val['image_backup']]
page_size = [val['page_size'] for _, val in pdf_info_dict.items()]
page_cnt = len(pdf_info_dict.keys()) # 一共多少页
header, footer, page_w, page_h = get_most_common_bbox(all_text_bboxes+image_bboxes, page_size, page_cnt)
""""
把范围扩展到页面水平的整个方向上
"""
if header:
header = [0, 0, page_w, header[3]+1]
if footer:
footer = [0, footer[1]-1, page_w, page_h]
# 找到footer, header范围之后，针对每一页pdf，从text、图片中删除这些范围内的内容
# 移除text block
for _, page_info in pdf_info_dict.items():
header_text_blk = []
footer_text_blk = []
for blk in page_info['preproc_blocks']:
blk_bbox = blk['bbox']
if header and blk_bbox[3]<=header[3]:
blk['tag'] = "header"
header_text_blk.append(blk)
elif footer and blk_bbox[1]>=footer[1]:
blk['tag'] = "footer"
footer_text_blk.append(blk)
# 放入text_block_droped中
page_info['droped_text_block'].extend(header_text_blk)
page_info['droped_text_block'].extend(footer_text_blk)
for blk in header_text_blk:
page_info['preproc_blocks'].remove(blk)
for blk in footer_text_blk:
page_info['preproc_blocks'].remove(blk)
"""接下来把footer、header上的图片也删除掉。图片包括正常的和backup的"""
header_image = []
footer_image = []
for image_info in page_info['images']:
img_bbox = image_info['bbox']
if header and img_bbox[3]<=header[3]:
image_info['tag'] = "header"
header_image.append(image_info)
elif footer and img_bbox[1]>=footer[1]:
image_info['tag'] = "footer"
footer_image.append(image_info)
page_info['droped_image_block'].extend(header_image)
page_info['droped_image_block'].extend(footer_image)
for img in header_image:
page_info['images'].remove(img)
for img in footer_image:
page_info['images'].remove(img)
"""接下来吧backup的图片也删除掉"""
header_image = []
footer_image = []
for image_info in page_info['image_backup']:
img_bbox = image_info['bbox']
if header and img_bbox[3]<=header[3]:
image_info['tag'] = "header"
header_image.append(image_info)
elif footer and img_bbox[1]>=footer[1]:
image_info['tag'] = "footer"
footer_image.append(image_info)
page_info['droped_image_block'].extend(header_image)
page_info['droped_image_block'].extend(footer_image)
for img in header_image:
page_info['image_backup'].remove(img)
for img in footer_image:
page_info['image_backup'].remove(img)
return header, footer
magic_pdf/pre_proc/detect_footnote.py.bak deleted 100644 → 0
View file @ 79b58a1e
from collections import Counter
from magic_pdf.libs.commons import fitz # pyMuPDF库
from magic_pdf.libs.coordinate_transform import get_scale_ratio
def parse_footnotes_by_model(page_ID: int, page: fitz.Page, json_from_DocXchain_obj: dict, md_bookname_save_path=None, debug_mode=False):
"""
:param page_ID: int类型，当前page在当前pdf文档中是第page_D页。
:param page :fitz读取的当前页的内容
:param res_dir_path: str类型，是每一个pdf文档，在当前.py文件的目录下生成一个与pdf文档同名的文件夹，res_dir_path就是文件夹的dir
:param json_from_DocXchain_obj: dict类型，把pdf文档送入DocXChain模型中后，提取bbox，结果保存到pdf文档同名文件夹下的 page_ID.json文件中了。json_from_DocXchain_obj就是打开后的dict
"""
#--------- 通过json_from_DocXchain来获取 footnote ---------#
footnote_bbox_from_DocXChain = []
xf_json = json_from_DocXchain_obj
horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(xf_json, page)
# {0: 'title', # 标题
# 1: 'figure', # 图片
# 2: 'plain text', # 文本
# 3: 'header', # 页眉
# 4: 'page number', # 页码
# 5: 'footnote', # 脚注
# 6: 'footer', # 页脚
# 7: 'table', # 表格
# 8: 'table caption', # 表格描述
# 9: 'figure caption', # 图片描述
# 10: 'equation', # 公式
# 11: 'full column', # 单栏
# 12: 'sub column', # 多栏
# 13: 'embedding', # 嵌入公式
# 14: 'isolated'} # 单行公式
for xf in xf_json['layout_dets']:
L = xf['poly'][0] / horizontal_scale_ratio
U = xf['poly'][1] / vertical_scale_ratio
R = xf['poly'][2] / horizontal_scale_ratio
D = xf['poly'][5] / vertical_scale_ratio
# L += pageL # 有的页面，artBox偏移了。不在（0,0）
# R += pageL
# U += pageU
# D += pageU
L, R = min(L, R), max(L, R)
U, D = min(U, D), max(U, D)
# if xf['category_id'] == 5 and xf['score'] >= 0.3:
if xf['category_id'] == 5 and xf['score'] >= 0.43: # 新的footnote阈值
footnote_bbox_from_DocXChain.append((L, U, R, D))
footnote_final_names = []
footnote_final_bboxs = []
footnote_ID = 0
for L, U, R, D in footnote_bbox_from_DocXChain:
if debug_mode:
# cur_footnote = page.get_pixmap(clip=(L,U,R,D))
new_footnote_name = "footnote_{}_{}.png".format(page_ID, footnote_ID) # 脚注name
# cur_footnote.save(md_bookname_save_path + '/' + new_footnote_name) # 把脚注存储在新建的文件夹，并命名
footnote_final_names.append(new_footnote_name) # 把脚注的名字存在list中
footnote_final_bboxs.append((L, U, R, D))
footnote_ID += 1
footnote_final_bboxs.sort(key = lambda LURD: (LURD[1], LURD[0]))
curPage_all_footnote_bboxs = footnote_final_bboxs
return curPage_all_footnote_bboxs
def need_remove(block):
if 'lines' in block and len(block['lines']) > 0:
# block中只有一行，且该行文本全是大写字母，或字体为粗体bold关键词，SB关键词，把这个block捞回来
if len(block['lines']) == 1:
if 'spans' in block['lines'][0] and len(block['lines'][0]['spans']) == 1:
font_keywords = ['SB', 'bold', 'Bold']
if block['lines'][0]['spans'][0]['text'].isupper() or any(keyword in block['lines'][0]['spans'][0]['font'] for keyword in font_keywords):
return True
for line in block['lines']:
if 'spans' in line and len(line['spans']) > 0:
for span in line['spans']:
# 检测"keyword"是否在span中，忽略大小写
if "keyword" in span['text'].lower():
return True
return False
def parse_footnotes_by_rule(remain_text_blocks, page_height, page_id, main_text_font):
"""
根据给定的文本块、页高和页码，解析出符合规则的脚注文本块，并返回其边界框。
Args:
remain_text_blocks (list): 包含所有待处理的文本块的列表。
page_height (float): 页面的高度。
page_id (int): 页面的ID。
Returns:
list: 符合规则的脚注文本块的边界框列表。
"""
# if page_id > 20:
if page_id > 2: # 为保证精确度，先只筛选前3页
return []
else:
# 存储每一行的文本块大小的列表
line_sizes = []
# 存储每个文本块的平均行大小
block_sizes = []
# 存储每一行的字体信息
# font_names = []
font_names = Counter()
if len(remain_text_blocks) > 0:
for block in remain_text_blocks:
block_line_sizes = []
# block_fonts = []
block_fonts = Counter()
for line in block['lines']:
# 提取每个span的size属性，并计算行大小
span_sizes = [span['size'] for span in line['spans'] if 'size' in span]
if span_sizes:
line_size = sum(span_sizes) / len(span_sizes)
line_sizes.append(line_size)
block_line_sizes.append(line_size)
span_font = [(span['font'], len(span['text'])) for span in line['spans'] if 'font' in span and len(span['text']) > 0]
if span_font:
# main_text_font应该用基于字数最多的字体而不是span级别的统计
# font_names.append(font_name for font_name in span_font)
# block_fonts.append(font_name for font_name in span_font)
for font, count in span_font:
# font_names.extend([font] * count)
# block_fonts.extend([font] * count)
font_names[font] += count
block_fonts[font] += count
if block_line_sizes:
# 计算文本块的平均行大小
block_size = sum(block_line_sizes) / len(block_line_sizes)
# block_font = collections.Counter(block_fonts).most_common(1)[0][0]
block_font = block_fonts.most_common(1)[0][0]
block_sizes.append((block, block_size, block_font))
# 计算main_text_size
main_text_size = Counter(line_sizes).most_common(1)[0][0]
# 计算main_text_font
# main_text_font = collections.Counter(font_names).most_common(1)[0][0]
# main_text_font = font_names.most_common(1)[0][0]
# 删除一些可能被误识别为脚注的文本块
block_sizes = [(block, block_size, block_font) for block, block_size, block_font in block_sizes if not need_remove(block)]
# 检测footnote_block 并返回 footnote_bboxes
# footnote_bboxes = [block['bbox'] for block, block_size, block_font in block_sizes if
# block['bbox'][1] > page_height * 0.6 and block_size < main_text_size
# and (len(block['lines']) < 5 or block_font != main_text_font)]
# and len(block['lines']) < 5]
footnote_bboxes = [block['bbox'] for block, block_size, block_font in block_sizes if
block['bbox'][1] > page_height * 0.6 and
# 较为严格的规则
block_size < main_text_size and
(len(block['lines']) < 5 or
block_font != main_text_font)]
# 较为宽松的规则
# sum([block_size < main_text_size,
# len(block['lines']) < 5,
# block_font != main_text_font])
# >= 2]
return footnote_bboxes
else:
return []
magic_pdf/pre_proc/detect_header.py.bak deleted 100644 → 0
View file @ 79b58a1e
from magic_pdf.libs.commons import fitz # pyMuPDF库
from magic_pdf.libs.coordinate_transform import get_scale_ratio
def parse_headers(page_ID: int, page: fitz.Page, json_from_DocXchain_obj: dict):
"""
:param page_ID: int类型，当前page在当前pdf文档中是第page_D页。
:param page :fitz读取的当前页的内容
:param res_dir_path: str类型，是每一个pdf文档，在当前.py文件的目录下生成一个与pdf文档同名的文件夹，res_dir_path就是文件夹的dir
:param json_from_DocXchain_obj: dict类型，把pdf文档送入DocXChain模型中后，提取bbox，结果保存到pdf文档同名文件夹下的 page_ID.json文件中了。json_from_DocXchain_obj就是打开后的dict
"""
#--------- 通过json_from_DocXchain来获取 header ---------#
header_bbox_from_DocXChain = []
xf_json = json_from_DocXchain_obj
horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(xf_json, page)
# {0: 'title', # 标题
# 1: 'figure', # 图片
# 2: 'plain text', # 文本
# 3: 'header', # 页眉
# 4: 'page number', # 页码
# 5: 'footnote', # 脚注
# 6: 'footer', # 页脚
# 7: 'table', # 表格
# 8: 'table caption', # 表格描述
# 9: 'figure caption', # 图片描述
# 10: 'equation', # 公式
# 11: 'full column', # 单栏
# 12: 'sub column', # 多栏
# 13: 'embedding', # 嵌入公式
# 14: 'isolated'} # 单行公式
for xf in xf_json['layout_dets']:
L = xf['poly'][0] / horizontal_scale_ratio
U = xf['poly'][1] / vertical_scale_ratio
R = xf['poly'][2] / horizontal_scale_ratio
D = xf['poly'][5] / vertical_scale_ratio
# L += pageL # 有的页面，artBox偏移了。不在（0,0）
# R += pageL
# U += pageU
# D += pageU
L, R = min(L, R), max(L, R)
U, D = min(U, D), max(U, D)
if xf['category_id'] == 3 and xf['score'] >= 0.3:
header_bbox_from_DocXChain.append((L, U, R, D))
header_final_names = []
header_final_bboxs = []
header_ID = 0
for L, U, R, D in header_bbox_from_DocXChain:
# cur_header = page.get_pixmap(clip=(L,U,R,D))
new_header_name = "header_{}_{}.png".format(page_ID, header_ID) # 页眉name
# cur_header.save(res_dir_path + '/' + new_header_name) # 把页眉存储在新建的文件夹，并命名
header_final_names.append(new_header_name) # 把页面的名字存在list中
header_final_bboxs.append((L, U, R, D))
header_ID += 1
header_final_bboxs.sort(key = lambda LURD: (LURD[1], LURD[0]))
curPage_all_header_bboxs = header_final_bboxs
return curPage_all_header_bboxs
magic_pdf/pre_proc/detect_images.py.bak deleted 100644 → 0
View file @ 79b58a1e
import collections # 统计库
import re
from magic_pdf.libs.commons import fitz # pyMuPDF库
#--------------------------------------- Tool Functions --------------------------------------#
# 正则化，输入文本，输出只保留a-z,A-Z,0-9
def remove_special_chars(s: str) -> str:
pattern = r"[^a-zA-Z0-9]"
res = re.sub(pattern, "", s)
return res
def check_rect1_sameWith_rect2(L1: float, U1: float, R1: float, D1: float, L2: float, U2: float, R2: float, D2: float) -> bool:
# 判断rect1和rect2是否一模一样
return L1 == L2 and U1 == U2 and R1 == R2 and D1 == D2
def check_rect1_contains_rect2(L1: float, U1: float, R1: float, D1: float, L2: float, U2: float, R2: float, D2: float) -> bool:
# 判断rect1包含了rect2
return (L1 <= L2 <= R2 <= R1) and (U1 <= U2 <= D2 <= D1)
def check_rect1_overlaps_rect2(L1: float, U1: float, R1: float, D1: float, L2: float, U2: float, R2: float, D2: float) -> bool:
# 判断rect1与rect2是否存在重叠（只有一条边重叠，也算重叠）
return max(L1, L2) <= min(R1, R2) and max(U1, U2) <= min(D1, D2)
def calculate_overlapRatio_between_rect1_and_rect2(L1: float, U1: float, R1: float, D1: float, L2: float, U2: float, R2: float, D2: float) -> (float, float):
# 计算两个rect，重叠面积各占2个rect面积的比例
if min(R1, R2) < max(L1, L2) or min(D1, D2) < max(U1, U2):
return 0, 0
square_1 = (R1 - L1) * (D1 - U1)
square_2 = (R2 - L2) * (D2 - U2)
if square_1 == 0 or square_2 == 0:
return 0, 0
square_overlap = (min(R1, R2) - max(L1, L2)) * (min(D1, D2) - max(U1, U2))
return square_overlap / square_1, square_overlap / square_2
def calculate_overlapRatio_between_line1_and_line2(L1: float, R1: float, L2: float, R2: float) -> (float, float):
# 计算两个line，重叠区间各占2个line长度的比例
if max(L1, L2) > min(R1, R2):
return 0, 0
if L1 == R1 or L2 == R2:
return 0, 0
overlap_line = min(R1, R2) - max(L1, L2)
return overlap_line / (R1 - L1), overlap_line / (R2 - L2)
# 判断rect其实是一条line
def check_rect_isLine(L: float, U: float, R: float, D: float) -> bool:
width = R - L
height = D - U
if width <= 3 or height <= 3:
return True
if width / height >= 30 or height / width >= 30:
return True
def parse_images(page_ID: int, page: fitz.Page, json_from_DocXchain_obj: dict, junk_img_bojids=[]):
"""
:param page_ID: int类型，当前page在当前pdf文档中是第page_D页。
:param page :fitz读取的当前页的内容
:param res_dir_path: str类型，是每一个pdf文档，在当前.py文件的目录下生成一个与pdf文档同名的文件夹，res_dir_path就是文件夹的dir
:param json_from_DocXchain_obj: dict类型，把pdf文档送入DocXChain模型中后，提取bbox，结果保存到pdf文档同名文件夹下的 page_ID.json文件中了。json_from_DocXchain_obj就是打开后的dict
"""
#### 通过fitz获取page信息
## 超越边界
DPI = 72 # use this resolution
pix = page.get_pixmap(dpi=DPI)
pageL = 0
pageR = int(pix.w)
pageU = 0
pageD = int(pix.h)
#----------------- 保存每一个文本块的LURD ------------------#
textLine_blocks = []
blocks = page.get_text(
"dict",
flags=fitz.TEXTFLAGS_TEXT,
#clip=clip,
)["blocks"]
for i in range(len(blocks)):
bbox = blocks[i]['bbox']
# print(bbox)
for tt in blocks[i]['lines']:
# 当前line
cur_line_bbox = None # 当前line，最右侧的section的bbox
for xf in tt['spans']:
L, U, R, D = xf['bbox']
L, R = min(L, R), max(L, R)
U, D = min(U, D), max(U, D)
textLine_blocks.append((L, U, R, D))
textLine_blocks.sort(key = lambda LURD: (LURD[1], LURD[0]))
#---------------------------------------------- 保存img --------------------------------------------------#
raw_imgs = page.get_images() # 获取所有的图片
imgs = []
img_names = [] # 保存图片的名字，方便在md中插入引用
img_bboxs = [] # 保存图片的location信息。
img_visited = [] # 记忆化，记录该图片是否在md中已经插入过了
img_ID = 0
## 获取、保存每张img的location信息(x1, y1, x2, y2， UL, DR坐标)
for i in range(len(raw_imgs)):
# 如果图片在junklist中则跳过
if raw_imgs[i][0] in junk_img_bojids:
continue
else:
try:
tt = page.get_image_rects(raw_imgs[i][0], transform = True)
rec = tt[0][0]
L, U, R, D = int(rec[0]), int(rec[1]), int(rec[2]), int(rec[3])
L, R = min(L, R), max(L, R)
U, D = min(U, D), max(U, D)
if not(pageL <= L < R <= pageR and pageU <= U < D <= pageD):
continue
if pageL == L and R == pageR:
continue
if pageU == U and D == pageD:
continue
# pix1 = page.get_Pixmap(clip=(L,U,R,D))
new_img_name = "{}_{}.png".format(page_ID, i) # 图片name
# pix1.save(res_dir_path + '/' + new_img_name) # 把图片存出在新建的文件夹，并命名
img_names.append(new_img_name)
img_bboxs.append((L, U, R, D))
img_visited.append(False)
imgs.append(raw_imgs[i])
except:
continue
#-------- 如果img之间有重叠。说明获取的img大小有问题，位置也不一定对。就扔掉--------#
imgs_ok = [True for _ in range(len(imgs))]
for i in range(len(imgs)):
L1, U1, R1, D1 = img_bboxs[i]
for j in range(i + 1, len(imgs)):
L2, U2, R2, D2 = img_bboxs[j]
ratio_1, ratio_2 = calculate_overlapRatio_between_rect1_and_rect2(L1, U1, R1, D1, L2, U2, R2, D2)
s1 = abs(R1 - L1) * abs(D1 - U1)
s2 = abs(R2 - L2) * abs(D2 - U2)
if ratio_1 > 0 and ratio_2 > 0:
if ratio_1 == 1 and ratio_2 > 0.8:
imgs_ok[i] = False
elif ratio_1 > 0.8 and ratio_2 == 1:
imgs_ok[j] = False
elif s1 > 20000 and s2 > 20000 and ratio_1 > 0.4 and ratio_2 > 0.4:
imgs_ok[i] = False
imgs_ok[j] = False
elif s1 / s2 > 5 and ratio_2 > 0.5:
imgs_ok[j] = False
elif s2 / s1 > 5 and ratio_1 > 0.5:
imgs_ok[i] = False
imgs = [imgs[i] for i in range(len(imgs)) if imgs_ok[i] == True]
img_names = [img_names[i] for i in range(len(imgs)) if imgs_ok[i] == True]
img_bboxs = [img_bboxs[i] for i in range(len(imgs)) if imgs_ok[i] == True]
img_visited = [img_visited[i] for i in range(len(imgs)) if imgs_ok[i] == True]
#*******************************************************************************#
#---------------------------------------- 通过fitz提取svg的信息 -----------------------------------------#
#
svgs = page.get_drawings()
#------------ preprocess, check一些大框，看是否是合理的 ----------#
## 去重。有时候会遇到rect1和rect2是完全一样的情形。
svg_rect_visited = set()
available_svgIdx = []
for i in range(len(svgs)):
L, U, R, D = svgs[i]['rect'].irect
L, R = min(L, R), max(L, R)
U, D = min(U, D), max(U, D)
tt = (L, U, R, D)
if tt not in svg_rect_visited:
svg_rect_visited.add(tt)
available_svgIdx.append(i)
svgs = [svgs[i] for i in available_svgIdx] # 去重后，有效的svgs
svg_childs = [[] for _ in range(len(svgs))]
svg_parents = [[] for _ in range(len(svgs))]
svg_overlaps = [[] for _ in range(len(svgs))] #svg_overlaps[i]是一个list，存的是与svg_i有重叠的svg的index。e.g., svg_overlaps[0] = [1, 2, 7, 9]
svg_visited = [False for _ in range(len(svgs))]
svg_exceedPage = [0 for _ in range(len(svgs))] # 是否超越边界（artbox），很大，但一般是一个svg的底。
for i in range(len(svgs)):
L, U, R, D = svgs[i]['rect'].irect
ratio_1, ratio_2 = calculate_overlapRatio_between_rect1_and_rect2(L, U, R, D, pageL, pageU, pageR, pageD)
if (pageL + 20 < L <= R < pageR - 20) and (pageU + 20 < U <= D < pageD - 20):
if ratio_2 >= 0.7:
svg_exceedPage[i] += 4
else:
if L <= pageL:
svg_exceedPage[i] += 1
if pageR <= R:
svg_exceedPage[i] += 1
if U <= pageU:
svg_exceedPage[i] += 1
if pageD <= D:
svg_exceedPage[i] += 1
#### 如果有≥2个的超边界的框，就不要手写规则判断svg了。很难写对。
if len([x for x in svg_exceedPage if x >= 1]) >= 2:
svgs = []
svg_childs = []
svg_parents = []
svg_overlaps = []
svg_visited = []
svg_exceedPage = []
#---------------------------- build graph ----------------------------#
for i, p in enumerate(svgs):
L1, U1, R1, D1 = svgs[i]["rect"].irect
for j in range(len(svgs)):
if i == j:
continue
L2, U2, R2, D2 = svgs[j]["rect"].irect
## 包含
if check_rect1_contains_rect2(L1, U1, R1, D1, L2, U2, R2, D2) == True:
svg_childs[i].append(j)
svg_parents[j].append(i)
else:
## 交叉
if check_rect1_overlaps_rect2(L1, U1, R1, D1, L2, U2, R2, D2) == True:
svg_overlaps[i].append(j)
#---------------- 确定最终的svg。连通块儿的外围 -------------------#
eps_ERROR = 5 # 给识别出的svg，四周留白（为了防止pyMuPDF的rect不准）
svg_ID = 0
svg_final_names = []
svg_final_bboxs = []
svg_final_visited = [] # 为下面，text识别左准备。作用同img_visited
svg_idxs = [i for i in range(len(svgs))]
svg_idxs.sort(key = lambda i: -(svgs[i]['rect'].irect[2] - svgs[i]['rect'].irect[0]) * (svgs[i]['rect'].irect[3] - svgs[i]['rect'].irect[1])) # 按照面积，从大到小排序
for i in svg_idxs:
if svg_visited[i] == True:
continue
svg_visited[i] = True
L, U, R, D = svgs[i]['rect'].irect
width = R - L
height = D - U
if check_rect_isLine(L, U, R, D) == True:
svg_visited[i] = False
continue
# if i == 4:
# print(i, L, U, R, D)
# print(svg_parents[i])
cur_block_element_cnt = 0 # 当前要判定为svg的区域中，有多少elements，最外围的最大svg框除外。
if len(svg_parents[i]) == 0:
## 是个普通框的情形
cur_block_element_cnt += len(svg_childs[i])
if svg_exceedPage[i] == 0:
## 误差。可能已经包含在某个框里面了
neglect_flag = False
for pL, pU, pR, pD in svg_final_bboxs:
if pL <= L <= R <= pR and pU <= U <= D <= pD:
neglect_flag = True
break
if neglect_flag == True:
continue
## 搜索连通域, bfs+记忆化
q = collections.deque()
for j in svg_overlaps[i]:
q.append(j)
while q:
j = q.popleft()
svg_visited[j] = True
L2, U2, R2, D2 = svgs[j]['rect'].irect
# width2 = R2 - L2
# height2 = D2 - U2
# if width2 <= 2 or height2 <= 2 or (height2 / width2) >= 30 or (width2 / height2) >= 30:
# continue
L = min(L, L2)
R = max(R, R2)
U = min(U, U2)
D = max(D, D2)
cur_block_element_cnt += 1
cur_block_element_cnt += len(svg_childs[j])
for k in svg_overlaps[j]:
if svg_visited[k] == False and svg_exceedPage[k] == 0:
svg_visited[k] = True
q.append(k)
elif svg_exceedPage[i] <= 2:
## 误差。可能已经包含在某个svg_final_bbox框里面了
neglect_flag = False
for sL, sU, sR, sD in svg_final_bboxs:
if sL <= L <= R <= sR and sU <= U <= D <= sD:
neglect_flag = True
break
if neglect_flag == True:
continue
L, U, R, D = pageR, pageD, pageL, pageU
## 所有孩子元素的最大边界
for j in svg_childs[i]:
if svg_visited[j] == True:
continue
if svg_exceedPage[j] >= 1:
continue
svg_visited[j] = True #### 这个位置考虑一下
L2, U2, R2, D2 = svgs[j]['rect'].irect
L = min(L, L2)
R = max(R, R2)
U = min(U, U2)
D = max(D, D2)
cur_block_element_cnt += 1
# 如果是条line，就不用保存了
if check_rect_isLine(L, U, R, D) == True:
continue
# 如果当前的svg，连2个elements都没有，就不用保存了
if cur_block_element_cnt < 3:
continue
## 当前svg，框住了多少文本框。如果框多了，可能就是错了
contain_textLineBlock_cnt = 0
for L2, U2, R2, D2 in textLine_blocks:
if check_rect1_contains_rect2(L, U, R, D, L2, U2, R2, D2) == True:
contain_textLineBlock_cnt += 1
if contain_textLineBlock_cnt >= 10:
continue
# L -= eps_ERROR * 2
# U -= eps_ERROR
# R += eps_ERROR * 2
# D += eps_ERROR
# # cur_svg = page.get_pixmap(matrix=fitz.Identity, dpi=None, colorspace=fitz.csRGB, clip=(U,L,R,D), alpha=False, annots=True)
# cur_svg = page.get_pixmap(clip=(L,U,R,D))
new_svg_name = "svg_{}_{}.png".format(page_ID, svg_ID) # 图片name
# cur_svg.save(res_dir_path + '/' + new_svg_name) # 把图片存出在新建的文件夹，并命名
svg_final_names.append(new_svg_name) # 把图片的名字存在list中，方便在md中插入引用
svg_final_bboxs.append((L, U, R, D))
svg_final_visited.append(False)
svg_ID += 1
## 识别出的svg，可能有 包含，相邻的情形。需要进一步合并
svg_idxs = [i for i in range(len(svg_final_bboxs))]
svg_idxs.sort(key = lambda i: (svg_final_bboxs[i][1], svg_final_bboxs[i][0])) # (U, L)
svg_final_names_2 = []
svg_final_bboxs_2 = []
svg_final_visited_2 = [] # 为下面，text识别左准备。作用同img_visited
svg_ID_2 = 0
for i in range(len(svg_final_bboxs)):
L1, U1, R1, D1 = svg_final_bboxs[i]
for j in range(i + 1, len(svg_final_bboxs)):
L2, U2, R2, D2 = svg_final_bboxs[j]
# 如果 rect1包含了rect2
if check_rect1_contains_rect2(L1, U1, R1, D1, L2, U2, R2, D2) == True:
svg_final_visited[j] = True
continue
# 水平并列
ratio_1, ratio_2 = calculate_overlapRatio_between_line1_and_line2(U1, D1, U2, D2)
if ratio_1 >= 0.7 and ratio_2 >= 0.7:
if abs(L2 - R1) >= 20:
continue
LL = min(L1, L2)
UU = min(U1, U2)
RR = max(R1, R2)
DD = max(D1, D2)
svg_final_bboxs[i] = (LL, UU, RR, DD)
svg_final_visited[j] = True
continue
# 竖直并列
ratio_1, ratio_2 = calculate_overlapRatio_between_line1_and_line2(L1, R2, L2, R2)
if ratio_1 >= 0.7 and ratio_2 >= 0.7:
if abs(U2 - D1) >= 20:
continue
LL = min(L1, L2)
UU = min(U1, U2)
RR = max(R1, R2)
DD = max(D1, D2)
svg_final_bboxs[i] = (LL, UU, RR, DD)
svg_final_visited[j] = True
for i in range(len(svg_final_bboxs)):
if svg_final_visited[i] == False:
L, U, R, D = svg_final_bboxs[i]
svg_final_bboxs_2.append((L, U, R, D))
L -= eps_ERROR * 2
U -= eps_ERROR
R += eps_ERROR * 2
D += eps_ERROR
# cur_svg = page.get_pixmap(clip=(L,U,R,D))
new_svg_name = "svg_{}_{}.png".format(page_ID, svg_ID_2) # 图片name
# cur_svg.save(res_dir_path + '/' + new_svg_name) # 把图片存出在新建的文件夹，并命名
svg_final_names_2.append(new_svg_name) # 把图片的名字存在list中，方便在md中插入引用
svg_final_bboxs_2.append((L, U, R, D))
svg_final_visited_2.append(False)
svg_ID_2 += 1
## svg收尾。识别为drawing，但是在上面没有拼成一张图的。
# 有收尾才comprehensive
# xxxx
# xxxx
# xxxx
# xxxx
#--------- 通过json_from_DocXchain来获取，figure, table, equation的bbox ---------#
figure_bbox_from_DocXChain = []
figure_from_DocXChain_visited = [] # 记忆化
figure_bbox_from_DocXChain_overlappedRatio = []
figure_only_from_DocXChain_bboxs = [] # 存储
figure_only_from_DocXChain_names = []
figure_only_from_DocXChain_visited = []
figure_only_ID = 0
xf_json = json_from_DocXchain_obj
width_from_json = xf_json['page_info']['width']
height_from_json = xf_json['page_info']['height']
LR_scaleRatio = width_from_json / (pageR - pageL)
UD_scaleRatio = height_from_json / (pageD - pageU)
for xf in xf_json['layout_dets']:
# {0: 'title', 1: 'figure', 2: 'plain text', 3: 'header', 4: 'page number', 5: 'footnote', 6: 'footer', 7: 'table', 8: 'table caption', 9: 'figure caption', 10: 'equation', 11: 'full column', 12: 'sub column'}
L = xf['poly'][0] / LR_scaleRatio
U = xf['poly'][1] / UD_scaleRatio
R = xf['poly'][2] / LR_scaleRatio
D = xf['poly'][5] / UD_scaleRatio
# L += pageL # 有的页面，artBox偏移了。不在（0,0）
# R += pageL
# U += pageU
# D += pageU
L, R = min(L, R), max(L, R)
U, D = min(U, D), max(U, D)
# figure
if xf["category_id"] == 1 and xf['score'] >= 0.3:
figure_bbox_from_DocXChain.append((L, U, R, D))
figure_from_DocXChain_visited.append(False)
figure_bbox_from_DocXChain_overlappedRatio.append(0.0)
#---------------------- 比对上面识别出来的img,svg 与DocXChain给的figure -----------------------#
## 比对imgs
for i, b1 in enumerate(figure_bbox_from_DocXChain):
# print('--------- DocXChain的图片', b1)
L1, U1, R1, D1 = b1
for b2 in img_bboxs:
# print('-------- igms得到的图', b2)
L2, U2, R2, D2 = b2
s1 = abs(R1 - L1) * abs(D1 - U1)
s2 = abs(R2 - L2) * abs(D2 - U2)
# 相同
if check_rect1_sameWith_rect2(L1, U1, R1, D1, L2, U2, R2, D2) == True:
figure_from_DocXChain_visited[i] = True
# 包含
elif check_rect1_contains_rect2(L1, U1, R1, D1, L2, U2, R2, D2) == True:
if s2 / s1 > 0.8:
figure_from_DocXChain_visited[i] = True
elif check_rect1_contains_rect2(L2, U2, R2, D2, L1, U1, R1, D1) == True:
if s1 / s2 > 0.8:
figure_from_DocXChain_visited[i] = True
else:
# 重叠了相当一部分
# print('进入第3部分')
ratio_1, ratio_2 = calculate_overlapRatio_between_rect1_and_rect2(L1, U1, R1, D1, L2, U2, R2, D2)
if (ratio_1 >= 0.6 and ratio_2 >= 0.6) or (ratio_1 >= 0.8 and s1/s2>0.8) or (ratio_2 >= 0.8 and s2/s1>0.8):
figure_from_DocXChain_visited[i] = True
else:
figure_bbox_from_DocXChain_overlappedRatio[i] += ratio_1
# print('图片的重叠率是{}'.format(ratio_1))
## 比对svgs
svg_final_bboxs_2_badIdxs = []
for i, b1 in enumerate(figure_bbox_from_DocXChain):
L1, U1, R1, D1 = b1
for j, b2 in enumerate(svg_final_bboxs_2):
L2, U2, R2, D2 = b2
s1 = abs(R1 - L1) * abs(D1 - U1)
s2 = abs(R2 - L2) * abs(D2 - U2)
# 相同
if check_rect1_sameWith_rect2(L1, U1, R1, D1, L2, U2, R2, D2) == True:
figure_from_DocXChain_visited[i] = True
# 包含
elif check_rect1_contains_rect2(L1, U1, R1, D1, L2, U2, R2, D2) == True:
figure_from_DocXChain_visited[i] = True
elif check_rect1_contains_rect2(L2, U2, R2, D2, L1, U1, R1, D1) == True:
if s1 / s2 > 0.7:
figure_from_DocXChain_visited[i] = True
else:
svg_final_bboxs_2_badIdxs.append(j) # svg丢弃。用DocXChain的结果。
else:
# 重叠了相当一部分
ratio_1, ratio_2 = calculate_overlapRatio_between_rect1_and_rect2(L1, U1, R1, D1, L2, U2, R2, D2)
if (ratio_1 >= 0.5 and ratio_2 >= 0.5) or (min(ratio_1, ratio_2) >= 0.4 and max(ratio_1, ratio_2) >= 0.6):
figure_from_DocXChain_visited[i] = True
else:
figure_bbox_from_DocXChain_overlappedRatio[i] += ratio_1
# 丢掉错误的svg
svg_final_bboxs_2 = [svg_final_bboxs_2[i] for i in range(len(svg_final_bboxs_2)) if i not in set(svg_final_bboxs_2_badIdxs)]
for i in range(len(figure_from_DocXChain_visited)):
if figure_bbox_from_DocXChain_overlappedRatio[i] >= 0.7:
figure_from_DocXChain_visited[i] = True
# DocXChain识别出来的figure，但是没被保存的。
for i in range(len(figure_from_DocXChain_visited)):
if figure_from_DocXChain_visited[i] == False:
figure_from_DocXChain_visited[i] = True
cur_bbox = figure_bbox_from_DocXChain[i]
# cur_figure = page.get_pixmap(clip=cur_bbox)
new_figure_name = "figure_only_{}_{}.png".format(page_ID, figure_only_ID) # 图片name
# cur_figure.save(res_dir_path + '/' + new_figure_name) # 把图片存出在新建的文件夹，并命名
figure_only_from_DocXChain_names.append(new_figure_name) # 把图片的名字存在list中，方便在md中插入引用
figure_only_from_DocXChain_bboxs.append(cur_bbox)
figure_only_from_DocXChain_visited.append(False)
figure_only_ID += 1
img_bboxs.sort(key = lambda LURD: (LURD[1], LURD[0]))
svg_final_bboxs_2.sort(key = lambda LURD: (LURD[1], LURD[0]))
figure_only_from_DocXChain_bboxs.sort(key = lambda LURD: (LURD[1], LURD[0]))
curPage_all_fig_bboxs = img_bboxs + svg_final_bboxs + figure_only_from_DocXChain_bboxs
#--------------------------- 最后统一去重 -----------------------------------#
curPage_all_fig_bboxs.sort(key = lambda LURD: ( (LURD[2]-LURD[0])*(LURD[3]-LURD[1]) , LURD[0], LURD[1]) )
#### 先考虑包含关系的小块
final_duplicate = set()
for i in range(len(curPage_all_fig_bboxs)):
L1, U1, R1, D1 = curPage_all_fig_bboxs[i]
for j in range(len(curPage_all_fig_bboxs)):
if i == j:
continue
L2, U2, R2, D2 = curPage_all_fig_bboxs[j]
s1 = abs(R1 - L1) * abs(D1 - U1)
s2 = abs(R2 - L2) * abs(D2 - U2)
if check_rect1_contains_rect2(L2, U2, R2, D2, L1, U1, R1, D1) == True:
final_duplicate.add((L1, U1, R1, D1))
else:
ratio_1, ratio_2 = calculate_overlapRatio_between_rect1_and_rect2(L1, U1, R1, D1, L2, U2, R2, D2)
if ratio_1 >= 0.8 and ratio_2 <= 0.6:
final_duplicate.add((L1, U1, R1, D1))
curPage_all_fig_bboxs = [LURD for LURD in curPage_all_fig_bboxs if LURD not in final_duplicate]
#### 再考虑重叠关系的块
final_duplicate = set()
final_synthetic_bboxs = []
for i in range(len(curPage_all_fig_bboxs)):
L1, U1, R1, D1 = curPage_all_fig_bboxs[i]
for j in range(len(curPage_all_fig_bboxs)):
if i == j:
continue
L2, U2, R2, D2 = curPage_all_fig_bboxs[j]
s1 = abs(R1 - L1) * abs(D1 - U1)
s2 = abs(R2 - L2) * abs(D2 - U2)
ratio_1, ratio_2 = calculate_overlapRatio_between_rect1_and_rect2(L1, U1, R1, D1, L2, U2, R2, D2)
union_ok = False
if (ratio_1 >= 0.8 and ratio_2 <= 0.6) or (ratio_1 > 0.6 and ratio_2 > 0.6):
union_ok = True
if (ratio_1 > 0.2 and s2 / s1 > 5):
union_ok = True
if (L1 <= (L2+R2)/2 <= R1) and (U1 <= (U2+D2)/2 <= D1):
union_ok = True
if (L2 <= (L1+R1)/2 <= R2) and (U2 <= (U1+D1)/2 <= D2):
union_ok = True
if union_ok == True:
final_duplicate.add((L1, U1, R1, D1))
final_duplicate.add((L2, U2, R2, D2))
L3, U3, R3, D3 = min(L1, L2), min(U1, U2), max(R1, R2), max(D1, D2)
final_synthetic_bboxs.append((L3, U3, R3, D3))
# print('---------- curPage_all_fig_bboxs ---------')
# print(curPage_all_fig_bboxs)
curPage_all_fig_bboxs = [b for b in curPage_all_fig_bboxs if b not in final_duplicate]
final_synthetic_bboxs = list(set(final_synthetic_bboxs))
## 再再考虑重叠关系。极端情况下会迭代式地2进1
new_images = []
droped_img_idx = []
image_bboxes = [[b[0], b[1], b[2], b[3]] for b in final_synthetic_bboxs]
for i in range(0, len(image_bboxes)):
for j in range(i+1, len(image_bboxes)):
if j not in droped_img_idx:
L2, U2, R2, D2 = image_bboxes[j]
s1 = abs(R1 - L1) * abs(D1 - U1)
s2 = abs(R2 - L2) * abs(D2 - U2)
ratio_1, ratio_2 = calculate_overlapRatio_between_rect1_and_rect2(L1, U1, R1, D1, L2, U2, R2, D2)
union_ok = False
if (ratio_1 >= 0.8 and ratio_2 <= 0.6) or (ratio_1 > 0.6 and ratio_2 > 0.6):
union_ok = True
if (ratio_1 > 0.2 and s2 / s1 > 5):
union_ok = True
if (L1 <= (L2+R2)/2 <= R1) and (U1 <= (U2+D2)/2 <= D1):
union_ok = True
if (L2 <= (L1+R1)/2 <= R2) and (U2 <= (U1+D1)/2 <= D2):
union_ok = True
if union_ok == True:
# 合并
image_bboxes[i][0], image_bboxes[i][1],image_bboxes[i][2],image_bboxes[i][3] = min(image_bboxes[i][0], image_bboxes[j][0]), min(image_bboxes[i][1], image_bboxes[j][1]), max(image_bboxes[i][2], image_bboxes[j][2]), max(image_bboxes[i][3], image_bboxes[j][3])
droped_img_idx.append(j)
for i in range(0, len(image_bboxes)):
if i not in droped_img_idx:
new_images.append(image_bboxes[i])
# find_union_FLAG = True
# while find_union_FLAG == True:
# find_union_FLAG = False
# final_duplicate = set()
# tmp = []
# for i in range(len(final_synthetic_bboxs)):
# L1, U1, R1, D1 = final_synthetic_bboxs[i]
# for j in range(len(final_synthetic_bboxs)):
# if i == j:
# continue
# L2, U2, R2, D2 = final_synthetic_bboxs[j]
# s1 = abs(R1 - L1) * abs(D1 - U1)
# s2 = abs(R2 - L2) * abs(D2 - U2)
# ratio_1, ratio_2 = calculate_overlapRatio_between_rect1_and_rect2(L1, U1, R1, D1, L2, U2, R2, D2)
# union_ok = False
# if (ratio_1 >= 0.8 and ratio_2 <= 0.6) or (ratio_1 > 0.6 and ratio_2 > 0.6):
# union_ok = True
# if (ratio_1 > 0.2 and s2 / s1 > 5):
# union_ok = True
# if (L1 <= (L2+R2)/2 <= R1) and (U1 <= (U2+D2)/2 <= D1):
# union_ok = True
# if (L2 <= (L1+R1)/2 <= R2) and (U2 <= (U1+D1)/2 <= D2):
# union_ok = True
# if union_ok == True:
# find_union_FLAG = True
# final_duplicate.add((L1, U1, R1, D1))
# final_duplicate.add((L2, U2, R2, D2))
# L3, U3, R3, D3 = min(L1, L2), min(U1, U2), max(R1, R2), max(D1, D2)
# tmp.append((L3, U3, R3, D3))
# if find_union_FLAG == True:
# tmp = list(set(tmp))
# final_synthetic_bboxs = tmp[:]
# curPage_all_fig_bboxs += final_synthetic_bboxs
# print('--------- final synthetic')
# print(final_synthetic_bboxs)
#**************************************************************************#
images1 = [[img[0], img[1], img[2], img[3]] for img in curPage_all_fig_bboxs]
images = images1 + new_images
return images
magic_pdf/pre_proc/detect_page_number.py.bak deleted 100644 → 0
View file @ 79b58a1e
from magic_pdf.libs.commons import fitz # pyMuPDF库
from magic_pdf.libs.coordinate_transform import get_scale_ratio
def parse_pageNos(page_ID: int, page: fitz.Page, json_from_DocXchain_obj: dict):
"""
:param page_ID: int类型，当前page在当前pdf文档中是第page_D页。
:param page :fitz读取的当前页的内容
:param res_dir_path: str类型，是每一个pdf文档，在当前.py文件的目录下生成一个与pdf文档同名的文件夹，res_dir_path就是文件夹的dir
:param json_from_DocXchain_obj: dict类型，把pdf文档送入DocXChain模型中后，提取bbox，结果保存到pdf文档同名文件夹下的 page_ID.json文件中了。json_from_DocXchain_obj就是打开后的dict
"""
#--------- 通过json_from_DocXchain来获取 pageNo ---------#
pageNo_bbox_from_DocXChain = []
xf_json = json_from_DocXchain_obj
horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(xf_json, page)
# {0: 'title', # 标题
# 1: 'figure', # 图片
# 2: 'plain text', # 文本
# 3: 'header', # 页眉
# 4: 'page number', # 页码
# 5: 'footnote', # 脚注
# 6: 'footer', # 页脚
# 7: 'table', # 表格
# 8: 'table caption', # 表格描述
# 9: 'figure caption', # 图片描述
# 10: 'equation', # 公式
# 11: 'full column', # 单栏
# 12: 'sub column', # 多栏
# 13: 'embedding', # 嵌入公式
# 14: 'isolated'} # 单行公式
for xf in xf_json['layout_dets']:
L = xf['poly'][0] / horizontal_scale_ratio
U = xf['poly'][1] / vertical_scale_ratio
R = xf['poly'][2] / horizontal_scale_ratio
D = xf['poly'][5] / vertical_scale_ratio
# L += pageL # 有的页面，artBox偏移了。不在（0,0）
# R += pageL
# U += pageU
# D += pageU
L, R = min(L, R), max(L, R)
U, D = min(U, D), max(U, D)
if xf['category_id'] == 4 and xf['score'] >= 0.3:
pageNo_bbox_from_DocXChain.append((L, U, R, D))
pageNo_final_names = []
pageNo_final_bboxs = []
pageNo_ID = 0
for L, U, R, D in pageNo_bbox_from_DocXChain:
# cur_pageNo = page.get_pixmap(clip=(L,U,R,D))
new_pageNo_name = "pageNo_{}_{}.png".format(page_ID, pageNo_ID) # 页码name
# cur_pageNo.save(res_dir_path + '/' + new_pageNo_name) # 把页码存储在新建的文件夹，并命名
pageNo_final_names.append(new_pageNo_name) # 把页码的名字存在list中
pageNo_final_bboxs.append((L, U, R, D))
pageNo_ID += 1
pageNo_final_bboxs.sort(key = lambda LURD: (LURD[1], LURD[0]))
curPage_all_pageNo_bboxs = pageNo_final_bboxs
return curPage_all_pageNo_bboxs
magic_pdf/pre_proc/detect_tables.py.bak deleted 100644 → 0
View file @ 79b58a1e
from magic_pdf.libs.commons import fitz # pyMuPDF库
def parse_tables(page_ID: int, page: fitz.Page, json_from_DocXchain_obj: dict):
"""
:param page_ID: int类型，当前page在当前pdf文档中是第page_D页。
:param page :fitz读取的当前页的内容
:param res_dir_path: str类型，是每一个pdf文档，在当前.py文件的目录下生成一个与pdf文档同名的文件夹，res_dir_path就是文件夹的dir
:param json_from_DocXchain_obj: dict类型，把pdf文档送入DocXChain模型中后，提取bbox，结果保存到pdf文档同名文件夹下的 page_ID.json文件中了。json_from_DocXchain_obj就是打开后的dict
"""
DPI = 72 # use this resolution
pix = page.get_pixmap(dpi=DPI)
pageL = 0
pageR = int(pix.w)
pageU = 0
pageD = int(pix.h)
#--------- 通过json_from_DocXchain来获取 table ---------#
table_bbox_from_DocXChain = []
xf_json = json_from_DocXchain_obj
width_from_json = xf_json['page_info']['width']
height_from_json = xf_json['page_info']['height']
LR_scaleRatio = width_from_json / (pageR - pageL)
UD_scaleRatio = height_from_json / (pageD - pageU)
for xf in xf_json['layout_dets']:
# {0: 'title', 1: 'figure', 2: 'plain text', 3: 'header', 4: 'page number', 5: 'footnote', 6: 'footer', 7: 'table', 8: 'table caption', 9: 'figure caption', 10: 'equation', 11: 'full column', 12: 'sub column'}
# 13: 'embedding', # 嵌入公式
# 14: 'isolated'} # 单行公式
L = xf['poly'][0] / LR_scaleRatio
U = xf['poly'][1] / UD_scaleRatio
R = xf['poly'][2] / LR_scaleRatio
D = xf['poly'][5] / UD_scaleRatio
# L += pageL # 有的页面，artBox偏移了。不在（0,0）
# R += pageL
# U += pageU
# D += pageU
L, R = min(L, R), max(L, R)
U, D = min(U, D), max(U, D)
if xf['category_id'] == 7 and xf['score'] >= 0.3:
table_bbox_from_DocXChain.append((L, U, R, D))
table_final_names = []
table_final_bboxs = []
table_ID = 0
for L, U, R, D in table_bbox_from_DocXChain:
# cur_table = page.get_pixmap(clip=(L,U,R,D))
new_table_name = "table_{}_{}.png".format(page_ID, table_ID) # 表格name
# cur_table.save(res_dir_path + '/' + new_table_name) # 把表格存出在新建的文件夹，并命名
table_final_names.append(new_table_name) # 把表格的名字存在list中，方便在md中插入引用
table_final_bboxs.append((L, U, R, D))
table_ID += 1
table_final_bboxs.sort(key = lambda LURD: (LURD[1], LURD[0]))
curPage_all_table_bboxs = table_final_bboxs
return curPage_all_table_bboxs
magic_pdf/pre_proc/equations_replace.py.bak deleted 100644 → 0
View file @ 79b58a1e
"""对pymupdf返回的结构里的公式进行替换，替换为模型识别的公式结果."""
import json
import os
from pathlib import Path
from loguru import logger
from magic_pdf.config.ocr_content_type import ContentType
from magic_pdf.libs.commons import fitz
TYPE_INLINE_EQUATION = ContentType.InlineEquation
TYPE_INTERLINE_EQUATION = ContentType.InterlineEquation
def combine_chars_to_pymudict(block_dict, char_dict):
"""把block级别的pymupdf 结构里加入char结构."""
# 因为block_dict 被裁剪过，因此先把他和char_dict文字块对齐，才能进行补充
char_map = {tuple(item['bbox']): item for item in char_dict}
for i in range(len(block_dict)): # block
block = block_dict[i]
key = block['bbox']
char_dict_item = char_map[tuple(key)]
char_dict_map = {tuple(item['bbox']): item for item in char_dict_item['lines']}
for j in range(len(block['lines'])):
lines = block['lines'][j]
with_char_lines = char_dict_map[lines['bbox']]
for k in range(len(lines['spans'])):
spans = lines['spans'][k]
try:
chars = with_char_lines['spans'][k]['chars']
except Exception:
logger.error(char_dict[i]['lines'][j])
spans['chars'] = chars
return block_dict
def calculate_overlap_area_2_minbox_area_ratio(bbox1, min_bbox):
"""计算box1和box2的重叠面积占最小面积的box的比例."""
# Determine the coordinates of the intersection rectangle
x_left = max(bbox1[0], min_bbox[0])
y_top = max(bbox1[1], min_bbox[1])
x_right = min(bbox1[2], min_bbox[2])
y_bottom = min(bbox1[3], min_bbox[3])
if x_right < x_left or y_bottom < y_top:
return 0.0
# The area of overlap area
intersection_area = (x_right - x_left) * (y_bottom - y_top)
min_box_area = (min_bbox[3] - min_bbox[1]) * (min_bbox[2] - min_bbox[0])
if min_box_area == 0:
return 0
else:
return intersection_area / min_box_area
def _is_xin(bbox1, bbox2):
area1 = abs(bbox1[2] - bbox1[0]) * abs(bbox1[3] - bbox1[1])
area2 = abs(bbox2[2] - bbox2[0]) * abs(bbox2[3] - bbox2[1])
if area1 < area2:
ratio = calculate_overlap_area_2_minbox_area_ratio(bbox2, bbox1)
else:
ratio = calculate_overlap_area_2_minbox_area_ratio(bbox1, bbox2)
return ratio > 0.6
def remove_text_block_in_interline_equation_bbox(interline_bboxes, text_blocks):
"""消除掉整个块都在行间公式块内部的文本块."""
for eq_bbox in interline_bboxes:
removed_txt_blk = []
for text_blk in text_blocks:
text_bbox = text_blk['bbox']
if (
calculate_overlap_area_2_minbox_area_ratio(eq_bbox['bbox'], text_bbox)
>= 0.7
):
removed_txt_blk.append(text_blk)
for blk in removed_txt_blk:
text_blocks.remove(blk)
return text_blocks
def _is_in_or_part_overlap(box1, box2) -> bool:
"""两个bbox是否有部分重叠或者包含."""
if box1 is None or box2 is None:
return False
x0_1, y0_1, x1_1, y1_1 = box1
x0_2, y0_2, x1_2, y1_2 = box2
return not (
x1_1 < x0_2 # box1在box2的左边
or x0_1 > x1_2 # box1在box2的右边
or y1_1 < y0_2 # box1在box2的上边
or y0_1 > y1_2
) # box1在box2的下边
def remove_text_block_overlap_interline_equation_bbox(
interline_eq_bboxes, pymu_block_list
):
"""消除掉行行内公式有部分重叠的文本块的内容。 同时重新计算消除重叠之后文本块的大小."""
deleted_block = []
for text_block in pymu_block_list:
deleted_line = []
for line in text_block['lines']:
deleted_span = []
for span in line['spans']:
deleted_chars = []
for char in span['chars']:
if any(
[
(
calculate_overlap_area_2_minbox_area_ratio(
eq_bbox['bbox'], char['bbox']
)
> 0.5
)
for eq_bbox in interline_eq_bboxes
]
):
deleted_chars.append(char)
# 检查span里没有char则删除这个span
for char in deleted_chars:
span['chars'].remove(char)
# 重新计算这个span的大小
if len(span['chars']) == 0: # 删除这个span
deleted_span.append(span)
else:
span['bbox'] = (
min([b['bbox'][0] for b in span['chars']]),
min([b['bbox'][1] for b in span['chars']]),
max([b['bbox'][2] for b in span['chars']]),
max([b['bbox'][3] for b in span['chars']]),
)
# 检查这个span
for span in deleted_span:
line['spans'].remove(span)
if len(line['spans']) == 0: # 删除这个line
deleted_line.append(line)
else:
line['bbox'] = (
min([b['bbox'][0] for b in line['spans']]),
min([b['bbox'][1] for b in line['spans']]),
max([b['bbox'][2] for b in line['spans']]),
max([b['bbox'][3] for b in line['spans']]),
)
# 检查这个block是否可以删除
for line in deleted_line:
text_block['lines'].remove(line)
if len(text_block['lines']) == 0: # 删除block
deleted_block.append(text_block)
else:
text_block['bbox'] = (
min([b['bbox'][0] for b in text_block['lines']]),
min([b['bbox'][1] for b in text_block['lines']]),
max([b['bbox'][2] for b in text_block['lines']]),
max([b['bbox'][3] for b in text_block['lines']]),
)
# 检查text block删除
for block in deleted_block:
pymu_block_list.remove(block)
if len(pymu_block_list) == 0:
return []
return pymu_block_list
def insert_interline_equations_textblock(interline_eq_bboxes, pymu_block_list):
"""在行间公式对应的地方插上一个伪造的block."""
for eq in interline_eq_bboxes:
bbox = eq['bbox']
latex_content = eq['latex']
text_block = {
'number': len(pymu_block_list),
'type': 0,
'bbox': bbox,
'lines': [
{
'spans': [
{
'size': 9.962599754333496,
'type': TYPE_INTERLINE_EQUATION,
'flags': 4,
'font': TYPE_INTERLINE_EQUATION,
'color': 0,
'ascender': 0.9409999847412109,
'descender': -0.3050000071525574,
'latex': latex_content,
'origin': [bbox[0], bbox[1]],
'bbox': bbox,
}
],
'wmode': 0,
'dir': [1.0, 0.0],
'bbox': bbox,
}
],
}
pymu_block_list.append(text_block)
def x_overlap_ratio(box1, box2):
a, _, c, _ = box1
e, _, g, _ = box2
# 计算重叠宽度
overlap_x = max(min(c, g) - max(a, e), 0)
# 计算box1的宽度
width1 = g - e
# 计算重叠比例
overlap_ratio = overlap_x / width1 if width1 != 0 else 0
return overlap_ratio
def __is_x_dir_overlap(bbox1, bbox2):
return not (bbox1[2] < bbox2[0] or bbox1[0] > bbox2[2])
def __y_overlap_ratio(box1, box2):
""""""
_, b, _, d = box1
_, f, _, h = box2
# 计算重叠高度
overlap_y = max(min(d, h) - max(b, f), 0)
# 计算box1的高度
height1 = d - b
# 计算重叠比例
overlap_ratio = overlap_y / height1 if height1 != 0 else 0
return overlap_ratio
def replace_line_v2(eqinfo, line):
"""扫描这一行所有的和公式框X方向重叠的char,然后计算char的左、右x0, x1,位于这个区间内的span删除掉。
最后与这个x0,x1有相交的span0, span1内部进行分割。"""
first_overlap_span = -1
first_overlap_span_idx = -1
last_overlap_span = -1
delete_chars = []
for i in range(0, len(line['spans'])):
if 'chars' not in line['spans'][i]:
continue
if line['spans'][i].get('_type', None) is not None:
continue # 忽略，因为已经是插入的伪造span公式了
for char in line['spans'][i]['chars']:
if __is_x_dir_overlap(eqinfo['bbox'], char['bbox']):
line_txt = ''
for span in line['spans']:
span_txt = '<span>'
for ch in span['chars']:
span_txt = span_txt + ch['c']
span_txt = span_txt + '</span>'
line_txt = line_txt + span_txt
if first_overlap_span_idx == -1:
first_overlap_span = line['spans'][i]
first_overlap_span_idx = i
last_overlap_span = line['spans'][i]
delete_chars.append(char)
# 第一个和最后一个char要进行检查，到底属于公式多还是属于正常span多
if len(delete_chars) > 0:
ch0_bbox = delete_chars[0]['bbox']
if x_overlap_ratio(eqinfo['bbox'], ch0_bbox) < 0.51:
delete_chars.remove(delete_chars[0])
if len(delete_chars) > 0:
ch0_bbox = delete_chars[-1]['bbox']
if x_overlap_ratio(eqinfo['bbox'], ch0_bbox) < 0.51:
delete_chars.remove(delete_chars[-1])
# 计算x方向上被删除区间内的char的真实x0, x1
if len(delete_chars):
x0, x1 = (
min([b['bbox'][0] for b in delete_chars]),
max([b['bbox'][2] for b in delete_chars]),
)
else:
# logger.debug(f"行内公式替换没有发生，尝试下一行匹配, eqinfo={eqinfo}")
return False
# 删除位于x0, x1这两个中间的span
delete_span = []
for span in line['spans']:
span_box = span['bbox']
if x0 <= span_box[0] and span_box[2] <= x1:
delete_span.append(span)
for span in delete_span:
line['spans'].remove(span)
equation_span = {
'size': 9.962599754333496,
'type': TYPE_INLINE_EQUATION,
'flags': 4,
'font': TYPE_INLINE_EQUATION,
'color': 0,
'ascender': 0.9409999847412109,
'descender': -0.3050000071525574,
'latex': '',
'origin': [337.1410153102337, 216.0205245153934],
'bbox': eqinfo['bbox'],
}
# equation_span = line['spans'][0].copy()
equation_span['latex'] = eqinfo['latex']
equation_span['bbox'] = [x0, equation_span['bbox'][1], x1, equation_span['bbox'][3]]
equation_span['origin'] = [equation_span['bbox'][0], equation_span['bbox'][1]]
equation_span['chars'] = delete_chars
equation_span['type'] = TYPE_INLINE_EQUATION
equation_span['_eq_bbox'] = eqinfo['bbox']
line['spans'].insert(first_overlap_span_idx + 1, equation_span) # 放入公式
# logger.info(f"==>text is 【{line_txt}】, equation is 【{eqinfo['latex_text']}】")
# 第一个、和最后一个有overlap的span进行分割,然后插入对应的位置
first_span_chars = [
char
for char in first_overlap_span['chars']
if (char['bbox'][2] + char['bbox'][0]) / 2 < x0
]
tail_span_chars = [
char
for char in last_overlap_span['chars']
if (char['bbox'][0] + char['bbox'][2]) / 2 > x1
]
if len(first_span_chars) > 0:
first_overlap_span['chars'] = first_span_chars
first_overlap_span['text'] = ''.join([char['c'] for char in first_span_chars])
first_overlap_span['bbox'] = (
first_overlap_span['bbox'][0],
first_overlap_span['bbox'][1],
max([chr['bbox'][2] for chr in first_span_chars]),
first_overlap_span['bbox'][3],
)
# first_overlap_span['_type'] = "first"
else:
# 删掉
if first_overlap_span not in delete_span:
line['spans'].remove(first_overlap_span)
if len(tail_span_chars) > 0:
min_of_tail_span_x0 = min([chr['bbox'][0] for chr in tail_span_chars])
min_of_tail_span_y0 = min([chr['bbox'][1] for chr in tail_span_chars])
max_of_tail_span_x1 = max([chr['bbox'][2] for chr in tail_span_chars])
max_of_tail_span_y1 = max([chr['bbox'][3] for chr in tail_span_chars])
if last_overlap_span == first_overlap_span: # 这个时候应该插入一个新的
tail_span_txt = ''.join([char['c'] for char in tail_span_chars]) # noqa: F841
last_span_to_insert = last_overlap_span.copy()
last_span_to_insert['chars'] = tail_span_chars
last_span_to_insert['text'] = ''.join(
[char['c'] for char in tail_span_chars]
)
if equation_span['bbox'][2] >= last_overlap_span['bbox'][2]:
last_span_to_insert['bbox'] = (
min_of_tail_span_x0,
min_of_tail_span_y0,
max_of_tail_span_x1,
max_of_tail_span_y1,
)
else:
last_span_to_insert['bbox'] = (
min([chr['bbox'][0] for chr in tail_span_chars]),
last_overlap_span['bbox'][1],
last_overlap_span['bbox'][2],
last_overlap_span['bbox'][3],
)
# 插入到公式对象之后
equation_idx = line['spans'].index(equation_span)
line['spans'].insert(equation_idx + 1, last_span_to_insert) # 放入公式
else: # 直接修改原来的span
last_overlap_span['chars'] = tail_span_chars
last_overlap_span['text'] = ''.join([char['c'] for char in tail_span_chars])
last_overlap_span['bbox'] = (
min([chr['bbox'][0] for chr in tail_span_chars]),
last_overlap_span['bbox'][1],
last_overlap_span['bbox'][2],
last_overlap_span['bbox'][3],
)
else:
# 删掉
if (
last_overlap_span not in delete_span
and last_overlap_span != first_overlap_span
):
line['spans'].remove(last_overlap_span)
remain_txt = ''
for span in line['spans']:
span_txt = '<span>'
for char in span['chars']:
span_txt = span_txt + char['c']
span_txt = span_txt + '</span>'
remain_txt = remain_txt + span_txt
# logger.info(f"<== succ replace, text is 【{remain_txt}】, equation is 【{eqinfo['latex_text']}】")
return True
def replace_eq_blk(eqinfo, text_block):
"""替换行内公式."""
for line in text_block['lines']:
line_bbox = line['bbox']
if (
_is_xin(eqinfo['bbox'], line_bbox)
or __y_overlap_ratio(eqinfo['bbox'], line_bbox) > 0.6
): # 定位到行, 使用y方向重合率是因为有的时候，一个行的宽度会小于公式位置宽度：行很高，公式很窄，
replace_succ = replace_line_v2(eqinfo, line)
if not replace_succ: # 有的时候，一个pdf的line高度从API里会计算的有问题，因此在行内span级别会替换不成功，这就需要继续重试下一行
continue
else:
break
else:
return False
return True
def replace_inline_equations(inline_equation_bboxes, raw_text_blocks):
"""替换行内公式."""
for eqinfo in inline_equation_bboxes:
eqbox = eqinfo['bbox']
for blk in raw_text_blocks:
if _is_xin(eqbox, blk['bbox']):
if not replace_eq_blk(eqinfo, blk):
logger.warning(f'行内公式没有替换成功：{eqinfo} ')
else:
break
return raw_text_blocks
def remove_chars_in_text_blocks(text_blocks):
"""删除text_blocks里的char."""
for blk in text_blocks:
for line in blk['lines']:
for span in line['spans']:
_ = span.pop('chars', 'no such key')
return text_blocks
def replace_equations_in_textblock(
raw_text_blocks, inline_equation_bboxes, interline_equation_bboxes
):
"""替换行间和和行内公式为latex."""
raw_text_blocks = remove_text_block_in_interline_equation_bbox(
interline_equation_bboxes, raw_text_blocks
) # 消除重叠：第一步，在公式内部的
raw_text_blocks = remove_text_block_overlap_interline_equation_bbox(
interline_equation_bboxes, raw_text_blocks
) # 消重，第二步，和公式覆盖的
insert_interline_equations_textblock(interline_equation_bboxes, raw_text_blocks)
raw_text_blocks = replace_inline_equations(inline_equation_bboxes, raw_text_blocks)
return raw_text_blocks
def draw_block_on_pdf_with_txt_replace_eq_bbox(json_path, pdf_path):
""""""
new_pdf = f'{Path(pdf_path).parent}/{Path(pdf_path).stem}.step3-消除行内公式text_block.pdf'
with open(json_path, 'r', encoding='utf-8') as f:
obj = json.loads(f.read())
if os.path.exists(new_pdf):
os.remove(new_pdf)
new_doc = fitz.open('')
doc = fitz.open(pdf_path) # noqa: F841
new_doc = fitz.open(pdf_path)
for i in range(len(new_doc)):
page = new_doc[i]
inline_equation_bboxes = obj[f'page_{i}']['inline_equations']
interline_equation_bboxes = obj[f'page_{i}']['interline_equations']
raw_text_blocks = obj[f'page_{i}']['preproc_blocks']
raw_text_blocks = remove_text_block_in_interline_equation_bbox(
interline_equation_bboxes, raw_text_blocks
) # 消除重叠：第一步，在公式内部的
raw_text_blocks = remove_text_block_overlap_interline_equation_bbox(
interline_equation_bboxes, raw_text_blocks
) # 消重，第二步，和公式覆盖的
insert_interline_equations_textblock(interline_equation_bboxes, raw_text_blocks)
raw_text_blocks = replace_inline_equations(
inline_equation_bboxes, raw_text_blocks
)
# 为了检验公式是否重复，把每一行里，含有公式的span背景改成黄色的
color_map = [fitz.pdfcolor['blue'], fitz.pdfcolor['green']] # noqa: F841
j = 0 # noqa: F841
for blk in raw_text_blocks:
for i, line in enumerate(blk['lines']):
# line_box = line['bbox']
# shape = page.new_shape()
# shape.draw_rect(line_box)
# shape.finish(color=fitz.pdfcolor['red'], fill=color_map[j%2], fill_opacity=0.3)
# shape.commit()
# j = j+1
for i, span in enumerate(line['spans']):
shape_page = page.new_shape()
span_type = span.get('_type')
color = fitz.pdfcolor['blue']
if span_type == 'first':
color = fitz.pdfcolor['blue']
elif span_type == 'tail':
color = fitz.pdfcolor['green']
elif span_type == TYPE_INLINE_EQUATION:
color = fitz.pdfcolor['black']
else:
color = None
b = span['bbox']
shape_page.draw_rect(b)
shape_page.finish(color=None, fill=color, fill_opacity=0.3)
shape_page.commit()
new_doc.save(new_pdf)
logger.info(f'save ok {new_pdf}')
final_json = json.dumps(obj, ensure_ascii=False, indent=2)
with open('equations_test/final_json.json', 'w') as f:
f.write(final_json)
return new_pdf
if __name__ == '__main__':
# draw_block_on_pdf_with_txt_replace_eq_bbox(new_json_path, equation_color_pdf)
pass
magic_pdf/pre_proc/fix_image.py.bak deleted 100644 → 0
View file @ 79b58a1e
import re
from magic_pdf.libs.boxbase import _is_in_or_part_overlap, _is_part_overlap, find_bottom_nearest_text_bbox, find_left_nearest_text_bbox, find_right_nearest_text_bbox, find_top_nearest_text_bbox
from magic_pdf.libs.textbase import get_text_block_base_info
def fix_image_vertical(image_bboxes:list, text_blocks:list):
"""
修正图片的位置
如果图片与文字block发生一定重叠（也就是图片切到了一部分文字），那么减少图片边缘，让文字和图片不再重叠。
只对垂直方向进行。
"""
for image_bbox in image_bboxes:
for text_block in text_blocks:
text_bbox = text_block["bbox"]
if _is_part_overlap(text_bbox, image_bbox) and any([text_bbox[0]>=image_bbox[0] and text_bbox[2]<=image_bbox[2], text_bbox[0]<=image_bbox[0] and text_bbox[2]>=image_bbox[2]]):
if text_bbox[1] < image_bbox[1]:#在图片上方
image_bbox[1] = text_bbox[3]+1
elif text_bbox[3]>image_bbox[3]:#在图片下方
image_bbox[3] = text_bbox[1]-1
return image_bboxes
def __merge_if_common_edge(bbox1, bbox2):
x_min_1, y_min_1, x_max_1, y_max_1 = bbox1
x_min_2, y_min_2, x_max_2, y_max_2 = bbox2
# 检查是否有公共的水平边
if y_min_1 == y_min_2 or y_max_1 == y_max_2:
# 确保一个框的x范围在另一个框的x范围内
if max(x_min_1, x_min_2) <= min(x_max_1, x_max_2):
return [min(x_min_1, x_min_2), min(y_min_1, y_min_2), max(x_max_1, x_max_2), max(y_max_1, y_max_2)]
# 检查是否有公共的垂直边
if x_min_1 == x_min_2 or x_max_1 == x_max_2:
# 确保一个框的y范围在另一个框的y范围内
if max(y_min_1, y_min_2) <= min(y_max_1, y_max_2):
return [min(x_min_1, x_min_2), min(y_min_1, y_min_2), max(x_max_1, x_max_2), max(y_max_1, y_max_2)]
# 如果没有公共边
return None
def fix_seperated_image(image_bboxes:list):
"""
如果2个图片有一个边重叠，那么合并2个图片
"""
new_images = []
droped_img_idx = []
for i in range(0, len(image_bboxes)):
for j in range(i+1, len(image_bboxes)):
new_img = __merge_if_common_edge(image_bboxes[i], image_bboxes[j])
if new_img is not None:
new_images.append(new_img)
droped_img_idx.append(i)
droped_img_idx.append(j)
break
for i in range(0, len(image_bboxes)):
if i not in droped_img_idx:
new_images.append(image_bboxes[i])
return new_images
def __check_img_title_pattern(text):
"""
检查文本段是否是表格的标题
"""
patterns = [r"^(fig|figure).*", r"^(scheme).*"]
text = text.strip()
for pattern in patterns:
match = re.match(pattern, text, re.IGNORECASE)
if match:
return True
return False
def __get_fig_caption_text(text_block):
txt = " ".join(span['text'] for line in text_block['lines'] for span in line['spans'])
line_cnt = len(text_block['lines'])
txt = txt.replace("Ž . ", '')
return txt, line_cnt
def __find_and_extend_bottom_caption(text_block, pymu_blocks, image_box):
"""
继续向下方寻找和图片caption字号，字体，颜色一样的文字框，合并入caption。
text_block是已经找到的图片catpion（这个caption可能不全，多行被划分到多个pymu block里了）
"""
combined_image_caption_text_block = list(text_block.copy()['bbox'])
base_font_color, base_font_size, base_font_type = get_text_block_base_info(text_block)
while True:
tb_add = find_bottom_nearest_text_bbox(pymu_blocks, combined_image_caption_text_block)
if not tb_add:
break
tb_font_color, tb_font_size, tb_font_type = get_text_block_base_info(tb_add)
if tb_font_color==base_font_color and tb_font_size==base_font_size and tb_font_type==base_font_type:
combined_image_caption_text_block[0] = min(combined_image_caption_text_block[0], tb_add['bbox'][0])
combined_image_caption_text_block[2] = max(combined_image_caption_text_block[2], tb_add['bbox'][2])
combined_image_caption_text_block[3] = tb_add['bbox'][3]
else:
break
image_box[0] = min(image_box[0], combined_image_caption_text_block[0])
image_box[1] = min(image_box[1], combined_image_caption_text_block[1])
image_box[2] = max(image_box[2], combined_image_caption_text_block[2])
image_box[3] = max(image_box[3], combined_image_caption_text_block[3])
text_block['_image_caption'] = True
def include_img_title(pymu_blocks, image_bboxes: list):
"""
向上方和下方寻找符合图片title的文本block，合并到图片里
如果图片上下都有fig的情况怎么办？寻找标题距离最近的那个。
---
增加对左侧和右侧图片标题的寻找
"""
for tb in image_bboxes:
# 优先找下方的
max_find_cnt = 3 # 向上，向下最多找3个就停止
temp_box = tb.copy()
while max_find_cnt>0:
text_block_btn = find_bottom_nearest_text_bbox(pymu_blocks, temp_box)
if text_block_btn:
txt, line_cnt = __get_fig_caption_text(text_block_btn)
if len(txt.strip())>0:
if not __check_img_title_pattern(txt) and max_find_cnt>0 and line_cnt<3: # 设置line_cnt<=2目的是为了跳过子标题，或者有时候图片下方文字没有被图片识别模型放入图片里
max_find_cnt = max_find_cnt - 1
temp_box[3] = text_block_btn['bbox'][3]
continue
else:
break
else:
temp_box[3] = text_block_btn['bbox'][3] # 宽度不变，扩大
max_find_cnt = max_find_cnt - 1
else:
break
max_find_cnt = 3 # 向上，向下最多找3个就停止
temp_box = tb.copy()
while max_find_cnt>0:
text_block_top = find_top_nearest_text_bbox(pymu_blocks, temp_box)
if text_block_top:
txt, line_cnt = __get_fig_caption_text(text_block_top)
if len(txt.strip())>0:
if not __check_img_title_pattern(txt) and max_find_cnt>0 and line_cnt <3:
max_find_cnt = max_find_cnt - 1
temp_box[1] = text_block_top['bbox'][1]
continue
else:
break
else:
b = text_block_top['bbox']
temp_box[1] = b[1] # 宽度不变，扩大
max_find_cnt = max_find_cnt - 1
else:
break
if text_block_btn and text_block_top and text_block_btn.get("_image_caption", False) is False and text_block_top.get("_image_caption", False) is False :
btn_text, _ = __get_fig_caption_text(text_block_btn)
top_text, _ = __get_fig_caption_text(text_block_top)
if __check_img_title_pattern(btn_text) and __check_img_title_pattern(top_text):
# 取距离图片最近的
btn_text_distance = text_block_btn['bbox'][1] - tb[3]
top_text_distance = tb[1] - text_block_top['bbox'][3]
if btn_text_distance<top_text_distance: # caption在下方
__find_and_extend_bottom_caption(text_block_btn, pymu_blocks, tb)
else:
text_block = text_block_top
tb[0] = min(tb[0], text_block['bbox'][0])
tb[1] = min(tb[1], text_block['bbox'][1])
tb[2] = max(tb[2], text_block['bbox'][2])
tb[3] = max(tb[3], text_block['bbox'][3])
text_block_btn['_image_caption'] = True
continue
text_block = text_block_btn # find_bottom_nearest_text_bbox(pymu_blocks, tb)
if text_block and text_block.get("_image_caption", False) is False:
first_text_line, _ = __get_fig_caption_text(text_block)
if __check_img_title_pattern(first_text_line):
# 发现特征之后，继续向相同方向寻找（想同颜色，想同大小，想同字体）的textblock
__find_and_extend_bottom_caption(text_block, pymu_blocks, tb)
continue
text_block = text_block_top # find_top_nearest_text_bbox(pymu_blocks, tb)
if text_block and text_block.get("_image_caption", False) is False:
first_text_line, _ = __get_fig_caption_text(text_block)
if __check_img_title_pattern(first_text_line):
tb[0] = min(tb[0], text_block['bbox'][0])
tb[1] = min(tb[1], text_block['bbox'][1])
tb[2] = max(tb[2], text_block['bbox'][2])
tb[3] = max(tb[3], text_block['bbox'][3])
text_block['_image_caption'] = True
continue
"""向左、向右寻找，暂时只寻找一次"""
left_text_block = find_left_nearest_text_bbox(pymu_blocks, tb)
if left_text_block and left_text_block.get("_image_caption", False) is False:
first_text_line, _ = __get_fig_caption_text(left_text_block)
if __check_img_title_pattern(first_text_line):
tb[0] = min(tb[0], left_text_block['bbox'][0])
tb[1] = min(tb[1], left_text_block['bbox'][1])
tb[2] = max(tb[2], left_text_block['bbox'][2])
tb[3] = max(tb[3], left_text_block['bbox'][3])
left_text_block['_image_caption'] = True
continue
right_text_block = find_right_nearest_text_bbox(pymu_blocks, tb)
if right_text_block and right_text_block.get("_image_caption", False) is False:
first_text_line, _ = __get_fig_caption_text(right_text_block)
if __check_img_title_pattern(first_text_line):
tb[0] = min(tb[0], right_text_block['bbox'][0])
tb[1] = min(tb[1], right_text_block['bbox'][1])
tb[2] = max(tb[2], right_text_block['bbox'][2])
tb[3] = max(tb[3], right_text_block['bbox'][3])
right_text_block['_image_caption'] = True
continue
return image_bboxes
def combine_images(image_bboxes:list):
"""
合并图片，如果图片有重叠，那么合并
"""
new_images = []
droped_img_idx = []
for i in range(0, len(image_bboxes)):
for j in range(i+1, len(image_bboxes)):
if j not in droped_img_idx and _is_in_or_part_overlap(image_bboxes[i], image_bboxes[j]):
# 合并
image_bboxes[i][0], image_bboxes[i][1],image_bboxes[i][2],image_bboxes[i][3] = min(image_bboxes[i][0], image_bboxes[j][0]), min(image_bboxes[i][1], image_bboxes[j][1]), max(image_bboxes[i][2], image_bboxes[j][2]), max(image_bboxes[i][3], image_bboxes[j][3])
droped_img_idx.append(j)
for i in range(0, len(image_bboxes)):
if i not in droped_img_idx:
new_images.append(image_bboxes[i])
return new_images
\ No newline at end of file
magic_pdf/pre_proc/fix_table.py.bak deleted 100644 → 0
View file @ 79b58a1e
from magic_pdf.libs.commons import fitz # pyMuPDF库
import re
from magic_pdf.libs.boxbase import _is_in_or_part_overlap, _is_part_overlap, find_bottom_nearest_text_bbox, find_left_nearest_text_bbox, find_right_nearest_text_bbox, find_top_nearest_text_bbox # json
## version 2
def get_merged_line(page):
"""
这个函数是为了从pymuPDF中提取出的矢量里筛出水平的横线，并且将断开的线段进行了合并。
:param page :fitz读取的当前页的内容
"""
drawings_bbox = []
drawings_line = []
drawings = page.get_drawings() # 提取所有的矢量
for p in drawings:
drawings_bbox.append(p["rect"].irect) # (L, U, R, D)
lines = []
for L, U, R, D in drawings_bbox:
if abs(D - U) <= 3: # 筛出水平的横线
lines.append((L, U, R, D))
U_groups = []
visited = [False for _ in range(len(lines))]
for i, (L1, U1, R1, D1) in enumerate(lines):
if visited[i] == True:
continue
tmp_g = [(L1, U1, R1, D1)]
for j, (L2, U2, R2, D2) in enumerate(lines):
if i == j:
continue
if visited[j] == True:
continue
if max(U1, D1, U2, D2) - min(U1, D1, U2, D2) <= 5: # 把高度一致的线放进一个group
tmp_g.append((L2, U2, R2, D2))
visited[j] = True
U_groups.append(tmp_g)
res = []
for group in U_groups:
group.sort(key = lambda LURD: (LURD[0], LURD[2]))
LL, UU, RR, DD = group[0]
for i, (L1, U1, R1, D1) in enumerate(group):
if (L1 - RR) >= 5:
cur_line = (LL, UU, RR, DD)
res.append(cur_line)
LL = L1
else:
RR = max(RR, R1)
cur_line = (LL, UU, RR, DD)
res.append(cur_line)
return res
def fix_tables(page: fitz.Page, table_bboxes: list, include_table_title: bool, scan_line_num: int):
"""
:param page :fitz读取的当前页的内容
:param table_bboxes: list类型，每一个元素是一个元祖 (L, U, R, D)
:param include_table_title: 是否将表格的标题也圈进来
:param scan_line_num: 在与表格框临近的上下几个文本框里扫描搜索标题
"""
drawings_lines = get_merged_line(page)
fix_table_bboxes = []
for table in table_bboxes:
(L, U, R, D) = table
fix_table_L = []
fix_table_U = []
fix_table_R = []
fix_table_D = []
width = R - L
width_range = width * 0.1 # 只看距离表格整体宽度10%之内偏差的线
height = D - U
height_range = height * 0.1 # 只看距离表格整体高度10%之内偏差的线
for line in drawings_lines:
if (L - width_range) <= line[0] <= (L + width_range) and (R - width_range) <= line[2] <= (R + width_range): # 相近的宽度
if (U - height_range) < line[1] < (U + height_range): # 上边界，在一定的高度范围内
fix_table_U.append(line[1])
fix_table_L.append(line[0])
fix_table_R.append(line[2])
elif (D - height_range) < line[1] < (D + height_range): # 下边界，在一定的高度范围内
fix_table_D.append(line[1])
fix_table_L.append(line[0])
fix_table_R.append(line[2])
if fix_table_U:
U = min(fix_table_U)
if fix_table_D:
D = max(fix_table_D)
if fix_table_L:
L = min(fix_table_L)
if fix_table_R:
R = max(fix_table_R)
if include_table_title: # 需要将表格标题包括
text_blocks = page.get_text("dict", flags=fitz.TEXTFLAGS_TEXT)["blocks"] # 所有的text的block
incolumn_text_blocks = [block for block in text_blocks if not ((block['bbox'][0] < L and block['bbox'][2] < L) or (block['bbox'][0] > R and block['bbox'][2] > R))] # 将与表格完全没有任何遮挡的文字筛除掉（比如另一栏的文字）
upper_text_blocks = [block for block in incolumn_text_blocks if (U - block['bbox'][3]) > 0] # 将在表格线以上的text block筛选出来
sorted_filtered_text_blocks = sorted(upper_text_blocks, key=lambda x: (U - x['bbox'][3], x['bbox'][0])) # 按照text block的下边界距离表格上边界的距离升序排序，如果是同一个高度，则先左再右
for idx in range(scan_line_num):
if idx+1 <= len(sorted_filtered_text_blocks):
line_temp = sorted_filtered_text_blocks[idx]['lines']
if line_temp:
text = line_temp[0]['spans'][0]['text'] # 提取出第一个span里的text内容
check_en = re.match('Table', text) # 检查是否有Table开头的(英文）
check_ch = re.match('表', text) # 检查是否有Table开头的(中文）
if check_en or check_ch:
if sorted_filtered_text_blocks[idx]['bbox'][1] < D: # 以防出现负的bbox
U = sorted_filtered_text_blocks[idx]['bbox'][1]
fix_table_bboxes.append([L-2, U-2, R+2, D+2])
return fix_table_bboxes
def __check_table_title_pattern(text):
"""
检查文本段是否是表格的标题
"""
patterns = [r'^table\s\d+']
for pattern in patterns:
match = re.match(pattern, text, re.IGNORECASE)
if match:
return True
else:
return False
def fix_table_text_block(pymu_blocks, table_bboxes: list):
"""
调整table, 如果table和上下的text block有相交区域，则将table的上下边界调整到text block的上下边界
例如 tmp/unittest/unittest_pdf/纯2列_ViLT_6_文字 表格.pdf
"""
for tb in table_bboxes:
(L, U, R, D) = tb
for block in pymu_blocks:
if _is_in_or_part_overlap((L, U, R, D), block['bbox']):
txt = " ".join(span['text'] for line in block['lines'] for span in line['spans'])
if not __check_table_title_pattern(txt) and block.get("_table", False) is False: # 如果是table的title，那么不调整。因为下一步会统一调整，如果这里进行了调整，后面的调整会造成调整到其他table的title上（在连续出现2个table的情况下）。
tb[0] = min(tb[0], block['bbox'][0])
tb[1] = min(tb[1], block['bbox'][1])
tb[2] = max(tb[2], block['bbox'][2])
tb[3] = max(tb[3], block['bbox'][3])
block['_table'] = True # 占位，防止其他table再次占用
"""如果是个table的title，但是有部分重叠，那么修正这个title,使得和table不重叠"""
if _is_part_overlap(tb, block['bbox']) and __check_table_title_pattern(txt):
block['bbox'] = list(block['bbox'])
if block['bbox'][3] > U:
block['bbox'][3] = U-1
if block['bbox'][1] < D:
block['bbox'][1] = D+1
return table_bboxes
def __get_table_caption_text(text_block):
txt = " ".join(span['text'] for line in text_block['lines'] for span in line['spans'])
line_cnt = len(text_block['lines'])
txt = txt.replace("Ž . ", '')
return txt, line_cnt
def include_table_title(pymu_blocks, table_bboxes: list):
"""
把表格的title也包含进来，扩展到table_bbox上
"""
for tb in table_bboxes:
max_find_cnt = 3 # 上上最多找3次
temp_box = tb.copy()
while max_find_cnt>0:
text_block_top = find_top_nearest_text_bbox(pymu_blocks, temp_box)
if text_block_top:
txt, line_cnt = __get_table_caption_text(text_block_top)
if len(txt.strip())>0:
if not __check_table_title_pattern(txt) and max_find_cnt>0 and line_cnt<3:
max_find_cnt = max_find_cnt -1
temp_box[1] = text_block_top['bbox'][1]
continue
else:
break
else:
temp_box[1] = text_block_top['bbox'][1] # 宽度不变，扩大
max_find_cnt = max_find_cnt - 1
else:
break
max_find_cnt = 3 # 向下找
temp_box = tb.copy()
while max_find_cnt>0:
text_block_bottom = find_bottom_nearest_text_bbox(pymu_blocks, temp_box)
if text_block_bottom:
txt, line_cnt = __get_table_caption_text(text_block_bottom)
if len(txt.strip())>0:
if not __check_table_title_pattern(txt) and max_find_cnt>0 and line_cnt<3:
max_find_cnt = max_find_cnt - 1
temp_box[3] = text_block_bottom['bbox'][3]
continue
else:
break
else:
temp_box[3] = text_block_bottom['bbox'][3]
max_find_cnt = max_find_cnt - 1
else:
break
if text_block_top and text_block_bottom and text_block_top.get("_table_caption", False) is False and text_block_bottom.get("_table_caption", False) is False :
btn_text, _ = __get_table_caption_text(text_block_bottom)
top_text, _ = __get_table_caption_text(text_block_top)
if __check_table_title_pattern(btn_text) and __check_table_title_pattern(top_text): # 上下都有一个tbale的caption
# 取距离最近的
btn_text_distance = text_block_bottom['bbox'][1] - tb[3]
top_text_distance = tb[1] - text_block_top['bbox'][3]
text_block = text_block_bottom if btn_text_distance<top_text_distance else text_block_top
tb[0] = min(tb[0], text_block['bbox'][0])
tb[1] = min(tb[1], text_block['bbox'][1])
tb[2] = max(tb[2], text_block['bbox'][2])
tb[3] = max(tb[3], text_block['bbox'][3])
text_block_bottom['_table_caption'] = True
continue
# 如果以上条件都不满足，那么就向下找
text_block = text_block_top
if text_block and text_block.get("_table_caption", False) is False:
first_text_line = " ".join(span['text'] for line in text_block['lines'] for span in line['spans'])
if __check_table_title_pattern(first_text_line) and text_block.get("_table", False) is False:
tb[0] = min(tb[0], text_block['bbox'][0])
tb[1] = min(tb[1], text_block['bbox'][1])
tb[2] = max(tb[2], text_block['bbox'][2])
tb[3] = max(tb[3], text_block['bbox'][3])
text_block['_table_caption'] = True
continue
text_block = text_block_bottom
if text_block and text_block.get("_table_caption", False) is False:
first_text_line, _ = __get_table_caption_text(text_block)
if __check_table_title_pattern(first_text_line) and text_block.get("_table", False) is False:
tb[0] = min(tb[0], text_block['bbox'][0])
tb[1] = min(tb[1], text_block['bbox'][1])
tb[2] = max(tb[2], text_block['bbox'][2])
tb[3] = max(tb[3], text_block['bbox'][3])
text_block['_table_caption'] = True
continue
"""向左、向右寻找，暂时只寻找一次"""
left_text_block = find_left_nearest_text_bbox(pymu_blocks, tb)
if left_text_block and left_text_block.get("_image_caption", False) is False:
first_text_line, _ = __get_table_caption_text(left_text_block)
if __check_table_title_pattern(first_text_line):
tb[0] = min(tb[0], left_text_block['bbox'][0])
tb[1] = min(tb[1], left_text_block['bbox'][1])
tb[2] = max(tb[2], left_text_block['bbox'][2])
tb[3] = max(tb[3], left_text_block['bbox'][3])
left_text_block['_image_caption'] = True
continue
right_text_block = find_right_nearest_text_bbox(pymu_blocks, tb)
if right_text_block and right_text_block.get("_image_caption", False) is False:
first_text_line, _ = __get_table_caption_text(right_text_block)
if __check_table_title_pattern(first_text_line):
tb[0] = min(tb[0], right_text_block['bbox'][0])
tb[1] = min(tb[1], right_text_block['bbox'][1])
tb[2] = max(tb[2], right_text_block['bbox'][2])
tb[3] = max(tb[3], right_text_block['bbox'][3])
right_text_block['_image_caption'] = True
continue
return table_bboxes
\ No newline at end of file
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