mirror of
https://github.com/opendatalab/MinerU.git
synced 2026-09-24 23:10:23 +08:00
Merge pull request #1099 from myhloli/dev
refactor(magic_pdf): remove unused functions and simplify code
This commit is contained in:
Xiaomeng Zhao
@@ -1,438 +0,0 @@
|
||||
import math
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from magic_pdf.config.ocr_content_type import ContentType
|
||||
from magic_pdf.libs.boxbase import (find_bottom_nearest_text_bbox,
|
||||
find_top_nearest_text_bbox)
|
||||
from magic_pdf.libs.commons import join_path
|
||||
|
||||
TYPE_INLINE_EQUATION = ContentType.InlineEquation
|
||||
TYPE_INTERLINE_EQUATION = ContentType.InterlineEquation
|
||||
UNI_FORMAT_TEXT_TYPE = ['text', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6']
|
||||
|
||||
|
||||
@DeprecationWarning
|
||||
def mk_nlp_markdown_1(para_dict: dict):
|
||||
"""对排序后的bboxes拼接内容."""
|
||||
content_lst = []
|
||||
for _, page_info in para_dict.items():
|
||||
para_blocks = page_info.get('para_blocks')
|
||||
if not para_blocks:
|
||||
continue
|
||||
|
||||
for block in para_blocks:
|
||||
item = block['paras']
|
||||
for _, p in item.items():
|
||||
para_text = p['para_text']
|
||||
is_title = p['is_para_title']
|
||||
title_level = p['para_title_level']
|
||||
md_title_prefix = '#' * title_level
|
||||
if is_title:
|
||||
content_lst.append(f'{md_title_prefix} {para_text}')
|
||||
else:
|
||||
content_lst.append(para_text)
|
||||
|
||||
content_text = '\n\n'.join(content_lst)
|
||||
|
||||
return content_text
|
||||
|
||||
|
||||
# 找到目标字符串在段落中的索引
|
||||
def __find_index(paragraph, target):
|
||||
index = paragraph.find(target)
|
||||
if index != -1:
|
||||
return index
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def __insert_string(paragraph, target, position):
|
||||
new_paragraph = paragraph[:position] + target + paragraph[position:]
|
||||
return new_paragraph
|
||||
|
||||
|
||||
def __insert_after(content, image_content, target):
|
||||
"""在content中找到target，将image_content插入到target后面."""
|
||||
index = content.find(target)
|
||||
if index != -1:
|
||||
content = (
|
||||
content[: index + len(target)]
|
||||
+ '\n\n'
|
||||
+ image_content
|
||||
+ '\n\n'
|
||||
+ content[index + len(target) :]
|
||||
)
|
||||
else:
|
||||
logger.error(
|
||||
f"Can't find the location of image {image_content} in the markdown file, search target is {target}"
|
||||
)
|
||||
return content
|
||||
|
||||
|
||||
def __insert_before(content, image_content, target):
|
||||
"""在content中找到target，将image_content插入到target前面."""
|
||||
index = content.find(target)
|
||||
if index != -1:
|
||||
content = content[:index] + '\n\n' + image_content + '\n\n' + content[index:]
|
||||
else:
|
||||
logger.error(
|
||||
f"Can't find the location of image {image_content} in the markdown file, search target is {target}"
|
||||
)
|
||||
return content
|
||||
|
||||
|
||||
@DeprecationWarning
|
||||
def mk_mm_markdown_1(para_dict: dict):
|
||||
"""拼装多模态markdown."""
|
||||
content_lst = []
|
||||
for _, page_info in para_dict.items():
|
||||
page_lst = [] # 一个page内的段落列表
|
||||
para_blocks = page_info.get('para_blocks')
|
||||
pymu_raw_blocks = page_info.get('preproc_blocks')
|
||||
|
||||
all_page_images = []
|
||||
all_page_images.extend(page_info.get('images', []))
|
||||
all_page_images.extend(page_info.get('image_backup', []))
|
||||
all_page_images.extend(page_info.get('tables', []))
|
||||
all_page_images.extend(page_info.get('table_backup', []))
|
||||
|
||||
if not para_blocks or not pymu_raw_blocks: # 只有图片的拼接的场景
|
||||
for img in all_page_images:
|
||||
page_lst.append(f"") # TODO 图片顺序
|
||||
page_md = '\n\n'.join(page_lst)
|
||||
|
||||
else:
|
||||
for block in para_blocks:
|
||||
item = block['paras']
|
||||
for _, p in item.items():
|
||||
para_text = p['para_text']
|
||||
is_title = p['is_para_title']
|
||||
title_level = p['para_title_level']
|
||||
md_title_prefix = '#' * title_level
|
||||
if is_title:
|
||||
page_lst.append(f'{md_title_prefix} {para_text}')
|
||||
else:
|
||||
page_lst.append(para_text)
|
||||
|
||||
"""拼装成一个页面的文本"""
|
||||
page_md = '\n\n'.join(page_lst)
|
||||
"""插入图片"""
|
||||
for img in all_page_images:
|
||||
imgbox = img['bbox']
|
||||
img_content = f""
|
||||
# 先看在哪个block内
|
||||
for block in pymu_raw_blocks:
|
||||
bbox = block['bbox']
|
||||
if (
|
||||
bbox[0] - 1 <= imgbox[0] < bbox[2] + 1
|
||||
and bbox[1] - 1 <= imgbox[1] < bbox[3] + 1
|
||||
): # 确定在block内
|
||||
for l in block['lines']: # noqa: E741
|
||||
line_box = l['bbox']
|
||||
if (
|
||||
line_box[0] - 1 <= imgbox[0] < line_box[2] + 1
|
||||
and line_box[1] - 1 <= imgbox[1] < line_box[3] + 1
|
||||
): # 在line内的，插入line前面
|
||||
line_txt = ''.join([s['text'] for s in l['spans']])
|
||||
page_md = __insert_before(
|
||||
page_md, img_content, line_txt
|
||||
)
|
||||
break
|
||||
break
|
||||
else: # 在行与行之间
|
||||
# 找到图片x0,y0与line的x0,y0最近的line
|
||||
min_distance = 100000
|
||||
min_line = None
|
||||
for l in block['lines']: # noqa: E741
|
||||
line_box = l['bbox']
|
||||
distance = math.sqrt(
|
||||
(line_box[0] - imgbox[0]) ** 2
|
||||
+ (line_box[1] - imgbox[1]) ** 2
|
||||
)
|
||||
if distance < min_distance:
|
||||
min_distance = distance
|
||||
min_line = l
|
||||
if min_line:
|
||||
line_txt = ''.join(
|
||||
[s['text'] for s in min_line['spans']]
|
||||
)
|
||||
img_h = imgbox[3] - imgbox[1]
|
||||
if min_distance < img_h: # 文字在图片前面
|
||||
page_md = __insert_after(
|
||||
page_md, img_content, line_txt
|
||||
)
|
||||
else:
|
||||
page_md = __insert_before(
|
||||
page_md, img_content, line_txt
|
||||
)
|
||||
else:
|
||||
logger.error(
|
||||
f"Can't find the location of image {img['image_path']} in the markdown file #1"
|
||||
)
|
||||
else: # 应当在两个block之间
|
||||
# 找到上方最近的block，如果上方没有就找大下方最近的block
|
||||
top_txt_block = find_top_nearest_text_bbox(pymu_raw_blocks, imgbox)
|
||||
if top_txt_block:
|
||||
line_txt = ''.join(
|
||||
[s['text'] for s in top_txt_block['lines'][-1]['spans']]
|
||||
)
|
||||
page_md = __insert_after(page_md, img_content, line_txt)
|
||||
else:
|
||||
bottom_txt_block = find_bottom_nearest_text_bbox(
|
||||
pymu_raw_blocks, imgbox
|
||||
)
|
||||
if bottom_txt_block:
|
||||
line_txt = ''.join(
|
||||
[
|
||||
s['text']
|
||||
for s in bottom_txt_block['lines'][0]['spans']
|
||||
]
|
||||
)
|
||||
page_md = __insert_before(page_md, img_content, line_txt)
|
||||
else:
|
||||
logger.error(
|
||||
f"Can't find the location of image {img['image_path']} in the markdown file #2"
|
||||
)
|
||||
|
||||
content_lst.append(page_md)
|
||||
|
||||
"""拼装成全部页面的文本"""
|
||||
content_text = '\n\n'.join(content_lst)
|
||||
|
||||
return content_text
|
||||
|
||||
|
||||
def __insert_after_para(text, type, element, content_list):
|
||||
"""在content_list中找到text，将image_path作为一个新的node插入到text后面."""
|
||||
for i, c in enumerate(content_list):
|
||||
content_type = c.get('type')
|
||||
if content_type in UNI_FORMAT_TEXT_TYPE and text in c.get('text', ''):
|
||||
if type == 'image':
|
||||
content_node = {
|
||||
'type': 'image',
|
||||
'img_path': element.get('image_path'),
|
||||
'img_alt': '',
|
||||
'img_title': '',
|
||||
'img_caption': '',
|
||||
}
|
||||
elif type == 'table':
|
||||
content_node = {
|
||||
'type': 'table',
|
||||
'img_path': element.get('image_path'),
|
||||
'table_latex': element.get('text'),
|
||||
'table_title': '',
|
||||
'table_caption': '',
|
||||
'table_quality': element.get('quality'),
|
||||
}
|
||||
content_list.insert(i + 1, content_node)
|
||||
break
|
||||
else:
|
||||
logger.error(
|
||||
f"Can't find the location of image {element.get('image_path')} in the markdown file, search target is {text}"
|
||||
)
|
||||
|
||||
|
||||
def __insert_before_para(text, type, element, content_list):
|
||||
"""在content_list中找到text，将image_path作为一个新的node插入到text前面."""
|
||||
for i, c in enumerate(content_list):
|
||||
content_type = c.get('type')
|
||||
if content_type in UNI_FORMAT_TEXT_TYPE and text in c.get('text', ''):
|
||||
if type == 'image':
|
||||
content_node = {
|
||||
'type': 'image',
|
||||
'img_path': element.get('image_path'),
|
||||
'img_alt': '',
|
||||
'img_title': '',
|
||||
'img_caption': '',
|
||||
}
|
||||
elif type == 'table':
|
||||
content_node = {
|
||||
'type': 'table',
|
||||
'img_path': element.get('image_path'),
|
||||
'table_latex': element.get('text'),
|
||||
'table_title': '',
|
||||
'table_caption': '',
|
||||
'table_quality': element.get('quality'),
|
||||
}
|
||||
content_list.insert(i, content_node)
|
||||
break
|
||||
else:
|
||||
logger.error(
|
||||
f"Can't find the location of image {element.get('image_path')} in the markdown file, search target is {text}"
|
||||
)
|
||||
|
||||
|
||||
def mk_universal_format(pdf_info_list: list, img_buket_path):
|
||||
"""构造统一格式 https://aicarrier.feishu.cn/wiki/FqmMwcH69iIdCWkkyjvcDwNUnTY."""
|
||||
content_lst = []
|
||||
for page_info in pdf_info_list:
|
||||
page_lst = [] # 一个page内的段落列表
|
||||
para_blocks = page_info.get('para_blocks')
|
||||
pymu_raw_blocks = page_info.get('preproc_blocks')
|
||||
|
||||
all_page_images = []
|
||||
all_page_images.extend(page_info.get('images', []))
|
||||
all_page_images.extend(page_info.get('image_backup', []))
|
||||
# all_page_images.extend(page_info.get("tables",[]))
|
||||
# all_page_images.extend(page_info.get("table_backup",[]) )
|
||||
all_page_tables = []
|
||||
all_page_tables.extend(page_info.get('tables', []))
|
||||
|
||||
if not para_blocks or not pymu_raw_blocks: # 只有图片的拼接的场景
|
||||
for img in all_page_images:
|
||||
content_node = {
|
||||
'type': 'image',
|
||||
'img_path': join_path(img_buket_path, img['image_path']),
|
||||
'img_alt': '',
|
||||
'img_title': '',
|
||||
'img_caption': '',
|
||||
}
|
||||
page_lst.append(content_node) # TODO 图片顺序
|
||||
for table in all_page_tables:
|
||||
content_node = {
|
||||
'type': 'table',
|
||||
'img_path': join_path(img_buket_path, table['image_path']),
|
||||
'table_latex': table.get('text'),
|
||||
'table_title': '',
|
||||
'table_caption': '',
|
||||
'table_quality': table.get('quality'),
|
||||
}
|
||||
page_lst.append(content_node) # TODO 图片顺序
|
||||
else:
|
||||
for block in para_blocks:
|
||||
item = block['paras']
|
||||
for _, p in item.items():
|
||||
font_type = p[
|
||||
'para_font_type'
|
||||
] # 对于文本来说，要么是普通文本，要么是个行间公式
|
||||
if font_type == TYPE_INTERLINE_EQUATION:
|
||||
content_node = {'type': 'equation', 'latex': p['para_text']}
|
||||
page_lst.append(content_node)
|
||||
else:
|
||||
para_text = p['para_text']
|
||||
is_title = p['is_para_title']
|
||||
title_level = p['para_title_level']
|
||||
|
||||
if is_title:
|
||||
content_node = {
|
||||
'type': f'h{title_level}',
|
||||
'text': para_text,
|
||||
}
|
||||
page_lst.append(content_node)
|
||||
else:
|
||||
content_node = {'type': 'text', 'text': para_text}
|
||||
page_lst.append(content_node)
|
||||
|
||||
content_lst.extend(page_lst)
|
||||
|
||||
"""插入图片"""
|
||||
for img in all_page_images:
|
||||
insert_img_or_table('image', img, pymu_raw_blocks, content_lst)
|
||||
|
||||
"""插入表格"""
|
||||
for table in all_page_tables:
|
||||
insert_img_or_table('table', table, pymu_raw_blocks, content_lst)
|
||||
# end for
|
||||
return content_lst
|
||||
|
||||
|
||||
def insert_img_or_table(type, element, pymu_raw_blocks, content_lst):
|
||||
element_bbox = element['bbox']
|
||||
# 先看在哪个block内
|
||||
for block in pymu_raw_blocks:
|
||||
bbox = block['bbox']
|
||||
if (
|
||||
bbox[0] - 1 <= element_bbox[0] < bbox[2] + 1
|
||||
and bbox[1] - 1 <= element_bbox[1] < bbox[3] + 1
|
||||
): # 确定在这个大的block内，然后进入逐行比较距离
|
||||
for l in block['lines']: # noqa: E741
|
||||
line_box = l['bbox']
|
||||
if (
|
||||
line_box[0] - 1 <= element_bbox[0] < line_box[2] + 1
|
||||
and line_box[1] - 1 <= element_bbox[1] < line_box[3] + 1
|
||||
): # 在line内的，插入line前面
|
||||
line_txt = ''.join([s['text'] for s in l['spans']])
|
||||
__insert_before_para(line_txt, type, element, content_lst)
|
||||
break
|
||||
break
|
||||
else: # 在行与行之间
|
||||
# 找到图片x0,y0与line的x0,y0最近的line
|
||||
min_distance = 100000
|
||||
min_line = None
|
||||
for l in block['lines']: # noqa: E741
|
||||
line_box = l['bbox']
|
||||
distance = math.sqrt(
|
||||
(line_box[0] - element_bbox[0]) ** 2
|
||||
+ (line_box[1] - element_bbox[1]) ** 2
|
||||
)
|
||||
if distance < min_distance:
|
||||
min_distance = distance
|
||||
min_line = l
|
||||
if min_line:
|
||||
line_txt = ''.join([s['text'] for s in min_line['spans']])
|
||||
img_h = element_bbox[3] - element_bbox[1]
|
||||
if min_distance < img_h: # 文字在图片前面
|
||||
__insert_after_para(line_txt, type, element, content_lst)
|
||||
else:
|
||||
__insert_before_para(line_txt, type, element, content_lst)
|
||||
break
|
||||
else:
|
||||
logger.error(
|
||||
f"Can't find the location of image {element.get('image_path')} in the markdown file #1"
|
||||
)
|
||||
else: # 应当在两个block之间
|
||||
# 找到上方最近的block，如果上方没有就找大下方最近的block
|
||||
top_txt_block = find_top_nearest_text_bbox(pymu_raw_blocks, element_bbox)
|
||||
if top_txt_block:
|
||||
line_txt = ''.join([s['text'] for s in top_txt_block['lines'][-1]['spans']])
|
||||
__insert_after_para(line_txt, type, element, content_lst)
|
||||
else:
|
||||
bottom_txt_block = find_bottom_nearest_text_bbox(
|
||||
pymu_raw_blocks, element_bbox
|
||||
)
|
||||
if bottom_txt_block:
|
||||
line_txt = ''.join(
|
||||
[s['text'] for s in bottom_txt_block['lines'][0]['spans']]
|
||||
)
|
||||
__insert_before_para(line_txt, type, element, content_lst)
|
||||
else: # TODO ，图片可能独占一列，这种情况上下是没有图片的
|
||||
logger.error(
|
||||
f"Can't find the location of image {element.get('image_path')} in the markdown file #2"
|
||||
)
|
||||
|
||||
|
||||
def mk_mm_markdown(content_list):
|
||||
"""基于同一格式的内容列表，构造markdown，含图片."""
|
||||
content_md = []
|
||||
for c in content_list:
|
||||
content_type = c.get('type')
|
||||
if content_type == 'text':
|
||||
content_md.append(c.get('text'))
|
||||
elif content_type == 'equation':
|
||||
content = c.get('latex')
|
||||
if content.startswith('$$') and content.endswith('$$'):
|
||||
content_md.append(content)
|
||||
else:
|
||||
content_md.append(f"\n$$\n{c.get('latex')}\n$$\n")
|
||||
elif content_type in UNI_FORMAT_TEXT_TYPE:
|
||||
content_md.append(f"{'#'*int(content_type[1])} {c.get('text')}")
|
||||
elif content_type == 'image':
|
||||
content_md.append(f"})")
|
||||
return '\n\n'.join(content_md)
|
||||
|
||||
|
||||
def mk_nlp_markdown(content_list):
|
||||
"""基于同一格式的内容列表，构造markdown，不含图片."""
|
||||
content_md = []
|
||||
for c in content_list:
|
||||
content_type = c.get('type')
|
||||
if content_type == 'text':
|
||||
content_md.append(c.get('text'))
|
||||
elif content_type == 'equation':
|
||||
content_md.append(f"$$\n{c.get('latex')}\n$$")
|
||||
elif content_type == 'table':
|
||||
content_md.append(f"$$$\n{c.get('table_latex')}\n$$$")
|
||||
elif content_type in UNI_FORMAT_TEXT_TYPE:
|
||||
content_md.append(f"{'#'*int(content_type[1])} {c.get('text')}")
|
||||
return '\n\n'.join(content_md)
|
||||
@@ -1,681 +0,0 @@
|
||||
# 定义这里的bbox是一个list [x0, y0, x1, y1, block_content, idx_x, idx_y, content_type, ext_x0, ext_y0, ext_x1, ext_y1], 初始时候idx_x, idx_y都是None
|
||||
# 其中x0, y0代表左上角坐标，x1, y1代表右下角坐标，坐标原点在左上角。
|
||||
|
||||
|
||||
|
||||
from magic_pdf.layout.layout_spiler_recog import get_spilter_of_page
|
||||
from magic_pdf.libs.boxbase import _is_in, _is_in_or_part_overlap, _is_vertical_full_overlap
|
||||
from magic_pdf.libs.commons import mymax
|
||||
|
||||
X0_IDX = 0
|
||||
Y0_IDX = 1
|
||||
X1_IDX = 2
|
||||
Y1_IDX = 3
|
||||
CONTENT_IDX = 4
|
||||
IDX_X = 5
|
||||
IDX_Y = 6
|
||||
CONTENT_TYPE_IDX = 7
|
||||
|
||||
X0_EXT_IDX = 8
|
||||
Y0_EXT_IDX = 9
|
||||
X1_EXT_IDX = 10
|
||||
Y1_EXT_IDX = 11
|
||||
|
||||
|
||||
def prepare_bboxes_for_layout_split(image_info, image_backup_info, table_info, inline_eq_info, interline_eq_info, text_raw_blocks: dict, page_boundry, page):
|
||||
"""
|
||||
text_raw_blocks:结构参考test/assets/papre/pymu_textblocks.json
|
||||
把bbox重新组装成一个list，每个元素[x0, y0, x1, y1, block_content, idx_x, idx_y, content_type, ext_x0, ext_y0, ext_x1, ext_y1], 初始时候idx_x, idx_y都是None. 对于图片、公式来说，block_content是图片的地址， 对于段落来说，block_content是pymupdf里的block结构
|
||||
"""
|
||||
all_bboxes = []
|
||||
|
||||
for image in image_info:
|
||||
box = image['bbox']
|
||||
# 由于没有实现横向的栏切分，因此在这里先过滤掉一些小的图片。这些图片有可能影响layout，造成没有横向栏切分的情况下，layout切分不准确。例如 scihub_76500000/libgen.scimag76570000-76570999.zip_10.1186/s13287-019-1355-1
|
||||
# 把长宽都小于50的去掉
|
||||
if abs(box[0]-box[2]) < 50 and abs(box[1]-box[3]) < 50:
|
||||
continue
|
||||
all_bboxes.append([box[0], box[1], box[2], box[3], None, None, None, 'image', None, None, None, None])
|
||||
|
||||
for table in table_info:
|
||||
box = table['bbox']
|
||||
all_bboxes.append([box[0], box[1], box[2], box[3], None, None, None, 'table', None, None, None, None])
|
||||
|
||||
"""由于公式与段落混合，因此公式不再参与layout划分，无需加入all_bboxes"""
|
||||
# 加入文本block
|
||||
text_block_temp = []
|
||||
for block in text_raw_blocks:
|
||||
bbox = block['bbox']
|
||||
text_block_temp.append([bbox[0], bbox[1], bbox[2], bbox[3], None, None, None, 'text', None, None, None, None])
|
||||
|
||||
text_block_new = resolve_bbox_overlap_for_layout_det(text_block_temp)
|
||||
text_block_new = filter_lines_bbox(text_block_new) # 去掉线条bbox，有可能让layout探测陷入无限循环
|
||||
|
||||
|
||||
"""找出会影响layout的色块、横向分割线"""
|
||||
spilter_bboxes = get_spilter_of_page(page, [b['bbox'] for b in image_info]+[b['bbox'] for b in image_backup_info], [b['bbox'] for b in table_info], )
|
||||
# 还要去掉存在于spilter_bboxes里的text_block
|
||||
if len(spilter_bboxes) > 0:
|
||||
text_block_new = [box for box in text_block_new if not any([_is_in_or_part_overlap(box[:4], spilter_bbox) for spilter_bbox in spilter_bboxes])]
|
||||
|
||||
for bbox in text_block_new:
|
||||
all_bboxes.append([bbox[0], bbox[1], bbox[2], bbox[3], None, None, None, 'text', None, None, None, None])
|
||||
|
||||
for bbox in spilter_bboxes:
|
||||
all_bboxes.append([bbox[0], bbox[1], bbox[2], bbox[3], None, None, None, 'spilter', None, None, None, None])
|
||||
|
||||
|
||||
return all_bboxes
|
||||
|
||||
def resolve_bbox_overlap_for_layout_det(bboxes:list):
|
||||
"""
|
||||
1. 去掉bbox互相包含的，去掉被包含的
|
||||
2. 上下方向上如果有重叠，就扩大大box范围，直到覆盖小box
|
||||
"""
|
||||
def _is_in_other_bbox(i:int):
|
||||
"""
|
||||
判断i个box是否被其他box有所包含
|
||||
"""
|
||||
for j in range(0, len(bboxes)):
|
||||
if j!=i and _is_in(bboxes[i][:4], bboxes[j][:4]):
|
||||
return True
|
||||
# elif j!=i and _is_bottom_full_overlap(bboxes[i][:4], bboxes[j][:4]):
|
||||
# return True
|
||||
|
||||
return False
|
||||
|
||||
# 首先去掉被包含的bbox
|
||||
new_bbox_1 = []
|
||||
for i in range(0, len(bboxes)):
|
||||
if not _is_in_other_bbox(i):
|
||||
new_bbox_1.append(bboxes[i])
|
||||
|
||||
# 其次扩展大的box
|
||||
new_box = []
|
||||
new_bbox_2 = []
|
||||
len_1 = len(new_bbox_2)
|
||||
while True:
|
||||
merged_idx = []
|
||||
for i in range(0, len(new_bbox_1)):
|
||||
if i in merged_idx:
|
||||
continue
|
||||
for j in range(i+1, len(new_bbox_1)):
|
||||
if j in merged_idx:
|
||||
continue
|
||||
bx1 = new_bbox_1[i]
|
||||
bx2 = new_bbox_1[j]
|
||||
if i!=j and _is_vertical_full_overlap(bx1[:4], bx2[:4]):
|
||||
merged_box = min([bx1[0], bx2[0]]), min([bx1[1], bx2[1]]), max([bx1[2], bx2[2]]), max([bx1[3], bx2[3]])
|
||||
new_bbox_2.append(merged_box)
|
||||
merged_idx.append(i)
|
||||
merged_idx.append(j)
|
||||
|
||||
for i in range(0, len(new_bbox_1)): # 没有合并的加入进来
|
||||
if i not in merged_idx:
|
||||
new_bbox_2.append(new_bbox_1[i])
|
||||
|
||||
if len(new_bbox_2)==0 or len_1==len(new_bbox_2):
|
||||
break
|
||||
else:
|
||||
len_1 = len(new_bbox_2)
|
||||
new_box = new_bbox_2
|
||||
new_bbox_1, new_bbox_2 = new_bbox_2, []
|
||||
|
||||
return new_box
|
||||
|
||||
|
||||
def filter_lines_bbox(bboxes: list):
|
||||
"""
|
||||
过滤掉bbox为空的行
|
||||
"""
|
||||
new_box = []
|
||||
for box in bboxes:
|
||||
x0, y0, x1, y1 = box[0], box[1], box[2], box[3]
|
||||
if abs(x0-x1)<=1 or abs(y0-y1)<=1:
|
||||
continue
|
||||
else:
|
||||
new_box.append(box)
|
||||
return new_box
|
||||
|
||||
|
||||
################################################################################
|
||||
# 第一种排序算法
|
||||
# 以下是基于延长线遮挡做的一个算法
|
||||
#
|
||||
################################################################################
|
||||
def find_all_left_bbox(this_bbox, all_bboxes) -> list:
|
||||
"""
|
||||
寻找this_bbox左边的所有bbox
|
||||
"""
|
||||
left_boxes = [box for box in all_bboxes if box[X1_IDX] <= this_bbox[X0_IDX]]
|
||||
return left_boxes
|
||||
|
||||
|
||||
def find_all_top_bbox(this_bbox, all_bboxes) -> list:
|
||||
"""
|
||||
寻找this_bbox上面的所有bbox
|
||||
"""
|
||||
top_boxes = [box for box in all_bboxes if box[Y1_IDX] <= this_bbox[Y0_IDX]]
|
||||
return top_boxes
|
||||
|
||||
|
||||
def get_and_set_idx_x(this_bbox, all_bboxes) -> int:
|
||||
"""
|
||||
寻找this_bbox在all_bboxes中的遮挡深度 idx_x
|
||||
"""
|
||||
if this_bbox[IDX_X] is not None:
|
||||
return this_bbox[IDX_X]
|
||||
else:
|
||||
all_left_bboxes = find_all_left_bbox(this_bbox, all_bboxes)
|
||||
if len(all_left_bboxes) == 0:
|
||||
this_bbox[IDX_X] = 0
|
||||
else:
|
||||
all_left_bboxes_idx = [get_and_set_idx_x(bbox, all_bboxes) for bbox in all_left_bboxes]
|
||||
max_idx_x = mymax(all_left_bboxes_idx)
|
||||
this_bbox[IDX_X] = max_idx_x + 1
|
||||
return this_bbox[IDX_X]
|
||||
|
||||
|
||||
def get_and_set_idx_y(this_bbox, all_bboxes) -> int:
|
||||
"""
|
||||
寻找this_bbox在all_bboxes中y方向的遮挡深度 idx_y
|
||||
"""
|
||||
if this_bbox[IDX_Y] is not None:
|
||||
return this_bbox[IDX_Y]
|
||||
else:
|
||||
all_top_bboxes = find_all_top_bbox(this_bbox, all_bboxes)
|
||||
if len(all_top_bboxes) == 0:
|
||||
this_bbox[IDX_Y] = 0
|
||||
else:
|
||||
all_top_bboxes_idx = [get_and_set_idx_y(bbox, all_bboxes) for bbox in all_top_bboxes]
|
||||
max_idx_y = mymax(all_top_bboxes_idx)
|
||||
this_bbox[IDX_Y] = max_idx_y + 1
|
||||
return this_bbox[IDX_Y]
|
||||
|
||||
|
||||
def bbox_sort(all_bboxes: list):
|
||||
"""
|
||||
排序
|
||||
"""
|
||||
all_bboxes_idx_x = [get_and_set_idx_x(bbox, all_bboxes) for bbox in all_bboxes]
|
||||
all_bboxes_idx_y = [get_and_set_idx_y(bbox, all_bboxes) for bbox in all_bboxes]
|
||||
all_bboxes_idx = [(idx_x, idx_y) for idx_x, idx_y in zip(all_bboxes_idx_x, all_bboxes_idx_y)]
|
||||
|
||||
all_bboxes_idx = [idx_x_y[0] * 100000 + idx_x_y[1] for idx_x_y in all_bboxes_idx] # 变换成一个点，保证能够先X，X相同时按Y排序
|
||||
all_bboxes_idx = list(zip(all_bboxes_idx, all_bboxes))
|
||||
all_bboxes_idx.sort(key=lambda x: x[0])
|
||||
sorted_bboxes = [bbox for idx, bbox in all_bboxes_idx]
|
||||
return sorted_bboxes
|
||||
|
||||
|
||||
################################################################################
|
||||
# 第二种排序算法
|
||||
# 下面的算法在计算idx_x和idx_y的时候不考虑延长线，而只考虑实际的长或者宽被遮挡的情况
|
||||
#
|
||||
################################################################################
|
||||
|
||||
def find_left_nearest_bbox(this_bbox, all_bboxes) -> list:
|
||||
"""
|
||||
在all_bboxes里找到所有右侧高度和this_bbox有重叠的bbox
|
||||
"""
|
||||
left_boxes = [box for box in all_bboxes if box[X1_IDX] <= this_bbox[X0_IDX] and any([
|
||||
box[Y0_IDX] < this_bbox[Y0_IDX] < box[Y1_IDX], box[Y0_IDX] < this_bbox[Y1_IDX] < box[Y1_IDX],
|
||||
this_bbox[Y0_IDX] < box[Y0_IDX] < this_bbox[Y1_IDX], this_bbox[Y0_IDX] < box[Y1_IDX] < this_bbox[Y1_IDX],
|
||||
box[Y0_IDX]==this_bbox[Y0_IDX] and box[Y1_IDX]==this_bbox[Y1_IDX]])]
|
||||
|
||||
# 然后再过滤一下，找到水平上距离this_bbox最近的那个
|
||||
if len(left_boxes) > 0:
|
||||
left_boxes.sort(key=lambda x: x[X1_IDX], reverse=True)
|
||||
left_boxes = [left_boxes[0]]
|
||||
else:
|
||||
left_boxes = []
|
||||
return left_boxes
|
||||
|
||||
|
||||
def get_and_set_idx_x_2(this_bbox, all_bboxes):
|
||||
"""
|
||||
寻找this_bbox在all_bboxes中的被直接遮挡的深度 idx_x
|
||||
这个遮挡深度不考虑延长线，而是被实际的长或者宽遮挡的情况
|
||||
"""
|
||||
if this_bbox[IDX_X] is not None:
|
||||
return this_bbox[IDX_X]
|
||||
else:
|
||||
left_nearest_bbox = find_left_nearest_bbox(this_bbox, all_bboxes)
|
||||
if len(left_nearest_bbox) == 0:
|
||||
this_bbox[IDX_X] = 0
|
||||
else:
|
||||
left_idx_x = get_and_set_idx_x_2(left_nearest_bbox[0], all_bboxes)
|
||||
this_bbox[IDX_X] = left_idx_x + 1
|
||||
return this_bbox[IDX_X]
|
||||
|
||||
|
||||
def find_top_nearest_bbox(this_bbox, all_bboxes) -> list:
|
||||
"""
|
||||
在all_bboxes里找到所有下侧宽度和this_bbox有重叠的bbox
|
||||
"""
|
||||
top_boxes = [box for box in all_bboxes if box[Y1_IDX] <= this_bbox[Y0_IDX] and any([
|
||||
box[X0_IDX] < this_bbox[X0_IDX] < box[X1_IDX], box[X0_IDX] < this_bbox[X1_IDX] < box[X1_IDX],
|
||||
this_bbox[X0_IDX] < box[X0_IDX] < this_bbox[X1_IDX], this_bbox[X0_IDX] < box[X1_IDX] < this_bbox[X1_IDX],
|
||||
box[X0_IDX]==this_bbox[X0_IDX] and box[X1_IDX]==this_bbox[X1_IDX]])]
|
||||
# 然后再过滤一下，找到水平上距离this_bbox最近的那个
|
||||
if len(top_boxes) > 0:
|
||||
top_boxes.sort(key=lambda x: x[Y1_IDX], reverse=True)
|
||||
top_boxes = [top_boxes[0]]
|
||||
else:
|
||||
top_boxes = []
|
||||
return top_boxes
|
||||
|
||||
|
||||
def get_and_set_idx_y_2(this_bbox, all_bboxes):
|
||||
"""
|
||||
寻找this_bbox在all_bboxes中的被直接遮挡的深度 idx_y
|
||||
这个遮挡深度不考虑延长线，而是被实际的长或者宽遮挡的情况
|
||||
"""
|
||||
if this_bbox[IDX_Y] is not None:
|
||||
return this_bbox[IDX_Y]
|
||||
else:
|
||||
top_nearest_bbox = find_top_nearest_bbox(this_bbox, all_bboxes)
|
||||
if len(top_nearest_bbox) == 0:
|
||||
this_bbox[IDX_Y] = 0
|
||||
else:
|
||||
top_idx_y = get_and_set_idx_y_2(top_nearest_bbox[0], all_bboxes)
|
||||
this_bbox[IDX_Y] = top_idx_y + 1
|
||||
return this_bbox[IDX_Y]
|
||||
|
||||
|
||||
def paper_bbox_sort(all_bboxes: list, page_width, page_height):
|
||||
all_bboxes_idx_x = [get_and_set_idx_x_2(bbox, all_bboxes) for bbox in all_bboxes]
|
||||
all_bboxes_idx_y = [get_and_set_idx_y_2(bbox, all_bboxes) for bbox in all_bboxes]
|
||||
all_bboxes_idx = [(idx_x, idx_y) for idx_x, idx_y in zip(all_bboxes_idx_x, all_bboxes_idx_y)]
|
||||
|
||||
all_bboxes_idx = [idx_x_y[0] * 100000 + idx_x_y[1] for idx_x_y in all_bboxes_idx] # 变换成一个点，保证能够先X，X相同时按Y排序
|
||||
all_bboxes_idx = list(zip(all_bboxes_idx, all_bboxes))
|
||||
all_bboxes_idx.sort(key=lambda x: x[0])
|
||||
sorted_bboxes = [bbox for idx, bbox in all_bboxes_idx]
|
||||
return sorted_bboxes
|
||||
|
||||
################################################################################
|
||||
"""
|
||||
第三种排序算法, 假设page的最左侧为X0，最右侧为X1，最上侧为Y0，最下侧为Y1
|
||||
这个排序算法在第二种算法基础上增加对bbox的预处理步骤。预处理思路如下：
|
||||
1. 首先在水平方向上对bbox进行扩展。扩展方法是：
|
||||
- 对每个bbox，找到其左边最近的bbox（也就是y方向有重叠），然后将其左边界扩展到左边最近bbox的右边界(x1+1),这里加1是为了避免重叠。如果没有左边的bbox，那么就将其左边界扩展到page的最左侧X0。
|
||||
- 对每个bbox，找到其右边最近的bbox（也就是y方向有重叠），然后将其右边界扩展到右边最近bbox的左边界(x0-1),这里减1是为了避免重叠。如果没有右边的bbox，那么就将其右边界扩展到page的最右侧X1。
|
||||
- 经过上面2个步骤，bbox扩展到了水平方向的最大范围。[左最近bbox.x1+1, 右最近bbox.x0-1]
|
||||
|
||||
2. 合并所有的连续水平方向的bbox, 合并方法是：
|
||||
- 对bbox进行y方向排序，然后从上到下遍历所有bbox，如果当前bbox和下一个bbox的x0, x1等于X0, X1，那么就合并这两个bbox。
|
||||
|
||||
3. 然后在垂直方向上对bbox进行扩展。扩展方法是：
|
||||
- 首先从page上切割掉合并后的水平bbox, 得到几个新的block
|
||||
针对每个block
|
||||
- x0: 扎到位于左侧x=x0延长线的左侧所有的bboxes, 找到最大的x1,让x0=x1+1。如果没有，则x0=X0
|
||||
- x1: 找到位于右侧x=x1延长线右侧所有的bboxes， 找到最小的x0, 让x1=x0-1。如果没有，则x1=X1
|
||||
随后在垂直方向上合并所有的连续的block，方法如下：
|
||||
- 对block进行x方向排序，然后从左到右遍历所有block，如果当前block和下一个block的x0, x1相等，那么就合并这两个block。
|
||||
如果垂直切分后所有小bbox都被分配到了一个block, 那么分割就完成了。这些合并后的block打上标签'GOOD_LAYOUT’
|
||||
如果在某个垂直方向上无法被完全分割到一个block，那么就将这个block打上标签'BAD_LAYOUT'。
|
||||
至此完成，一个页面的预处理，天然的block要么属于'GOOD_LAYOUT'，要么属于'BAD_LAYOUT'。针对含有'BAD_LAYOUT'的页面，可以先按照自上而下，自左到右进行天然排序，也可以先过滤掉这种书籍。
|
||||
(完成条件下次加强：进行水平方向切分，把混乱的layout部分尽可能切割出去)
|
||||
"""
|
||||
################################################################################
|
||||
def find_left_neighbor_bboxes(this_bbox, all_bboxes) -> list:
|
||||
"""
|
||||
在all_bboxes里找到所有右侧高度和this_bbox有重叠的bbox
|
||||
这里使用扩展之后的bbox
|
||||
"""
|
||||
left_boxes = [box for box in all_bboxes if box[X1_EXT_IDX] <= this_bbox[X0_EXT_IDX] and any([
|
||||
box[Y0_EXT_IDX] < this_bbox[Y0_EXT_IDX] < box[Y1_EXT_IDX], box[Y0_EXT_IDX] < this_bbox[Y1_EXT_IDX] < box[Y1_EXT_IDX],
|
||||
this_bbox[Y0_EXT_IDX] < box[Y0_EXT_IDX] < this_bbox[Y1_EXT_IDX], this_bbox[Y0_EXT_IDX] < box[Y1_EXT_IDX] < this_bbox[Y1_EXT_IDX],
|
||||
box[Y0_EXT_IDX]==this_bbox[Y0_EXT_IDX] and box[Y1_EXT_IDX]==this_bbox[Y1_EXT_IDX]])]
|
||||
|
||||
# 然后再过滤一下，找到水平上距离this_bbox最近的那个
|
||||
if len(left_boxes) > 0:
|
||||
left_boxes.sort(key=lambda x: x[X1_EXT_IDX], reverse=True)
|
||||
left_boxes = left_boxes
|
||||
else:
|
||||
left_boxes = []
|
||||
return left_boxes
|
||||
|
||||
def find_top_neighbor_bboxes(this_bbox, all_bboxes) -> list:
|
||||
"""
|
||||
在all_bboxes里找到所有下侧宽度和this_bbox有重叠的bbox
|
||||
这里使用扩展之后的bbox
|
||||
"""
|
||||
top_boxes = [box for box in all_bboxes if box[Y1_EXT_IDX] <= this_bbox[Y0_EXT_IDX] and any([
|
||||
box[X0_EXT_IDX] < this_bbox[X0_EXT_IDX] < box[X1_EXT_IDX], box[X0_EXT_IDX] < this_bbox[X1_EXT_IDX] < box[X1_EXT_IDX],
|
||||
this_bbox[X0_EXT_IDX] < box[X0_EXT_IDX] < this_bbox[X1_EXT_IDX], this_bbox[X0_EXT_IDX] < box[X1_EXT_IDX] < this_bbox[X1_EXT_IDX],
|
||||
box[X0_EXT_IDX]==this_bbox[X0_EXT_IDX] and box[X1_EXT_IDX]==this_bbox[X1_EXT_IDX]])]
|
||||
# 然后再过滤一下，找到水平上距离this_bbox最近的那个
|
||||
if len(top_boxes) > 0:
|
||||
top_boxes.sort(key=lambda x: x[Y1_EXT_IDX], reverse=True)
|
||||
top_boxes = top_boxes
|
||||
else:
|
||||
top_boxes = []
|
||||
return top_boxes
|
||||
|
||||
def get_and_set_idx_x_2_ext(this_bbox, all_bboxes):
|
||||
"""
|
||||
寻找this_bbox在all_bboxes中的被直接遮挡的深度 idx_x
|
||||
这个遮挡深度不考虑延长线，而是被实际的长或者宽遮挡的情况
|
||||
"""
|
||||
if this_bbox[IDX_X] is not None:
|
||||
return this_bbox[IDX_X]
|
||||
else:
|
||||
left_nearest_bbox = find_left_neighbor_bboxes(this_bbox, all_bboxes)
|
||||
if len(left_nearest_bbox) == 0:
|
||||
this_bbox[IDX_X] = 0
|
||||
else:
|
||||
left_idx_x = [get_and_set_idx_x_2(b, all_bboxes) for b in left_nearest_bbox]
|
||||
this_bbox[IDX_X] = mymax(left_idx_x) + 1
|
||||
return this_bbox[IDX_X]
|
||||
|
||||
def get_and_set_idx_y_2_ext(this_bbox, all_bboxes):
|
||||
"""
|
||||
寻找this_bbox在all_bboxes中的被直接遮挡的深度 idx_y
|
||||
这个遮挡深度不考虑延长线，而是被实际的长或者宽遮挡的情况
|
||||
"""
|
||||
if this_bbox[IDX_Y] is not None:
|
||||
return this_bbox[IDX_Y]
|
||||
else:
|
||||
top_nearest_bbox = find_top_neighbor_bboxes(this_bbox, all_bboxes)
|
||||
if len(top_nearest_bbox) == 0:
|
||||
this_bbox[IDX_Y] = 0
|
||||
else:
|
||||
top_idx_y = [get_and_set_idx_y_2_ext(b, all_bboxes) for b in top_nearest_bbox]
|
||||
this_bbox[IDX_Y] = mymax(top_idx_y) + 1
|
||||
return this_bbox[IDX_Y]
|
||||
|
||||
def _paper_bbox_sort_ext(all_bboxes: list):
|
||||
all_bboxes_idx_x = [get_and_set_idx_x_2_ext(bbox, all_bboxes) for bbox in all_bboxes]
|
||||
all_bboxes_idx_y = [get_and_set_idx_y_2_ext(bbox, all_bboxes) for bbox in all_bboxes]
|
||||
all_bboxes_idx = [(idx_x, idx_y) for idx_x, idx_y in zip(all_bboxes_idx_x, all_bboxes_idx_y)]
|
||||
|
||||
all_bboxes_idx = [idx_x_y[0] * 100000 + idx_x_y[1] for idx_x_y in all_bboxes_idx] # 变换成一个点，保证能够先X，X相同时按Y排序
|
||||
all_bboxes_idx = list(zip(all_bboxes_idx, all_bboxes))
|
||||
all_bboxes_idx.sort(key=lambda x: x[0])
|
||||
sorted_bboxes = [bbox for idx, bbox in all_bboxes_idx]
|
||||
return sorted_bboxes
|
||||
|
||||
# ===============================================================================================
|
||||
def find_left_bbox_ext_line(this_bbox, all_bboxes) -> list:
|
||||
"""
|
||||
寻找this_bbox左边的所有bbox, 使用延长线
|
||||
"""
|
||||
left_boxes = [box for box in all_bboxes if box[X1_IDX] <= this_bbox[X0_IDX]]
|
||||
if len(left_boxes):
|
||||
left_boxes.sort(key=lambda x: x[X1_IDX], reverse=True)
|
||||
left_boxes = left_boxes[0]
|
||||
else:
|
||||
left_boxes = None
|
||||
|
||||
return left_boxes
|
||||
|
||||
def find_right_bbox_ext_line(this_bbox, all_bboxes) -> list:
|
||||
"""
|
||||
寻找this_bbox右边的所有bbox, 使用延长线
|
||||
"""
|
||||
right_boxes = [box for box in all_bboxes if box[X0_IDX] >= this_bbox[X1_IDX]]
|
||||
if len(right_boxes):
|
||||
right_boxes.sort(key=lambda x: x[X0_IDX])
|
||||
right_boxes = right_boxes[0]
|
||||
else:
|
||||
right_boxes = None
|
||||
return right_boxes
|
||||
|
||||
# =============================================================================================
|
||||
|
||||
def find_left_nearest_bbox_direct(this_bbox, all_bboxes) -> list:
|
||||
"""
|
||||
在all_bboxes里找到所有右侧高度和this_bbox有重叠的bbox， 不用延长线并且不能像
|
||||
"""
|
||||
left_boxes = [box for box in all_bboxes if box[X1_IDX] <= this_bbox[X0_IDX] and any([
|
||||
box[Y0_IDX] < this_bbox[Y0_IDX] < box[Y1_IDX], box[Y0_IDX] < this_bbox[Y1_IDX] < box[Y1_IDX],
|
||||
this_bbox[Y0_IDX] < box[Y0_IDX] < this_bbox[Y1_IDX], this_bbox[Y0_IDX] < box[Y1_IDX] < this_bbox[Y1_IDX],
|
||||
box[Y0_IDX]==this_bbox[Y0_IDX] and box[Y1_IDX]==this_bbox[Y1_IDX]])]
|
||||
|
||||
# 然后再过滤一下，找到水平上距离this_bbox最近的那个——x1最大的那个
|
||||
if len(left_boxes) > 0:
|
||||
left_boxes.sort(key=lambda x: x[X1_EXT_IDX] if x[X1_EXT_IDX] else x[X1_IDX], reverse=True)
|
||||
left_boxes = left_boxes[0]
|
||||
else:
|
||||
left_boxes = None
|
||||
return left_boxes
|
||||
|
||||
def find_right_nearst_bbox_direct(this_bbox, all_bboxes) -> list:
|
||||
"""
|
||||
找到在this_bbox右侧且距离this_bbox距离最近的bbox.必须是直接遮挡的那种
|
||||
"""
|
||||
right_bboxes = [box for box in all_bboxes if box[X0_IDX] >= this_bbox[X1_IDX] and any([
|
||||
this_bbox[Y0_IDX] < box[Y0_IDX] < this_bbox[Y1_IDX], this_bbox[Y0_IDX] < box[Y1_IDX] < this_bbox[Y1_IDX],
|
||||
box[Y0_IDX] < this_bbox[Y0_IDX] < box[Y1_IDX], box[Y0_IDX] < this_bbox[Y1_IDX] < box[Y1_IDX],
|
||||
box[Y0_IDX]==this_bbox[Y0_IDX] and box[Y1_IDX]==this_bbox[Y1_IDX]])]
|
||||
|
||||
if len(right_bboxes)>0:
|
||||
right_bboxes.sort(key=lambda x: x[X0_EXT_IDX] if x[X0_EXT_IDX] else x[X0_IDX])
|
||||
right_bboxes = right_bboxes[0]
|
||||
else:
|
||||
right_bboxes = None
|
||||
return right_bboxes
|
||||
|
||||
def reset_idx_x_y(all_boxes:list)->list:
|
||||
for box in all_boxes:
|
||||
box[IDX_X] = None
|
||||
box[IDX_Y] = None
|
||||
|
||||
return all_boxes
|
||||
|
||||
# ===================================================================================================
|
||||
def find_top_nearest_bbox_direct(this_bbox, bboxes_collection) -> list:
|
||||
"""
|
||||
找到在this_bbox上方且距离this_bbox距离最近的bbox.必须是直接遮挡的那种
|
||||
"""
|
||||
top_bboxes = [box for box in bboxes_collection if box[Y1_IDX] <= this_bbox[Y0_IDX] and any([
|
||||
box[X0_IDX] < this_bbox[X0_IDX] < box[X1_IDX], box[X0_IDX] < this_bbox[X1_IDX] < box[X1_IDX],
|
||||
this_bbox[X0_IDX] < box[X0_IDX] < this_bbox[X1_IDX], this_bbox[X0_IDX] < box[X1_IDX] < this_bbox[X1_IDX],
|
||||
box[X0_IDX]==this_bbox[X0_IDX] and box[X1_IDX]==this_bbox[X1_IDX]])]
|
||||
# 然后再过滤一下，找到上方距离this_bbox最近的那个
|
||||
if len(top_bboxes) > 0:
|
||||
top_bboxes.sort(key=lambda x: x[Y1_IDX], reverse=True)
|
||||
top_bboxes = top_bboxes[0]
|
||||
else:
|
||||
top_bboxes = None
|
||||
return top_bboxes
|
||||
|
||||
def find_bottom_nearest_bbox_direct(this_bbox, bboxes_collection) -> list:
|
||||
"""
|
||||
找到在this_bbox下方且距离this_bbox距离最近的bbox.必须是直接遮挡的那种
|
||||
"""
|
||||
bottom_bboxes = [box for box in bboxes_collection if box[Y0_IDX] >= this_bbox[Y1_IDX] and any([
|
||||
box[X0_IDX] < this_bbox[X0_IDX] < box[X1_IDX], box[X0_IDX] < this_bbox[X1_IDX] < box[X1_IDX],
|
||||
this_bbox[X0_IDX] < box[X0_IDX] < this_bbox[X1_IDX], this_bbox[X0_IDX] < box[X1_IDX] < this_bbox[X1_IDX],
|
||||
box[X0_IDX]==this_bbox[X0_IDX] and box[X1_IDX]==this_bbox[X1_IDX]])]
|
||||
# 然后再过滤一下，找到水平上距离this_bbox最近的那个
|
||||
if len(bottom_bboxes) > 0:
|
||||
bottom_bboxes.sort(key=lambda x: x[Y0_IDX])
|
||||
bottom_bboxes = bottom_bboxes[0]
|
||||
else:
|
||||
bottom_bboxes = None
|
||||
return bottom_bboxes
|
||||
|
||||
def find_boundry_bboxes(bboxes:list) -> tuple:
|
||||
"""
|
||||
找到bboxes的边界——找到所有bbox里最小的(x0, y0), 最大的(x1, y1)
|
||||
"""
|
||||
x0, y0, x1, y1 = bboxes[0][X0_IDX], bboxes[0][Y0_IDX], bboxes[0][X1_IDX], bboxes[0][Y1_IDX]
|
||||
for box in bboxes:
|
||||
x0 = min(box[X0_IDX], x0)
|
||||
y0 = min(box[Y0_IDX], y0)
|
||||
x1 = max(box[X1_IDX], x1)
|
||||
y1 = max(box[Y1_IDX], y1)
|
||||
|
||||
return x0, y0, x1, y1
|
||||
|
||||
|
||||
def extend_bbox_vertical(bboxes:list, boundry_x0, boundry_y0, boundry_x1, boundry_y1) -> list:
|
||||
"""
|
||||
在垂直方向上扩展能够直接垂直打通的bbox,也就是那些上下都没有其他box的bbox
|
||||
"""
|
||||
for box in bboxes:
|
||||
top_nearest_bbox = find_top_nearest_bbox_direct(box, bboxes)
|
||||
bottom_nearest_bbox = find_bottom_nearest_bbox_direct(box, bboxes)
|
||||
if top_nearest_bbox is None and bottom_nearest_bbox is None: # 独占一列
|
||||
box[X0_EXT_IDX] = box[X0_IDX]
|
||||
box[Y0_EXT_IDX] = boundry_y0
|
||||
box[X1_EXT_IDX] = box[X1_IDX]
|
||||
box[Y1_EXT_IDX] = boundry_y1
|
||||
# else:
|
||||
# if top_nearest_bbox is None:
|
||||
# box[Y0_EXT_IDX] = boundry_y0
|
||||
# else:
|
||||
# box[Y0_EXT_IDX] = top_nearest_bbox[Y1_IDX] + 1
|
||||
# if bottom_nearest_bbox is None:
|
||||
# box[Y1_EXT_IDX] = boundry_y1
|
||||
# else:
|
||||
# box[Y1_EXT_IDX] = bottom_nearest_bbox[Y0_IDX] - 1
|
||||
# box[X0_EXT_IDX] = box[X0_IDX]
|
||||
# box[X1_EXT_IDX] = box[X1_IDX]
|
||||
return bboxes
|
||||
|
||||
|
||||
# ===================================================================================================
|
||||
|
||||
def paper_bbox_sort_v2(all_bboxes: list, page_width:int, page_height:int):
|
||||
"""
|
||||
增加预处理行为的排序:
|
||||
return:
|
||||
[
|
||||
{
|
||||
"layout_bbox": [x0, y0, x1, y1],
|
||||
"layout_label":"GOOD_LAYOUT/BAD_LAYOUT",
|
||||
"content_bboxes": [] #每个元素都是[x0, y0, x1, y1, block_content, idx_x, idx_y, content_type, ext_x0, ext_y0, ext_x1, ext_y1], 并且顺序就是阅读顺序
|
||||
}
|
||||
]
|
||||
"""
|
||||
sorted_layouts = [] # 最后的返回结果
|
||||
page_x0, page_y0, page_x1, page_y1 = 1, 1, page_width-1, page_height-1
|
||||
|
||||
all_bboxes = paper_bbox_sort(all_bboxes) # 大致拍下序
|
||||
# 首先在水平方向上扩展独占一行的bbox
|
||||
for bbox in all_bboxes:
|
||||
left_nearest_bbox = find_left_nearest_bbox_direct(bbox, all_bboxes) # 非扩展线
|
||||
right_nearest_bbox = find_right_nearst_bbox_direct(bbox, all_bboxes)
|
||||
if left_nearest_bbox is None and right_nearest_bbox is None: # 独占一行
|
||||
bbox[X0_EXT_IDX] = page_x0
|
||||
bbox[Y0_EXT_IDX] = bbox[Y0_IDX]
|
||||
bbox[X1_EXT_IDX] = page_x1
|
||||
bbox[Y1_EXT_IDX] = bbox[Y1_IDX]
|
||||
|
||||
# 此时独占一行的被成功扩展到指定的边界上，这个时候利用边界条件合并连续的bbox，成为一个group
|
||||
if len(all_bboxes)==1:
|
||||
return [{"layout_bbox": [page_x0, page_y0, page_x1, page_y1], "layout_label":"GOOD_LAYOUT", "content_bboxes": all_bboxes}]
|
||||
if len(all_bboxes)==0:
|
||||
return []
|
||||
|
||||
"""
|
||||
然后合并所有连续水平方向的bbox.
|
||||
|
||||
"""
|
||||
all_bboxes.sort(key=lambda x: x[Y0_IDX])
|
||||
h_bboxes = []
|
||||
h_bbox_group = []
|
||||
v_boxes = []
|
||||
|
||||
for bbox in all_bboxes:
|
||||
if bbox[X0_IDX] == page_x0 and bbox[X1_IDX] == page_x1:
|
||||
h_bbox_group.append(bbox)
|
||||
else:
|
||||
if len(h_bbox_group)>0:
|
||||
h_bboxes.append(h_bbox_group)
|
||||
h_bbox_group = []
|
||||
# 最后一个group
|
||||
if len(h_bbox_group)>0:
|
||||
h_bboxes.append(h_bbox_group)
|
||||
|
||||
"""
|
||||
现在h_bboxes里面是所有的group了，每个group都是一个list
|
||||
对h_bboxes里的每个group进行计算放回到sorted_layouts里
|
||||
"""
|
||||
for gp in h_bboxes:
|
||||
gp.sort(key=lambda x: x[Y0_IDX])
|
||||
block_info = {"layout_label":"GOOD_LAYOUT", "content_bboxes": gp}
|
||||
# 然后计算这个group的layout_bbox，也就是最小的x0,y0, 最大的x1,y1
|
||||
x0, y0, x1, y1 = gp[0][X0_EXT_IDX], gp[0][Y0_EXT_IDX], gp[-1][X1_EXT_IDX], gp[-1][Y1_EXT_IDX]
|
||||
block_info["layout_bbox"] = [x0, y0, x1, y1]
|
||||
sorted_layouts.append(block_info)
|
||||
|
||||
# 接下来利用这些连续的水平bbox的layout_bbox的y0, y1，从水平上切分开其余的为几个部分
|
||||
h_split_lines = [page_y0]
|
||||
for gp in h_bboxes:
|
||||
layout_bbox = gp['layout_bbox']
|
||||
y0, y1 = layout_bbox[1], layout_bbox[3]
|
||||
h_split_lines.append(y0)
|
||||
h_split_lines.append(y1)
|
||||
h_split_lines.append(page_y1)
|
||||
|
||||
unsplited_bboxes = []
|
||||
for i in range(0, len(h_split_lines), 2):
|
||||
start_y0, start_y1 = h_split_lines[i:i+2]
|
||||
# 然后找出[start_y0, start_y1]之间的其他bbox，这些组成一个未分割板块
|
||||
bboxes_in_block = [bbox for bbox in all_bboxes if bbox[Y0_IDX]>=start_y0 and bbox[Y1_IDX]<=start_y1]
|
||||
unsplited_bboxes.append(bboxes_in_block)
|
||||
# ================== 至此，水平方向的 已经切分排序完毕====================================
|
||||
"""
|
||||
接下来针对每个非水平的部分切分垂直方向的
|
||||
此时，只剩下了无法被完全水平打通的bbox了。对这些box，优先进行垂直扩展，然后进行垂直切分.
|
||||
分3步：
|
||||
1. 先把能完全垂直打通的隔离出去当做一个layout
|
||||
2. 其余的先垂直切分
|
||||
3. 垂直切分之后的部分再尝试水平切分
|
||||
4. 剩下的不能被切分的各个部分当成一个layout
|
||||
"""
|
||||
# 对每部分进行垂直切分
|
||||
for bboxes_in_block in unsplited_bboxes:
|
||||
# 首先对这个block的bbox进行垂直方向上的扩展
|
||||
boundry_x0, boundry_y0, boundry_x1, boundry_y1 = find_boundry_bboxes(bboxes_in_block)
|
||||
# 进行垂直方向上的扩展
|
||||
extended_vertical_bboxes = extend_bbox_vertical(bboxes_in_block, boundry_x0, boundry_y0, boundry_x1, boundry_y1)
|
||||
# 然后对这个block进行垂直方向上的切分
|
||||
extend_bbox_vertical.sort(key=lambda x: x[X0_IDX]) # x方向上从小到大，代表了从左到右读取
|
||||
v_boxes_group = []
|
||||
for bbox in extended_vertical_bboxes:
|
||||
if bbox[Y0_IDX]==boundry_y0 and bbox[Y1_IDX]==boundry_y1:
|
||||
v_boxes_group.append(bbox)
|
||||
else:
|
||||
if len(v_boxes_group)>0:
|
||||
v_boxes.append(v_boxes_group)
|
||||
v_boxes_group = []
|
||||
|
||||
if len(v_boxes_group)>0:
|
||||
|
||||
v_boxes.append(v_boxes_group)
|
||||
|
||||
# 把连续的垂直部分加入到sorted_layouts里。注意这个时候已经是连续的垂直部分了，因为上面已经做了
|
||||
for gp in v_boxes:
|
||||
gp.sort(key=lambda x: x[X0_IDX])
|
||||
block_info = {"layout_label":"GOOD_LAYOUT", "content_bboxes": gp}
|
||||
# 然后计算这个group的layout_bbox，也就是最小的x0,y0, 最大的x1,y1
|
||||
x0, y0, x1, y1 = gp[0][X0_EXT_IDX], gp[0][Y0_EXT_IDX], gp[-1][X1_EXT_IDX], gp[-1][Y1_EXT_IDX]
|
||||
block_info["layout_bbox"] = [x0, y0, x1, y1]
|
||||
sorted_layouts.append(block_info)
|
||||
|
||||
# 在垂直方向上，划分子块，也就是用贯通的垂直线进行切分。这些被切分出来的块，极大可能是可被垂直切分的，如果不能完全的垂直切分，那么尝试水平切分。都不能的则当成一个layout
|
||||
v_split_lines = [boundry_x0]
|
||||
for gp in v_boxes:
|
||||
layout_bbox = gp['layout_bbox']
|
||||
x0, x1 = layout_bbox[0], layout_bbox[2]
|
||||
v_split_lines.append(x0)
|
||||
v_split_lines.append(x1)
|
||||
v_split_lines.append(boundry_x1)
|
||||
|
||||
reset_idx_x_y(all_bboxes)
|
||||
all_boxes = _paper_bbox_sort_ext(all_bboxes)
|
||||
return all_boxes
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -1,182 +0,0 @@
|
||||
from magic_pdf.layout.bbox_sort import X0_EXT_IDX, X0_IDX, X1_EXT_IDX, X1_IDX, Y0_IDX, Y1_EXT_IDX, Y1_IDX
|
||||
from magic_pdf.libs.boxbase import _is_bottom_full_overlap, _left_intersect, _right_intersect
|
||||
|
||||
|
||||
def find_all_left_bbox_direct(this_bbox, all_bboxes) -> list:
|
||||
"""
|
||||
在all_bboxes里找到所有右侧垂直方向上和this_bbox有重叠的bbox， 不用延长线
|
||||
并且要考虑两个box左右相交的情况，如果相交了，那么右侧的box就不算最左侧。
|
||||
"""
|
||||
left_boxes = [box for box in all_bboxes if box[X1_IDX] <= this_bbox[X0_IDX]
|
||||
and any([
|
||||
box[Y0_IDX] < this_bbox[Y0_IDX] < box[Y1_IDX], box[Y0_IDX] < this_bbox[Y1_IDX] < box[Y1_IDX],
|
||||
this_bbox[Y0_IDX] < box[Y0_IDX] < this_bbox[Y1_IDX], this_bbox[Y0_IDX] < box[Y1_IDX] < this_bbox[Y1_IDX],
|
||||
box[Y0_IDX]==this_bbox[Y0_IDX] and box[Y1_IDX]==this_bbox[Y1_IDX]]) or _left_intersect(box[:4], this_bbox[:4])]
|
||||
|
||||
# 然后再过滤一下，找到水平上距离this_bbox最近的那个——x1最大的那个
|
||||
if len(left_boxes) > 0:
|
||||
left_boxes.sort(key=lambda x: x[X1_EXT_IDX] if x[X1_EXT_IDX] else x[X1_IDX], reverse=True)
|
||||
left_boxes = left_boxes[0]
|
||||
else:
|
||||
left_boxes = None
|
||||
return left_boxes
|
||||
|
||||
def find_all_right_bbox_direct(this_bbox, all_bboxes) -> list:
|
||||
"""
|
||||
找到在this_bbox右侧且距离this_bbox距离最近的bbox.必须是直接遮挡的那种
|
||||
"""
|
||||
right_bboxes = [box for box in all_bboxes if box[X0_IDX] >= this_bbox[X1_IDX]
|
||||
and any([
|
||||
this_bbox[Y0_IDX] < box[Y0_IDX] < this_bbox[Y1_IDX], this_bbox[Y0_IDX] < box[Y1_IDX] < this_bbox[Y1_IDX],
|
||||
box[Y0_IDX] < this_bbox[Y0_IDX] < box[Y1_IDX], box[Y0_IDX] < this_bbox[Y1_IDX] < box[Y1_IDX],
|
||||
box[Y0_IDX]==this_bbox[Y0_IDX] and box[Y1_IDX]==this_bbox[Y1_IDX]]) or _right_intersect(this_bbox[:4], box[:4])]
|
||||
|
||||
if len(right_bboxes)>0:
|
||||
right_bboxes.sort(key=lambda x: x[X0_EXT_IDX] if x[X0_EXT_IDX] else x[X0_IDX])
|
||||
right_bboxes = right_bboxes[0]
|
||||
else:
|
||||
right_bboxes = None
|
||||
return right_bboxes
|
||||
|
||||
def find_all_top_bbox_direct(this_bbox, all_bboxes) -> list:
|
||||
"""
|
||||
找到在this_bbox上侧且距离this_bbox距离最近的bbox.必须是直接遮挡的那种
|
||||
"""
|
||||
top_bboxes = [box for box in all_bboxes if box[Y1_IDX] <= this_bbox[Y0_IDX] and any([
|
||||
box[X0_IDX] < this_bbox[X0_IDX] < box[X1_IDX], box[X0_IDX] < this_bbox[X1_IDX] < box[X1_IDX],
|
||||
this_bbox[X0_IDX] < box[X0_IDX] < this_bbox[X1_IDX], this_bbox[X0_IDX] < box[X1_IDX] < this_bbox[X1_IDX],
|
||||
box[X0_IDX]==this_bbox[X0_IDX] and box[X1_IDX]==this_bbox[X1_IDX]])]
|
||||
|
||||
if len(top_bboxes)>0:
|
||||
top_bboxes.sort(key=lambda x: x[Y1_EXT_IDX] if x[Y1_EXT_IDX] else x[Y1_IDX], reverse=True)
|
||||
top_bboxes = top_bboxes[0]
|
||||
else:
|
||||
top_bboxes = None
|
||||
return top_bboxes
|
||||
|
||||
def find_all_bottom_bbox_direct(this_bbox, all_bboxes) -> list:
|
||||
"""
|
||||
找到在this_bbox下侧且距离this_bbox距离最近的bbox.必须是直接遮挡的那种
|
||||
"""
|
||||
bottom_bboxes = [box for box in all_bboxes if box[Y0_IDX] >= this_bbox[Y1_IDX] and any([
|
||||
this_bbox[X0_IDX] < box[X0_IDX] < this_bbox[X1_IDX], this_bbox[X0_IDX] < box[X1_IDX] < this_bbox[X1_IDX],
|
||||
box[X0_IDX] < this_bbox[X0_IDX] < box[X1_IDX], box[X0_IDX] < this_bbox[X1_IDX] < box[X1_IDX],
|
||||
box[X0_IDX]==this_bbox[X0_IDX] and box[X1_IDX]==this_bbox[X1_IDX]])]
|
||||
|
||||
if len(bottom_bboxes)>0:
|
||||
bottom_bboxes.sort(key=lambda x: x[Y0_IDX])
|
||||
bottom_bboxes = bottom_bboxes[0]
|
||||
else:
|
||||
bottom_bboxes = None
|
||||
return bottom_bboxes
|
||||
|
||||
# ===================================================================================================================
|
||||
def find_bottom_bbox_direct_from_right_edge(this_bbox, all_bboxes) -> list:
|
||||
"""
|
||||
找到在this_bbox下侧且距离this_bbox距离最近的bbox.必须是直接遮挡的那种
|
||||
"""
|
||||
bottom_bboxes = [box for box in all_bboxes if box[Y0_IDX] >= this_bbox[Y1_IDX] and any([
|
||||
this_bbox[X0_IDX] < box[X0_IDX] < this_bbox[X1_IDX], this_bbox[X0_IDX] < box[X1_IDX] < this_bbox[X1_IDX],
|
||||
box[X0_IDX] < this_bbox[X0_IDX] < box[X1_IDX], box[X0_IDX] < this_bbox[X1_IDX] < box[X1_IDX],
|
||||
box[X0_IDX]==this_bbox[X0_IDX] and box[X1_IDX]==this_bbox[X1_IDX]])]
|
||||
|
||||
if len(bottom_bboxes)>0:
|
||||
# y0最小， X1最大的那个,也就是box上边缘最靠近this_bbox的那个,并且还最靠右
|
||||
bottom_bboxes.sort(key=lambda x: x[Y0_IDX])
|
||||
bottom_bboxes = [box for box in bottom_bboxes if box[Y0_IDX]==bottom_bboxes[0][Y0_IDX]]
|
||||
# 然后再y1相同的情况下，找到x1最大的那个
|
||||
bottom_bboxes.sort(key=lambda x: x[X1_IDX], reverse=True)
|
||||
bottom_bboxes = bottom_bboxes[0]
|
||||
else:
|
||||
bottom_bboxes = None
|
||||
return bottom_bboxes
|
||||
|
||||
def find_bottom_bbox_direct_from_left_edge(this_bbox, all_bboxes) -> list:
|
||||
"""
|
||||
找到在this_bbox下侧且距离this_bbox距离最近的bbox.必须是直接遮挡的那种
|
||||
"""
|
||||
bottom_bboxes = [box for box in all_bboxes if box[Y0_IDX] >= this_bbox[Y1_IDX] and any([
|
||||
this_bbox[X0_IDX] < box[X0_IDX] < this_bbox[X1_IDX], this_bbox[X0_IDX] < box[X1_IDX] < this_bbox[X1_IDX],
|
||||
box[X0_IDX] < this_bbox[X0_IDX] < box[X1_IDX], box[X0_IDX] < this_bbox[X1_IDX] < box[X1_IDX],
|
||||
box[X0_IDX]==this_bbox[X0_IDX] and box[X1_IDX]==this_bbox[X1_IDX]])]
|
||||
|
||||
if len(bottom_bboxes)>0:
|
||||
# y0最小， X0最小的那个
|
||||
bottom_bboxes.sort(key=lambda x: x[Y0_IDX])
|
||||
bottom_bboxes = [box for box in bottom_bboxes if box[Y0_IDX]==bottom_bboxes[0][Y0_IDX]]
|
||||
# 然后再y0相同的情况下，找到x0最小的那个
|
||||
bottom_bboxes.sort(key=lambda x: x[X0_IDX])
|
||||
bottom_bboxes = bottom_bboxes[0]
|
||||
else:
|
||||
bottom_bboxes = None
|
||||
return bottom_bboxes
|
||||
|
||||
def find_top_bbox_direct_from_left_edge(this_bbox, all_bboxes) -> list:
|
||||
"""
|
||||
找到在this_bbox上侧且距离this_bbox距离最近的bbox.必须是直接遮挡的那种
|
||||
"""
|
||||
top_bboxes = [box for box in all_bboxes if box[Y1_IDX] <= this_bbox[Y0_IDX] and any([
|
||||
box[X0_IDX] < this_bbox[X0_IDX] < box[X1_IDX], box[X0_IDX] < this_bbox[X1_IDX] < box[X1_IDX],
|
||||
this_bbox[X0_IDX] < box[X0_IDX] < this_bbox[X1_IDX], this_bbox[X0_IDX] < box[X1_IDX] < this_bbox[X1_IDX],
|
||||
box[X0_IDX]==this_bbox[X0_IDX] and box[X1_IDX]==this_bbox[X1_IDX]])]
|
||||
|
||||
if len(top_bboxes)>0:
|
||||
# y1最大， X0最小的那个
|
||||
top_bboxes.sort(key=lambda x: x[Y1_IDX], reverse=True)
|
||||
top_bboxes = [box for box in top_bboxes if box[Y1_IDX]==top_bboxes[0][Y1_IDX]]
|
||||
# 然后再y1相同的情况下，找到x0最小的那个
|
||||
top_bboxes.sort(key=lambda x: x[X0_IDX])
|
||||
top_bboxes = top_bboxes[0]
|
||||
else:
|
||||
top_bboxes = None
|
||||
return top_bboxes
|
||||
|
||||
def find_top_bbox_direct_from_right_edge(this_bbox, all_bboxes) -> list:
|
||||
"""
|
||||
找到在this_bbox上侧且距离this_bbox距离最近的bbox.必须是直接遮挡的那种
|
||||
"""
|
||||
top_bboxes = [box for box in all_bboxes if box[Y1_IDX] <= this_bbox[Y0_IDX] and any([
|
||||
box[X0_IDX] < this_bbox[X0_IDX] < box[X1_IDX], box[X0_IDX] < this_bbox[X1_IDX] < box[X1_IDX],
|
||||
this_bbox[X0_IDX] < box[X0_IDX] < this_bbox[X1_IDX], this_bbox[X0_IDX] < box[X1_IDX] < this_bbox[X1_IDX],
|
||||
box[X0_IDX]==this_bbox[X0_IDX] and box[X1_IDX]==this_bbox[X1_IDX]])]
|
||||
|
||||
if len(top_bboxes)>0:
|
||||
# y1最大， X1最大的那个
|
||||
top_bboxes.sort(key=lambda x: x[Y1_IDX], reverse=True)
|
||||
top_bboxes = [box for box in top_bboxes if box[Y1_IDX]==top_bboxes[0][Y1_IDX]]
|
||||
# 然后再y1相同的情况下，找到x1最大的那个
|
||||
top_bboxes.sort(key=lambda x: x[X1_IDX], reverse=True)
|
||||
top_bboxes = top_bboxes[0]
|
||||
else:
|
||||
top_bboxes = None
|
||||
return top_bboxes
|
||||
|
||||
# ===================================================================================================================
|
||||
|
||||
def get_left_edge_bboxes(all_bboxes) -> list:
|
||||
"""
|
||||
返回最左边的bbox
|
||||
"""
|
||||
left_bboxes = [box for box in all_bboxes if find_all_left_bbox_direct(box, all_bboxes) is None]
|
||||
return left_bboxes
|
||||
|
||||
def get_right_edge_bboxes(all_bboxes) -> list:
|
||||
"""
|
||||
返回最右边的bbox
|
||||
"""
|
||||
right_bboxes = [box for box in all_bboxes if find_all_right_bbox_direct(box, all_bboxes) is None]
|
||||
return right_bboxes
|
||||
|
||||
def fix_vertical_bbox_pos(bboxes:list):
|
||||
"""
|
||||
检查这批bbox在垂直方向是否有轻微的重叠，如果重叠了，就把重叠的bbox往下移动一点
|
||||
在x方向上必须一个包含或者被包含，或者完全重叠，不能只有部分重叠
|
||||
"""
|
||||
bboxes.sort(key=lambda x: x[Y0_IDX]) # 从上向下排列
|
||||
for i in range(0, len(bboxes)):
|
||||
for j in range(i+1, len(bboxes)):
|
||||
if _is_bottom_full_overlap(bboxes[i][:4], bboxes[j][:4]):
|
||||
# 如果两个bbox有部分重叠，那么就把下面的bbox往下移动一点
|
||||
bboxes[j][Y0_IDX] = bboxes[i][Y1_IDX] + 2 # 2是个经验值
|
||||
break
|
||||
return bboxes
|
||||
@@ -1,921 +0,0 @@
|
||||
"""对pdf上的box进行layout识别，并对内部组成的box进行排序."""
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from magic_pdf.layout.bbox_sort import (CONTENT_IDX, CONTENT_TYPE_IDX,
|
||||
X0_EXT_IDX, X0_IDX, X1_EXT_IDX, X1_IDX,
|
||||
Y0_EXT_IDX, Y0_IDX, Y1_EXT_IDX, Y1_IDX,
|
||||
paper_bbox_sort)
|
||||
from magic_pdf.layout.layout_det_utils import (
|
||||
find_all_bottom_bbox_direct, find_all_left_bbox_direct,
|
||||
find_all_right_bbox_direct, find_all_top_bbox_direct,
|
||||
find_bottom_bbox_direct_from_left_edge,
|
||||
find_bottom_bbox_direct_from_right_edge,
|
||||
find_top_bbox_direct_from_left_edge, find_top_bbox_direct_from_right_edge,
|
||||
get_left_edge_bboxes, get_right_edge_bboxes)
|
||||
from magic_pdf.libs.boxbase import get_bbox_in_boundary
|
||||
|
||||
LAYOUT_V = 'V'
|
||||
LAYOUT_H = 'H'
|
||||
LAYOUT_UNPROC = 'U'
|
||||
LAYOUT_BAD = 'B'
|
||||
|
||||
|
||||
def _is_single_line_text(bbox):
|
||||
"""检查bbox里面的文字是否只有一行."""
|
||||
return True # TODO
|
||||
box_type = bbox[CONTENT_TYPE_IDX]
|
||||
if box_type != 'text':
|
||||
return False
|
||||
paras = bbox[CONTENT_IDX]['paras']
|
||||
text_content = ''
|
||||
for para_id, para in paras.items(): # 拼装内部的段落文本
|
||||
is_title = para['is_title']
|
||||
if is_title != 0:
|
||||
text_content += f"## {para['text']}"
|
||||
else:
|
||||
text_content += para['text']
|
||||
text_content += '\n\n'
|
||||
|
||||
return bbox[CONTENT_TYPE_IDX] == 'text' and len(text_content.split('\n\n')) <= 1
|
||||
|
||||
|
||||
def _horizontal_split(bboxes: list, boundary: tuple, avg_font_size=20) -> list:
|
||||
"""
|
||||
对bboxes进行水平切割
|
||||
方法是：找到左侧和右侧都没有被直接遮挡的box，然后进行扩展，之后进行切割
|
||||
return:
|
||||
返回几个大的Layout区域 [[x0, y0, x1, y1, "h|u|v"], ], h代表水平，u代表未探测的，v代表垂直布局
|
||||
"""
|
||||
sorted_layout_blocks = [] # 这是要最终返回的值
|
||||
|
||||
bound_x0, bound_y0, bound_x1, bound_y1 = boundary
|
||||
all_bboxes = get_bbox_in_boundary(bboxes, boundary)
|
||||
# all_bboxes = paper_bbox_sort(all_bboxes, abs(bound_x1-bound_x0), abs(bound_y1-bound_x0)) # 大致拍下序, 这个是基于直接遮挡的。
|
||||
"""
|
||||
首先在水平方向上扩展独占一行的bbox
|
||||
|
||||
"""
|
||||
last_h_split_line_y1 = bound_y0 # 记录下上次的水平分割线
|
||||
for i, bbox in enumerate(all_bboxes):
|
||||
left_nearest_bbox = find_all_left_bbox_direct(bbox, all_bboxes) # 非扩展线
|
||||
right_nearest_bbox = find_all_right_bbox_direct(bbox, all_bboxes)
|
||||
if left_nearest_bbox is None and right_nearest_bbox is None: # 独占一行
|
||||
"""
|
||||
然而，如果只是孤立的一行文字，那么就还要满足以下几个条件才可以：
|
||||
1. bbox和中心线相交。或者
|
||||
2. 上方或者下方也存在同类水平的独占一行的bbox。 或者
|
||||
3. TODO 加强条件：这个bbox上方和下方是同一列column，那么就不能算作独占一行
|
||||
"""
|
||||
# 先检查这个bbox里是否只包含一行文字
|
||||
# is_single_line = _is_single_line_text(bbox)
|
||||
"""
|
||||
这里有个点需要注意，当页面内容不是居中的时候，第一次调用传递的是page的boundary，这个时候mid_x就不是中心线了.
|
||||
所以这里计算出最紧致的boundary，然后再计算mid_x
|
||||
"""
|
||||
boundary_real_x0, boundary_real_x1 = min(
|
||||
[bbox[X0_IDX] for bbox in all_bboxes]
|
||||
), max([bbox[X1_IDX] for bbox in all_bboxes])
|
||||
mid_x = (boundary_real_x0 + boundary_real_x1) / 2
|
||||
# 检查这个box是否内容在中心线有交
|
||||
# 必须跨过去2个字符的宽度
|
||||
is_cross_boundary_mid_line = (
|
||||
min(mid_x - bbox[X0_IDX], bbox[X1_IDX] - mid_x) > avg_font_size * 2
|
||||
)
|
||||
"""
|
||||
检查条件2
|
||||
"""
|
||||
is_belong_to_col = False
|
||||
"""
|
||||
检查是否能被上方col吸收，方法是：
|
||||
1. 上方非空且不是独占一行的，并且
|
||||
2. 从上个水平分割的最大y=y1开始到当前bbox,最左侧的bbox的[min_x0, max_x1],能够覆盖当前box的[x0, x1]
|
||||
"""
|
||||
"""
|
||||
以迭代的方式向上找，查找范围是[bound_x0, last_h_sp, bound_x1, bbox[Y0_IDX]]
|
||||
"""
|
||||
# 先确定上方的y0, y0
|
||||
b_y0, b_y1 = last_h_split_line_y1, bbox[Y0_IDX]
|
||||
# 然后从box开始逐个向上找到所有与box在x上有交集的box
|
||||
box_to_check = [bound_x0, b_y0, bound_x1, b_y1]
|
||||
bbox_in_bound_check = get_bbox_in_boundary(all_bboxes, box_to_check)
|
||||
|
||||
bboxes_on_top = []
|
||||
virtual_box = bbox
|
||||
while True:
|
||||
b_on_top = find_all_top_bbox_direct(virtual_box, bbox_in_bound_check)
|
||||
if b_on_top is not None:
|
||||
bboxes_on_top.append(b_on_top)
|
||||
virtual_box = [
|
||||
min([virtual_box[X0_IDX], b_on_top[X0_IDX]]),
|
||||
min(virtual_box[Y0_IDX], b_on_top[Y0_IDX]),
|
||||
max([virtual_box[X1_IDX], b_on_top[X1_IDX]]),
|
||||
b_y1,
|
||||
]
|
||||
else:
|
||||
break
|
||||
|
||||
# 随后确定这些box的最小x0, 最大x1
|
||||
if len(bboxes_on_top) > 0 and len(bboxes_on_top) != len(
|
||||
bbox_in_bound_check
|
||||
): # virtual_box可能会膨胀到占满整个区域，这实际上就不能属于一个col了。
|
||||
min_x0, max_x1 = virtual_box[X0_IDX], virtual_box[X1_IDX]
|
||||
# 然后采用一种比较粗糙的方法，看min_x0，max_x1是否与位于[bound_x0, last_h_sp, bound_x1, bbox[Y0_IDX]]之间的box有相交
|
||||
|
||||
if not any(
|
||||
[
|
||||
b[X0_IDX] <= min_x0 - 1 <= b[X1_IDX]
|
||||
or b[X0_IDX] <= max_x1 + 1 <= b[X1_IDX]
|
||||
for b in bbox_in_bound_check
|
||||
]
|
||||
):
|
||||
# 其上，下都不能被扩展成行，暂时只检查一下上方 TODO
|
||||
top_nearest_bbox = find_all_top_bbox_direct(bbox, bboxes)
|
||||
bottom_nearest_bbox = find_all_bottom_bbox_direct(bbox, bboxes)
|
||||
if not any(
|
||||
[
|
||||
top_nearest_bbox is not None
|
||||
and (
|
||||
find_all_left_bbox_direct(top_nearest_bbox, bboxes)
|
||||
is None
|
||||
and find_all_right_bbox_direct(top_nearest_bbox, bboxes)
|
||||
is None
|
||||
),
|
||||
bottom_nearest_bbox is not None
|
||||
and (
|
||||
find_all_left_bbox_direct(bottom_nearest_bbox, bboxes)
|
||||
is None
|
||||
and find_all_right_bbox_direct(
|
||||
bottom_nearest_bbox, bboxes
|
||||
)
|
||||
is None
|
||||
),
|
||||
top_nearest_bbox is None or bottom_nearest_bbox is None,
|
||||
]
|
||||
):
|
||||
is_belong_to_col = True
|
||||
|
||||
# 检查是否能被下方col吸收 TODO
|
||||
"""
|
||||
这里为什么没有is_cross_boundary_mid_line的条件呢？
|
||||
确实有些杂志左右两栏宽度不是对称的。
|
||||
"""
|
||||
if not is_belong_to_col or is_cross_boundary_mid_line:
|
||||
bbox[X0_EXT_IDX] = bound_x0
|
||||
bbox[Y0_EXT_IDX] = bbox[Y0_IDX]
|
||||
bbox[X1_EXT_IDX] = bound_x1
|
||||
bbox[Y1_EXT_IDX] = bbox[Y1_IDX]
|
||||
last_h_split_line_y1 = bbox[Y1_IDX] # 更新这条线
|
||||
else:
|
||||
continue
|
||||
"""
|
||||
此时独占一行的被成功扩展到指定的边界上，这个时候利用边界条件合并连续的bbox，成为一个group
|
||||
然后合并所有连续水平方向的bbox.
|
||||
"""
|
||||
all_bboxes.sort(key=lambda x: x[Y0_IDX])
|
||||
h_bboxes = []
|
||||
h_bbox_group = []
|
||||
|
||||
for bbox in all_bboxes:
|
||||
if bbox[X0_EXT_IDX] == bound_x0 and bbox[X1_EXT_IDX] == bound_x1:
|
||||
h_bbox_group.append(bbox)
|
||||
else:
|
||||
if len(h_bbox_group) > 0:
|
||||
h_bboxes.append(h_bbox_group)
|
||||
h_bbox_group = []
|
||||
# 最后一个group
|
||||
if len(h_bbox_group) > 0:
|
||||
h_bboxes.append(h_bbox_group)
|
||||
"""
|
||||
现在h_bboxes里面是所有的group了，每个group都是一个list
|
||||
对h_bboxes里的每个group进行计算放回到sorted_layouts里
|
||||
"""
|
||||
h_layouts = []
|
||||
for gp in h_bboxes:
|
||||
gp.sort(key=lambda x: x[Y0_IDX])
|
||||
# 然后计算这个group的layout_bbox，也就是最小的x0,y0, 最大的x1,y1
|
||||
x0, y0, x1, y1 = (
|
||||
gp[0][X0_EXT_IDX],
|
||||
gp[0][Y0_EXT_IDX],
|
||||
gp[-1][X1_EXT_IDX],
|
||||
gp[-1][Y1_EXT_IDX],
|
||||
)
|
||||
h_layouts.append([x0, y0, x1, y1, LAYOUT_H]) # 水平的布局
|
||||
"""
|
||||
接下来利用这些连续的水平bbox的layout_bbox的y0, y1，从水平上切分开其余的为几个部分
|
||||
"""
|
||||
h_split_lines = [bound_y0]
|
||||
for gp in h_bboxes: # gp是一个list[bbox_list]
|
||||
y0, y1 = gp[0][1], gp[-1][3]
|
||||
h_split_lines.append(y0)
|
||||
h_split_lines.append(y1)
|
||||
h_split_lines.append(bound_y1)
|
||||
|
||||
unsplited_bboxes = []
|
||||
for i in range(0, len(h_split_lines), 2):
|
||||
start_y0, start_y1 = h_split_lines[i : i + 2]
|
||||
# 然后找出[start_y0, start_y1]之间的其他bbox，这些组成一个未分割板块
|
||||
bboxes_in_block = [
|
||||
bbox
|
||||
for bbox in all_bboxes
|
||||
if bbox[Y0_IDX] >= start_y0 and bbox[Y1_IDX] <= start_y1
|
||||
]
|
||||
unsplited_bboxes.append(bboxes_in_block)
|
||||
# 接着把未处理的加入到h_layouts里
|
||||
for bboxes_in_block in unsplited_bboxes:
|
||||
if len(bboxes_in_block) == 0:
|
||||
continue
|
||||
x0, y0, x1, y1 = (
|
||||
bound_x0,
|
||||
min([bbox[Y0_IDX] for bbox in bboxes_in_block]),
|
||||
bound_x1,
|
||||
max([bbox[Y1_IDX] for bbox in bboxes_in_block]),
|
||||
)
|
||||
h_layouts.append([x0, y0, x1, y1, LAYOUT_UNPROC])
|
||||
|
||||
h_layouts.sort(key=lambda x: x[1]) # 按照y0排序, 也就是从上到下的顺序
|
||||
"""
|
||||
转换成如下格式返回
|
||||
"""
|
||||
for layout in h_layouts:
|
||||
sorted_layout_blocks.append(
|
||||
{
|
||||
'layout_bbox': layout[:4],
|
||||
'layout_label': layout[4],
|
||||
'sub_layout': [],
|
||||
}
|
||||
)
|
||||
return sorted_layout_blocks
|
||||
|
||||
|
||||
###############################################################################################
|
||||
#
|
||||
# 垂直方向的处理
|
||||
#
|
||||
#
|
||||
###############################################################################################
|
||||
def _vertical_align_split_v1(bboxes: list, boundary: tuple) -> list:
|
||||
"""
|
||||
计算垂直方向上的对齐， 并分割bboxes成layout。负责对一列多行的进行列维度分割。
|
||||
如果不能完全分割，剩余部分作为layout_lable为u的layout返回
|
||||
-----------------------
|
||||
| | |
|
||||
| | |
|
||||
| | |
|
||||
| | |
|
||||
-------------------------
|
||||
此函数会将：以上布局将会切分出来2列
|
||||
"""
|
||||
sorted_layout_blocks = [] # 这是要最终返回的值
|
||||
new_boundary = [boundary[0], boundary[1], boundary[2], boundary[3]]
|
||||
|
||||
v_blocks = []
|
||||
"""
|
||||
先从左到右切分
|
||||
"""
|
||||
while True:
|
||||
all_bboxes = get_bbox_in_boundary(bboxes, new_boundary)
|
||||
left_edge_bboxes = get_left_edge_bboxes(all_bboxes)
|
||||
if len(left_edge_bboxes) == 0:
|
||||
break
|
||||
right_split_line_x1 = max([bbox[X1_IDX] for bbox in left_edge_bboxes]) + 1
|
||||
# 然后检查这条线能不与其他bbox的左边界相交或者重合
|
||||
if any(
|
||||
[bbox[X0_IDX] <= right_split_line_x1 <= bbox[X1_IDX] for bbox in all_bboxes]
|
||||
):
|
||||
# 垂直切分线与某些box发生相交，说明无法完全垂直方向切分。
|
||||
break
|
||||
else: # 说明成功分割出一列
|
||||
# 找到左侧边界最靠左的bbox作为layout的x0
|
||||
layout_x0 = min(
|
||||
[bbox[X0_IDX] for bbox in left_edge_bboxes]
|
||||
) # 这里主要是为了画出来有一定间距
|
||||
v_blocks.append(
|
||||
[
|
||||
layout_x0,
|
||||
new_boundary[1],
|
||||
right_split_line_x1,
|
||||
new_boundary[3],
|
||||
LAYOUT_V,
|
||||
]
|
||||
)
|
||||
new_boundary[0] = right_split_line_x1 # 更新边界
|
||||
"""
|
||||
再从右到左切， 此时如果还是无法完全切分，那么剩余部分作为layout_lable为u的layout返回
|
||||
"""
|
||||
unsplited_block = []
|
||||
while True:
|
||||
all_bboxes = get_bbox_in_boundary(bboxes, new_boundary)
|
||||
right_edge_bboxes = get_right_edge_bboxes(all_bboxes)
|
||||
if len(right_edge_bboxes) == 0:
|
||||
break
|
||||
left_split_line_x0 = min([bbox[X0_IDX] for bbox in right_edge_bboxes]) - 1
|
||||
# 然后检查这条线能不与其他bbox的左边界相交或者重合
|
||||
if any(
|
||||
[bbox[X0_IDX] <= left_split_line_x0 <= bbox[X1_IDX] for bbox in all_bboxes]
|
||||
):
|
||||
# 这里是余下的
|
||||
unsplited_block.append(
|
||||
[
|
||||
new_boundary[0],
|
||||
new_boundary[1],
|
||||
new_boundary[2],
|
||||
new_boundary[3],
|
||||
LAYOUT_UNPROC,
|
||||
]
|
||||
)
|
||||
break
|
||||
else:
|
||||
# 找到右侧边界最靠右的bbox作为layout的x1
|
||||
layout_x1 = max([bbox[X1_IDX] for bbox in right_edge_bboxes])
|
||||
v_blocks.append(
|
||||
[
|
||||
left_split_line_x0,
|
||||
new_boundary[1],
|
||||
layout_x1,
|
||||
new_boundary[3],
|
||||
LAYOUT_V,
|
||||
]
|
||||
)
|
||||
new_boundary[2] = left_split_line_x0 # 更新右边界
|
||||
"""
|
||||
最后拼装成layout格式返回
|
||||
"""
|
||||
for block in v_blocks:
|
||||
sorted_layout_blocks.append(
|
||||
{
|
||||
'layout_bbox': block[:4],
|
||||
'layout_label': block[4],
|
||||
'sub_layout': [],
|
||||
}
|
||||
)
|
||||
for block in unsplited_block:
|
||||
sorted_layout_blocks.append(
|
||||
{
|
||||
'layout_bbox': block[:4],
|
||||
'layout_label': block[4],
|
||||
'sub_layout': [],
|
||||
}
|
||||
)
|
||||
|
||||
# 按照x0排序
|
||||
sorted_layout_blocks.sort(key=lambda x: x['layout_bbox'][0])
|
||||
return sorted_layout_blocks
|
||||
|
||||
|
||||
def _vertical_align_split_v2(bboxes: list, boundary: tuple) -> list:
|
||||
"""改进的
|
||||
_vertical_align_split算法，原算法会因为第二列的box由于左侧没有遮挡被认为是左侧的一部分，导致整个layout多列被识别为一列。
|
||||
利用从左上角的box开始向下看的方法，不断扩展w_x0, w_x1，直到不能继续向下扩展，或者到达边界下边界。"""
|
||||
sorted_layout_blocks = [] # 这是要最终返回的值
|
||||
new_boundary = [boundary[0], boundary[1], boundary[2], boundary[3]]
|
||||
bad_boxes = [] # 被割中的box
|
||||
v_blocks = []
|
||||
while True:
|
||||
all_bboxes = get_bbox_in_boundary(bboxes, new_boundary)
|
||||
if len(all_bboxes) == 0:
|
||||
break
|
||||
left_top_box = min(
|
||||
all_bboxes, key=lambda x: (x[X0_IDX], x[Y0_IDX])
|
||||
) # 这里应该加强，检查一下必须是在第一列的 TODO
|
||||
start_box = [
|
||||
left_top_box[X0_IDX],
|
||||
left_top_box[Y0_IDX],
|
||||
left_top_box[X1_IDX],
|
||||
left_top_box[Y1_IDX],
|
||||
]
|
||||
w_x0, w_x1 = left_top_box[X0_IDX], left_top_box[X1_IDX]
|
||||
"""
|
||||
然后沿着这个box线向下找最近的那个box, 然后扩展w_x0, w_x1
|
||||
扩展之后，宽度会增加，随后用x=w_x1来检测在边界内是否有box与相交，如果相交，那么就说明不能再扩展了。
|
||||
当不能扩展的时候就要看是否到达下边界：
|
||||
1. 达到，那么更新左边界继续分下一个列
|
||||
2. 没有达到，那么此时开始从右侧切分进入下面的循环里
|
||||
"""
|
||||
while left_top_box is not None: # 向下去找
|
||||
virtual_box = [w_x0, left_top_box[Y0_IDX], w_x1, left_top_box[Y1_IDX]]
|
||||
left_top_box = find_bottom_bbox_direct_from_left_edge(
|
||||
virtual_box, all_bboxes
|
||||
)
|
||||
if left_top_box:
|
||||
w_x0, w_x1 = min(virtual_box[X0_IDX], left_top_box[X0_IDX]), max(
|
||||
[virtual_box[X1_IDX], left_top_box[X1_IDX]]
|
||||
)
|
||||
# 万一这个初始的box在column中间，那么还要向上看
|
||||
start_box = [
|
||||
w_x0,
|
||||
start_box[Y0_IDX],
|
||||
w_x1,
|
||||
start_box[Y1_IDX],
|
||||
] # 扩展一下宽度更鲁棒
|
||||
left_top_box = find_top_bbox_direct_from_left_edge(start_box, all_bboxes)
|
||||
while left_top_box is not None: # 向上去找
|
||||
virtual_box = [w_x0, left_top_box[Y0_IDX], w_x1, left_top_box[Y1_IDX]]
|
||||
left_top_box = find_top_bbox_direct_from_left_edge(virtual_box, all_bboxes)
|
||||
if left_top_box:
|
||||
w_x0, w_x1 = min(virtual_box[X0_IDX], left_top_box[X0_IDX]), max(
|
||||
[virtual_box[X1_IDX], left_top_box[X1_IDX]]
|
||||
)
|
||||
|
||||
# 检查相交
|
||||
if any([bbox[X0_IDX] <= w_x1 + 1 <= bbox[X1_IDX] for bbox in all_bboxes]):
|
||||
for b in all_bboxes:
|
||||
if b[X0_IDX] <= w_x1 + 1 <= b[X1_IDX]:
|
||||
bad_boxes.append([b[X0_IDX], b[Y0_IDX], b[X1_IDX], b[Y1_IDX]])
|
||||
break
|
||||
else: # 说明成功分割出一列
|
||||
v_blocks.append([w_x0, new_boundary[1], w_x1, new_boundary[3], LAYOUT_V])
|
||||
new_boundary[0] = w_x1 # 更新边界
|
||||
"""
|
||||
接着开始从右上角的box扫描
|
||||
"""
|
||||
w_x0, w_x1 = 0, 0
|
||||
unsplited_block = []
|
||||
while True:
|
||||
all_bboxes = get_bbox_in_boundary(bboxes, new_boundary)
|
||||
if len(all_bboxes) == 0:
|
||||
break
|
||||
# 先找到X1最大的
|
||||
bbox_list_sorted = sorted(
|
||||
all_bboxes, key=lambda bbox: bbox[X1_IDX], reverse=True
|
||||
)
|
||||
# Then, find the boxes with the smallest Y0 value
|
||||
bigest_x1 = bbox_list_sorted[0][X1_IDX]
|
||||
boxes_with_bigest_x1 = [
|
||||
bbox for bbox in bbox_list_sorted if bbox[X1_IDX] == bigest_x1
|
||||
] # 也就是最靠右的那些
|
||||
right_top_box = min(
|
||||
boxes_with_bigest_x1, key=lambda bbox: bbox[Y0_IDX]
|
||||
) # y0最小的那个
|
||||
start_box = [
|
||||
right_top_box[X0_IDX],
|
||||
right_top_box[Y0_IDX],
|
||||
right_top_box[X1_IDX],
|
||||
right_top_box[Y1_IDX],
|
||||
]
|
||||
w_x0, w_x1 = right_top_box[X0_IDX], right_top_box[X1_IDX]
|
||||
|
||||
while right_top_box is not None:
|
||||
virtual_box = [w_x0, right_top_box[Y0_IDX], w_x1, right_top_box[Y1_IDX]]
|
||||
right_top_box = find_bottom_bbox_direct_from_right_edge(
|
||||
virtual_box, all_bboxes
|
||||
)
|
||||
if right_top_box:
|
||||
w_x0, w_x1 = min([w_x0, right_top_box[X0_IDX]]), max(
|
||||
[w_x1, right_top_box[X1_IDX]]
|
||||
)
|
||||
# 在向上扫描
|
||||
start_box = [
|
||||
w_x0,
|
||||
start_box[Y0_IDX],
|
||||
w_x1,
|
||||
start_box[Y1_IDX],
|
||||
] # 扩展一下宽度更鲁棒
|
||||
right_top_box = find_top_bbox_direct_from_right_edge(start_box, all_bboxes)
|
||||
while right_top_box is not None:
|
||||
virtual_box = [w_x0, right_top_box[Y0_IDX], w_x1, right_top_box[Y1_IDX]]
|
||||
right_top_box = find_top_bbox_direct_from_right_edge(
|
||||
virtual_box, all_bboxes
|
||||
)
|
||||
if right_top_box:
|
||||
w_x0, w_x1 = min([w_x0, right_top_box[X0_IDX]]), max(
|
||||
[w_x1, right_top_box[X1_IDX]]
|
||||
)
|
||||
|
||||
# 检查是否与其他box相交， 垂直切分线与某些box发生相交，说明无法完全垂直方向切分。
|
||||
if any([bbox[X0_IDX] <= w_x0 - 1 <= bbox[X1_IDX] for bbox in all_bboxes]):
|
||||
unsplited_block.append(
|
||||
[
|
||||
new_boundary[0],
|
||||
new_boundary[1],
|
||||
new_boundary[2],
|
||||
new_boundary[3],
|
||||
LAYOUT_UNPROC,
|
||||
]
|
||||
)
|
||||
for b in all_bboxes:
|
||||
if b[X0_IDX] <= w_x0 - 1 <= b[X1_IDX]:
|
||||
bad_boxes.append([b[X0_IDX], b[Y0_IDX], b[X1_IDX], b[Y1_IDX]])
|
||||
break
|
||||
else: # 说明成功分割出一列
|
||||
v_blocks.append([w_x0, new_boundary[1], w_x1, new_boundary[3], LAYOUT_V])
|
||||
new_boundary[2] = w_x0
|
||||
"""转换数据结构"""
|
||||
for block in v_blocks:
|
||||
sorted_layout_blocks.append(
|
||||
{
|
||||
'layout_bbox': block[:4],
|
||||
'layout_label': block[4],
|
||||
'sub_layout': [],
|
||||
}
|
||||
)
|
||||
|
||||
for block in unsplited_block:
|
||||
sorted_layout_blocks.append(
|
||||
{
|
||||
'layout_bbox': block[:4],
|
||||
'layout_label': block[4],
|
||||
'sub_layout': [],
|
||||
'bad_boxes': bad_boxes, # 记录下来，这个box是被割中的
|
||||
}
|
||||
)
|
||||
|
||||
# 按照x0排序
|
||||
sorted_layout_blocks.sort(key=lambda x: x['layout_bbox'][0])
|
||||
return sorted_layout_blocks
|
||||
|
||||
|
||||
def _try_horizontal_mult_column_split(bboxes: list, boundary: tuple) -> list:
|
||||
"""
|
||||
尝试水平切分，如果切分不动，那就当一个BAD_LAYOUT返回
|
||||
------------------
|
||||
| | |
|
||||
------------------
|
||||
| | | | <- 这里是此函数要切分的场景
|
||||
------------------
|
||||
| | |
|
||||
| | |
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
def _vertical_split(bboxes: list, boundary: tuple) -> list:
|
||||
"""
|
||||
从垂直方向进行切割，分block
|
||||
这个版本里，如果垂直切分不动，那就当一个BAD_LAYOUT返回
|
||||
|
||||
--------------------------
|
||||
| | |
|
||||
| | |
|
||||
| |
|
||||
这种列是此函数要切分的 -> | |
|
||||
| |
|
||||
| | |
|
||||
| | |
|
||||
-------------------------
|
||||
"""
|
||||
sorted_layout_blocks = [] # 这是要最终返回的值
|
||||
|
||||
bound_x0, bound_y0, bound_x1, bound_y1 = boundary
|
||||
all_bboxes = get_bbox_in_boundary(bboxes, boundary)
|
||||
"""
|
||||
all_bboxes = fix_vertical_bbox_pos(all_bboxes) # 垂直方向解覆盖
|
||||
all_bboxes = fix_hor_bbox_pos(all_bboxes) # 水平解覆盖
|
||||
|
||||
这两行代码目前先不执行，因为公式检测，表格检测还不是很成熟，导致非常多的textblock参与了运算，时间消耗太大。
|
||||
这两行代码的作用是：
|
||||
如果遇到互相重叠的bbox, 那么会把面积较小的box进行压缩，从而避免重叠。对布局切分来说带来正反馈。
|
||||
"""
|
||||
|
||||
# all_bboxes = paper_bbox_sort(all_bboxes, abs(bound_x1-bound_x0), abs(bound_y1-bound_x0)) # 大致拍下序, 这个是基于直接遮挡的。
|
||||
"""
|
||||
首先在垂直方向上扩展独占一行的bbox
|
||||
|
||||
"""
|
||||
for bbox in all_bboxes:
|
||||
top_nearest_bbox = find_all_top_bbox_direct(bbox, all_bboxes) # 非扩展线
|
||||
bottom_nearest_bbox = find_all_bottom_bbox_direct(bbox, all_bboxes)
|
||||
if (
|
||||
top_nearest_bbox is None
|
||||
and bottom_nearest_bbox is None
|
||||
and not any(
|
||||
[
|
||||
b[X0_IDX] < bbox[X1_IDX] < b[X1_IDX]
|
||||
or b[X0_IDX] < bbox[X0_IDX] < b[X1_IDX]
|
||||
for b in all_bboxes
|
||||
]
|
||||
)
|
||||
): # 独占一列, 且不和其他重叠
|
||||
bbox[X0_EXT_IDX] = bbox[X0_IDX]
|
||||
bbox[Y0_EXT_IDX] = bound_y0
|
||||
bbox[X1_EXT_IDX] = bbox[X1_IDX]
|
||||
bbox[Y1_EXT_IDX] = bound_y1
|
||||
"""
|
||||
此时独占一列的被成功扩展到指定的边界上，这个时候利用边界条件合并连续的bbox，成为一个group
|
||||
然后合并所有连续垂直方向的bbox.
|
||||
"""
|
||||
all_bboxes.sort(key=lambda x: x[X0_IDX])
|
||||
# fix: 这里水平方向的列不要合并成一个行，因为需要保证返回给下游的最小block，总是可以无脑从上到下阅读文字。
|
||||
v_bboxes = []
|
||||
for box in all_bboxes:
|
||||
if box[Y0_EXT_IDX] == bound_y0 and box[Y1_EXT_IDX] == bound_y1:
|
||||
v_bboxes.append(box)
|
||||
"""
|
||||
现在v_bboxes里面是所有的group了，每个group都是一个list
|
||||
对v_bboxes里的每个group进行计算放回到sorted_layouts里
|
||||
"""
|
||||
v_layouts = []
|
||||
for vbox in v_bboxes:
|
||||
# gp.sort(key=lambda x: x[X0_IDX])
|
||||
# 然后计算这个group的layout_bbox，也就是最小的x0,y0, 最大的x1,y1
|
||||
x0, y0, x1, y1 = (
|
||||
vbox[X0_EXT_IDX],
|
||||
vbox[Y0_EXT_IDX],
|
||||
vbox[X1_EXT_IDX],
|
||||
vbox[Y1_EXT_IDX],
|
||||
)
|
||||
v_layouts.append([x0, y0, x1, y1, LAYOUT_V]) # 垂直的布局
|
||||
"""
|
||||
接下来利用这些连续的垂直bbox的layout_bbox的x0, x1，从垂直上切分开其余的为几个部分
|
||||
"""
|
||||
v_split_lines = [bound_x0]
|
||||
for gp in v_bboxes:
|
||||
x0, x1 = gp[X0_IDX], gp[X1_IDX]
|
||||
v_split_lines.append(x0)
|
||||
v_split_lines.append(x1)
|
||||
v_split_lines.append(bound_x1)
|
||||
|
||||
unsplited_bboxes = []
|
||||
for i in range(0, len(v_split_lines), 2):
|
||||
start_x0, start_x1 = v_split_lines[i : i + 2]
|
||||
# 然后找出[start_x0, start_x1]之间的其他bbox，这些组成一个未分割板块
|
||||
bboxes_in_block = [
|
||||
bbox
|
||||
for bbox in all_bboxes
|
||||
if bbox[X0_IDX] >= start_x0 and bbox[X1_IDX] <= start_x1
|
||||
]
|
||||
unsplited_bboxes.append(bboxes_in_block)
|
||||
# 接着把未处理的加入到v_layouts里
|
||||
for bboxes_in_block in unsplited_bboxes:
|
||||
if len(bboxes_in_block) == 0:
|
||||
continue
|
||||
x0, y0, x1, y1 = (
|
||||
min([bbox[X0_IDX] for bbox in bboxes_in_block]),
|
||||
bound_y0,
|
||||
max([bbox[X1_IDX] for bbox in bboxes_in_block]),
|
||||
bound_y1,
|
||||
)
|
||||
v_layouts.append(
|
||||
[x0, y0, x1, y1, LAYOUT_UNPROC]
|
||||
) # 说明这篇区域未能够分析出可靠的版面
|
||||
|
||||
v_layouts.sort(key=lambda x: x[0]) # 按照x0排序, 也就是从左到右的顺序
|
||||
|
||||
for layout in v_layouts:
|
||||
sorted_layout_blocks.append(
|
||||
{
|
||||
'layout_bbox': layout[:4],
|
||||
'layout_label': layout[4],
|
||||
'sub_layout': [],
|
||||
}
|
||||
)
|
||||
"""
|
||||
至此，垂直方向切成了2种类型，其一是独占一列的，其二是未处理的。
|
||||
下面对这些未处理的进行垂直方向切分，这个切分要切出来类似“吕”这种类型的垂直方向的布局
|
||||
"""
|
||||
for i, layout in enumerate(sorted_layout_blocks):
|
||||
if layout['layout_label'] == LAYOUT_UNPROC:
|
||||
x0, y0, x1, y1 = layout['layout_bbox']
|
||||
v_split_layouts = _vertical_align_split_v2(bboxes, [x0, y0, x1, y1])
|
||||
sorted_layout_blocks[i] = {
|
||||
'layout_bbox': [x0, y0, x1, y1],
|
||||
'layout_label': LAYOUT_H,
|
||||
'sub_layout': v_split_layouts,
|
||||
}
|
||||
layout['layout_label'] = LAYOUT_H # 被垂线切分成了水平布局
|
||||
|
||||
return sorted_layout_blocks
|
||||
|
||||
|
||||
def split_layout(bboxes: list, boundary: tuple, page_num: int) -> list:
|
||||
"""
|
||||
把bboxes切割成layout
|
||||
return:
|
||||
[
|
||||
{
|
||||
"layout_bbox": [x0,y0,x1,y1],
|
||||
"layout_label":"u|v|h|b", 未处理|垂直|水平|BAD_LAYOUT
|
||||
"sub_layout":[] #每个元素都是[
|
||||
x0,y0,
|
||||
x1,y1,
|
||||
block_content,
|
||||
idx_x,idx_y,
|
||||
content_type,
|
||||
ext_x0,ext_y0,
|
||||
ext_x1,ext_y1
|
||||
], 并且顺序就是阅读顺序
|
||||
}
|
||||
]
|
||||
example:
|
||||
[
|
||||
{
|
||||
"layout_bbox": [0, 0, 100, 100],
|
||||
"layout_label":"u|v|h|b",
|
||||
"sub_layout":[
|
||||
|
||||
]
|
||||
},
|
||||
{
|
||||
"layout_bbox": [0, 0, 100, 100],
|
||||
"layout_label":"u|v|h|b",
|
||||
"sub_layout":[
|
||||
{
|
||||
"layout_bbox": [0, 0, 100, 100],
|
||||
"layout_label":"u|v|h|b",
|
||||
"content_bboxes":[
|
||||
[],
|
||||
[],
|
||||
[]
|
||||
]
|
||||
},
|
||||
{
|
||||
"layout_bbox": [0, 0, 100, 100],
|
||||
"layout_label":"u|v|h|b",
|
||||
"sub_layout":[
|
||||
|
||||
]
|
||||
}
|
||||
}
|
||||
]
|
||||
"""
|
||||
sorted_layouts = [] # 最终返回的结果
|
||||
|
||||
boundary_x0, boundary_y0, boundary_x1, boundary_y1 = boundary
|
||||
if len(bboxes) <= 1:
|
||||
return [
|
||||
{
|
||||
'layout_bbox': [boundary_x0, boundary_y0, boundary_x1, boundary_y1],
|
||||
'layout_label': LAYOUT_V,
|
||||
'sub_layout': [],
|
||||
}
|
||||
]
|
||||
"""
|
||||
接下来按照先水平后垂直的顺序进行切分
|
||||
"""
|
||||
bboxes = paper_bbox_sort(
|
||||
bboxes, boundary_x1 - boundary_x0, boundary_y1 - boundary_y0
|
||||
)
|
||||
sorted_layouts = _horizontal_split(bboxes, boundary) # 通过水平分割出来的layout
|
||||
for i, layout in enumerate(sorted_layouts):
|
||||
x0, y0, x1, y1 = layout['layout_bbox']
|
||||
layout_type = layout['layout_label']
|
||||
if layout_type == LAYOUT_UNPROC: # 说明是非独占单行的，这些需要垂直切分
|
||||
v_split_layouts = _vertical_split(bboxes, [x0, y0, x1, y1])
|
||||
"""
|
||||
最后这里有个逻辑问题：如果这个函数只分离出来了一个column layout，那么这个layout分割肯定超出了算法能力范围。因为我们假定的是传进来的
|
||||
box已经把行全部剥离了，所以这里必须十多个列才可以。如果只剥离出来一个layout，并且是多个box，那么就说明这个layout是无法分割的，标记为LAYOUT_UNPROC
|
||||
"""
|
||||
layout_label = LAYOUT_V
|
||||
if len(v_split_layouts) == 1:
|
||||
if len(v_split_layouts[0]['sub_layout']) == 0:
|
||||
layout_label = LAYOUT_UNPROC
|
||||
# logger.warning(f"WARNING: pageno={page_num}, 无法分割的layout: ", v_split_layouts)
|
||||
"""
|
||||
组合起来最终的layout
|
||||
"""
|
||||
sorted_layouts[i] = {
|
||||
'layout_bbox': [x0, y0, x1, y1],
|
||||
'layout_label': layout_label,
|
||||
'sub_layout': v_split_layouts,
|
||||
}
|
||||
layout['layout_label'] = LAYOUT_H
|
||||
"""
|
||||
水平和垂直方向都切分完毕了。此时还有一些未处理的，这些未处理的可能是因为水平和垂直方向都无法切分。
|
||||
这些最后调用_try_horizontal_mult_block_split做一次水平多个block的联合切分，如果也不能切分最终就当做BAD_LAYOUT返回
|
||||
"""
|
||||
# TODO
|
||||
|
||||
return sorted_layouts
|
||||
|
||||
|
||||
def get_bboxes_layout(all_boxes: list, boundary: tuple, page_id: int):
|
||||
"""
|
||||
对利用layout排序之后的box，进行排序
|
||||
return:
|
||||
[
|
||||
{
|
||||
"layout_bbox": [x0, y0, x1, y1],
|
||||
"layout_label":"u|v|h|b", 未处理|垂直|水平|BAD_LAYOUT
|
||||
}，
|
||||
]
|
||||
"""
|
||||
|
||||
def _preorder_traversal(layout):
|
||||
"""对sorted_layouts的叶子节点，也就是len(sub_layout)==0的节点进行排序。排序按照前序遍历的顺序，也就是从上到
|
||||
下，从左到右的顺序."""
|
||||
sorted_layout_blocks = []
|
||||
for layout in layout:
|
||||
sub_layout = layout['sub_layout']
|
||||
if len(sub_layout) == 0:
|
||||
sorted_layout_blocks.append(layout)
|
||||
else:
|
||||
s = _preorder_traversal(sub_layout)
|
||||
sorted_layout_blocks.extend(s)
|
||||
return sorted_layout_blocks
|
||||
|
||||
# -------------------------------------------------------------------------------------------------------------------------
|
||||
sorted_layouts = split_layout(
|
||||
all_boxes, boundary, page_id
|
||||
) # 先切分成layout，得到一个Tree
|
||||
total_sorted_layout_blocks = _preorder_traversal(sorted_layouts)
|
||||
return total_sorted_layout_blocks, sorted_layouts
|
||||
|
||||
|
||||
def get_columns_cnt_of_layout(layout_tree):
|
||||
"""获取一个layout的宽度."""
|
||||
max_width_list = [0] # 初始化一个元素，防止max,min函数报错
|
||||
|
||||
for items in layout_tree: # 针对每一层（横切）计算列数，横着的算一列
|
||||
layout_type = items['layout_label']
|
||||
sub_layouts = items['sub_layout']
|
||||
if len(sub_layouts) == 0:
|
||||
max_width_list.append(1)
|
||||
else:
|
||||
if layout_type == LAYOUT_H:
|
||||
max_width_list.append(1)
|
||||
else:
|
||||
width = 0
|
||||
for sub_layout in sub_layouts:
|
||||
if len(sub_layout['sub_layout']) == 0:
|
||||
width += 1
|
||||
else:
|
||||
for lay in sub_layout['sub_layout']:
|
||||
width += get_columns_cnt_of_layout([lay])
|
||||
max_width_list.append(width)
|
||||
|
||||
return max(max_width_list)
|
||||
|
||||
|
||||
def sort_with_layout(bboxes: list, page_width, page_height) -> (list, list):
|
||||
"""输入是一个bbox的list.
|
||||
|
||||
获取到输入之后，先进行layout切分，然后对这些bbox进行排序。返回排序后的bboxes
|
||||
"""
|
||||
|
||||
new_bboxes = []
|
||||
for box in bboxes:
|
||||
# new_bboxes.append([box[0], box[1], box[2], box[3], None, None, None, 'text', None, None, None, None])
|
||||
new_bboxes.append(
|
||||
[
|
||||
box[0],
|
||||
box[1],
|
||||
box[2],
|
||||
box[3],
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
'text',
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
box[4],
|
||||
]
|
||||
)
|
||||
|
||||
layout_bboxes, _ = get_bboxes_layout(
|
||||
new_bboxes, tuple([0, 0, page_width, page_height]), 0
|
||||
)
|
||||
if any([lay['layout_label'] == LAYOUT_UNPROC for lay in layout_bboxes]):
|
||||
logger.warning('drop this pdf, reason: 复杂版面')
|
||||
return None, None
|
||||
|
||||
sorted_bboxes = []
|
||||
# 利用layout bbox每次框定一些box，然后排序
|
||||
for layout in layout_bboxes:
|
||||
lbox = layout['layout_bbox']
|
||||
bbox_in_layout = get_bbox_in_boundary(new_bboxes, lbox)
|
||||
sorted_bbox = paper_bbox_sort(
|
||||
bbox_in_layout, lbox[2] - lbox[0], lbox[3] - lbox[1]
|
||||
)
|
||||
sorted_bboxes.extend(sorted_bbox)
|
||||
|
||||
return sorted_bboxes, layout_bboxes
|
||||
|
||||
|
||||
def sort_text_block(text_block, layout_bboxes):
|
||||
"""对一页的text_block进行排序."""
|
||||
sorted_text_bbox = []
|
||||
all_text_bbox = []
|
||||
# 做一个box=>text的映射
|
||||
box_to_text = {}
|
||||
for blk in text_block:
|
||||
box = blk['bbox']
|
||||
box_to_text[(box[0], box[1], box[2], box[3])] = blk
|
||||
all_text_bbox.append(box)
|
||||
|
||||
# text_blocks_to_sort = []
|
||||
# for box in box_to_text.keys():
|
||||
# text_blocks_to_sort.append([box[0], box[1], box[2], box[3], None, None, None, 'text', None, None, None, None])
|
||||
|
||||
# 按照layout_bboxes的顺序，对text_block进行排序
|
||||
for layout in layout_bboxes:
|
||||
layout_box = layout['layout_bbox']
|
||||
text_bbox_in_layout = get_bbox_in_boundary(
|
||||
all_text_bbox,
|
||||
[
|
||||
layout_box[0] - 1,
|
||||
layout_box[1] - 1,
|
||||
layout_box[2] + 1,
|
||||
layout_box[3] + 1,
|
||||
],
|
||||
)
|
||||
# sorted_bbox = paper_bbox_sort(text_bbox_in_layout, layout_box[2]-layout_box[0], layout_box[3]-layout_box[1])
|
||||
text_bbox_in_layout.sort(
|
||||
key=lambda x: x[1]
|
||||
) # 一个layout内部的box，按照y0自上而下排序
|
||||
# sorted_bbox = [[b] for b in text_blocks_to_sort]
|
||||
for sb in text_bbox_in_layout:
|
||||
sorted_text_bbox.append(box_to_text[(sb[0], sb[1], sb[2], sb[3])])
|
||||
|
||||
return sorted_text_bbox
|
||||
@@ -1,101 +0,0 @@
|
||||
"""
|
||||
找到能分割布局的水平的横线、色块
|
||||
"""
|
||||
|
||||
import os
|
||||
from magic_pdf.libs.commons import fitz
|
||||
from magic_pdf.libs.boxbase import _is_in_or_part_overlap
|
||||
|
||||
|
||||
def __rect_filter_by_width(rect, page_w, page_h):
|
||||
mid_x = page_w/2
|
||||
if rect[0]< mid_x < rect[2]:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def __rect_filter_by_pos(rect, image_bboxes, table_bboxes):
|
||||
"""
|
||||
不能出现在table和image的位置
|
||||
"""
|
||||
for box in image_bboxes:
|
||||
if _is_in_or_part_overlap(rect, box):
|
||||
return False
|
||||
|
||||
for box in table_bboxes:
|
||||
if _is_in_or_part_overlap(rect, box):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def __debug_show_page(page, bboxes1: list,bboxes2: list,bboxes3: list,):
|
||||
save_path = "./tmp/debug.pdf"
|
||||
if os.path.exists(save_path):
|
||||
# 删除已经存在的文件
|
||||
os.remove(save_path)
|
||||
# 创建一个新的空白 PDF 文件
|
||||
doc = fitz.open('')
|
||||
|
||||
width = page.rect.width
|
||||
height = page.rect.height
|
||||
new_page = doc.new_page(width=width, height=height)
|
||||
|
||||
shape = new_page.new_shape()
|
||||
for bbox in bboxes1:
|
||||
# 原始box画上去
|
||||
rect = fitz.Rect(*bbox[0:4])
|
||||
shape = new_page.new_shape()
|
||||
shape.draw_rect(rect)
|
||||
shape.finish(color=fitz.pdfcolor['red'], fill=fitz.pdfcolor['blue'], fill_opacity=0.2)
|
||||
shape.finish()
|
||||
shape.commit()
|
||||
|
||||
for bbox in bboxes2:
|
||||
# 原始box画上去
|
||||
rect = fitz.Rect(*bbox[0:4])
|
||||
shape = new_page.new_shape()
|
||||
shape.draw_rect(rect)
|
||||
shape.finish(color=None, fill=fitz.pdfcolor['yellow'], fill_opacity=0.2)
|
||||
shape.finish()
|
||||
shape.commit()
|
||||
|
||||
for bbox in bboxes3:
|
||||
# 原始box画上去
|
||||
rect = fitz.Rect(*bbox[0:4])
|
||||
shape = new_page.new_shape()
|
||||
shape.draw_rect(rect)
|
||||
shape.finish(color=fitz.pdfcolor['red'], fill=None)
|
||||
shape.finish()
|
||||
shape.commit()
|
||||
|
||||
parent_dir = os.path.dirname(save_path)
|
||||
if not os.path.exists(parent_dir):
|
||||
os.makedirs(parent_dir)
|
||||
|
||||
doc.save(save_path)
|
||||
doc.close()
|
||||
|
||||
def get_spilter_of_page(page, image_bboxes, table_bboxes):
|
||||
"""
|
||||
获取到色块和横线
|
||||
"""
|
||||
cdrawings = page.get_cdrawings()
|
||||
|
||||
spilter_bbox = []
|
||||
for block in cdrawings:
|
||||
if 'fill' in block:
|
||||
fill = block['fill']
|
||||
if 'fill' in block and block['fill'] and block['fill']!=(1.0,1.0,1.0):
|
||||
rect = block['rect']
|
||||
if __rect_filter_by_width(rect, page.rect.width, page.rect.height) and __rect_filter_by_pos(rect, image_bboxes, table_bboxes):
|
||||
spilter_bbox.append(list(rect))
|
||||
|
||||
"""过滤、修正一下这些box。因为有时候会有一些矩形，高度为0或者为负数，造成layout计算无限循环。如果是负高度或者0高度，统一修正为高度为1"""
|
||||
for box in spilter_bbox:
|
||||
if box[3]-box[1] <= 0:
|
||||
box[3] = box[1] + 1
|
||||
|
||||
#__debug_show_page(page, spilter_bbox, [], [])
|
||||
|
||||
return spilter_bbox
|
||||
@@ -1,336 +0,0 @@
|
||||
"""
|
||||
This is an advanced PyMuPDF utility for detecting multi-column pages.
|
||||
It can be used in a shell script, or its main function can be imported and
|
||||
invoked as descript below.
|
||||
|
||||
Features
|
||||
---------
|
||||
- Identify text belonging to (a variable number of) columns on the page.
|
||||
- Text with different background color is handled separately, allowing for
|
||||
easier treatment of side remarks, comment boxes, etc.
|
||||
- Uses text block detection capability to identify text blocks and
|
||||
uses the block bboxes as primary structuring principle.
|
||||
- Supports ignoring footers via a footer margin parameter.
|
||||
- Returns re-created text boundary boxes (integer coordinates), sorted ascending
|
||||
by the top, then by the left coordinates.
|
||||
|
||||
Restrictions
|
||||
-------------
|
||||
- Only supporting horizontal, left-to-right text
|
||||
- Returns a list of text boundary boxes - not the text itself. The caller is
|
||||
expected to extract text from within the returned boxes.
|
||||
- Text written above images is ignored altogether (option).
|
||||
- This utility works as expected in most cases. The following situation cannot
|
||||
be handled correctly:
|
||||
* overlapping (non-disjoint) text blocks
|
||||
* image captions are not recognized and are handled like normal text
|
||||
|
||||
Usage
|
||||
------
|
||||
- As a CLI shell command use
|
||||
|
||||
python multi_column.py input.pdf footer_margin
|
||||
|
||||
Where footer margin is the height of the bottom stripe to ignore on each page.
|
||||
This code is intended to be modified according to your need.
|
||||
|
||||
- Use in a Python script as follows:
|
||||
|
||||
----------------------------------------------------------------------------------
|
||||
from multi_column import column_boxes
|
||||
|
||||
# for each page execute
|
||||
bboxes = column_boxes(page, footer_margin=50, no_image_text=True)
|
||||
|
||||
# bboxes is a list of fitz.IRect objects, that are sort ascending by their y0,
|
||||
# then x0 coordinates. Their text content can be extracted by all PyMuPDF
|
||||
# get_text() variants, like for instance the following:
|
||||
for rect in bboxes:
|
||||
print(page.get_text(clip=rect, sort=True))
|
||||
----------------------------------------------------------------------------------
|
||||
"""
|
||||
import sys
|
||||
from magic_pdf.libs.commons import fitz
|
||||
|
||||
|
||||
def column_boxes(page, footer_margin=50, header_margin=50, no_image_text=True):
|
||||
"""Determine bboxes which wrap a column."""
|
||||
paths = page.get_drawings()
|
||||
bboxes = []
|
||||
|
||||
# path rectangles
|
||||
path_rects = []
|
||||
|
||||
# image bboxes
|
||||
img_bboxes = []
|
||||
|
||||
# bboxes of non-horizontal text
|
||||
# avoid when expanding horizontal text boxes
|
||||
vert_bboxes = []
|
||||
|
||||
# compute relevant page area
|
||||
clip = +page.rect
|
||||
clip.y1 -= footer_margin # Remove footer area
|
||||
clip.y0 += header_margin # Remove header area
|
||||
|
||||
def can_extend(temp, bb, bboxlist):
|
||||
"""Determines whether rectangle 'temp' can be extended by 'bb'
|
||||
without intersecting any of the rectangles contained in 'bboxlist'.
|
||||
|
||||
Items of bboxlist may be None if they have been removed.
|
||||
|
||||
Returns:
|
||||
True if 'temp' has no intersections with items of 'bboxlist'.
|
||||
"""
|
||||
for b in bboxlist:
|
||||
if not intersects_bboxes(temp, vert_bboxes) and (
|
||||
b == None or b == bb or (temp & b).is_empty
|
||||
):
|
||||
continue
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def in_bbox(bb, bboxes):
|
||||
"""Return 1-based number if a bbox contains bb, else return 0."""
|
||||
for i, bbox in enumerate(bboxes):
|
||||
if bb in bbox:
|
||||
return i + 1
|
||||
return 0
|
||||
|
||||
def intersects_bboxes(bb, bboxes):
|
||||
"""Return True if a bbox intersects bb, else return False."""
|
||||
for bbox in bboxes:
|
||||
if not (bb & bbox).is_empty:
|
||||
return True
|
||||
return False
|
||||
|
||||
def extend_right(bboxes, width, path_bboxes, vert_bboxes, img_bboxes):
|
||||
"""Extend a bbox to the right page border.
|
||||
|
||||
Whenever there is no text to the right of a bbox, enlarge it up
|
||||
to the right page border.
|
||||
|
||||
Args:
|
||||
bboxes: (list[IRect]) bboxes to check
|
||||
width: (int) page width
|
||||
path_bboxes: (list[IRect]) bboxes with a background color
|
||||
vert_bboxes: (list[IRect]) bboxes with vertical text
|
||||
img_bboxes: (list[IRect]) bboxes of images
|
||||
Returns:
|
||||
Potentially modified bboxes.
|
||||
"""
|
||||
for i, bb in enumerate(bboxes):
|
||||
# do not extend text with background color
|
||||
if in_bbox(bb, path_bboxes):
|
||||
continue
|
||||
|
||||
# do not extend text in images
|
||||
if in_bbox(bb, img_bboxes):
|
||||
continue
|
||||
|
||||
# temp extends bb to the right page border
|
||||
temp = +bb
|
||||
temp.x1 = width
|
||||
|
||||
# do not cut through colored background or images
|
||||
if intersects_bboxes(temp, path_bboxes + vert_bboxes + img_bboxes):
|
||||
continue
|
||||
|
||||
# also, do not intersect other text bboxes
|
||||
check = can_extend(temp, bb, bboxes)
|
||||
if check:
|
||||
bboxes[i] = temp # replace with enlarged bbox
|
||||
|
||||
return [b for b in bboxes if b != None]
|
||||
|
||||
def clean_nblocks(nblocks):
|
||||
"""Do some elementary cleaning."""
|
||||
|
||||
# 1. remove any duplicate blocks.
|
||||
blen = len(nblocks)
|
||||
if blen < 2:
|
||||
return nblocks
|
||||
start = blen - 1
|
||||
for i in range(start, -1, -1):
|
||||
bb1 = nblocks[i]
|
||||
bb0 = nblocks[i - 1]
|
||||
if bb0 == bb1:
|
||||
del nblocks[i]
|
||||
|
||||
# 2. repair sequence in special cases:
|
||||
# consecutive bboxes with almost same bottom value are sorted ascending
|
||||
# by x-coordinate.
|
||||
y1 = nblocks[0].y1 # first bottom coordinate
|
||||
i0 = 0 # its index
|
||||
i1 = -1 # index of last bbox with same bottom
|
||||
|
||||
# Iterate over bboxes, identifying segments with approx. same bottom value.
|
||||
# Replace every segment by its sorted version.
|
||||
for i in range(1, len(nblocks)):
|
||||
b1 = nblocks[i]
|
||||
if abs(b1.y1 - y1) > 10: # different bottom
|
||||
if i1 > i0: # segment length > 1? Sort it!
|
||||
nblocks[i0 : i1 + 1] = sorted(
|
||||
nblocks[i0 : i1 + 1], key=lambda b: b.x0
|
||||
)
|
||||
y1 = b1.y1 # store new bottom value
|
||||
i0 = i # store its start index
|
||||
i1 = i # store current index
|
||||
if i1 > i0: # segment waiting to be sorted
|
||||
nblocks[i0 : i1 + 1] = sorted(nblocks[i0 : i1 + 1], key=lambda b: b.x0)
|
||||
return nblocks
|
||||
|
||||
# extract vector graphics
|
||||
for p in paths:
|
||||
path_rects.append(p["rect"].irect)
|
||||
path_bboxes = path_rects
|
||||
|
||||
# sort path bboxes by ascending top, then left coordinates
|
||||
path_bboxes.sort(key=lambda b: (b.y0, b.x0))
|
||||
|
||||
# bboxes of images on page, no need to sort them
|
||||
for item in page.get_images():
|
||||
img_bboxes.extend(page.get_image_rects(item[0]))
|
||||
|
||||
# blocks of text on page
|
||||
blocks = page.get_text(
|
||||
"dict",
|
||||
flags=fitz.TEXTFLAGS_TEXT,
|
||||
clip=clip,
|
||||
)["blocks"]
|
||||
|
||||
# Make block rectangles, ignoring non-horizontal text
|
||||
for b in blocks:
|
||||
bbox = fitz.IRect(b["bbox"]) # bbox of the block
|
||||
|
||||
# ignore text written upon images
|
||||
if no_image_text and in_bbox(bbox, img_bboxes):
|
||||
continue
|
||||
|
||||
# confirm first line to be horizontal
|
||||
line0 = b["lines"][0] # get first line
|
||||
if line0["dir"] != (1, 0): # only accept horizontal text
|
||||
vert_bboxes.append(bbox)
|
||||
continue
|
||||
|
||||
srect = fitz.EMPTY_IRECT()
|
||||
for line in b["lines"]:
|
||||
lbbox = fitz.IRect(line["bbox"])
|
||||
text = "".join([s["text"].strip() for s in line["spans"]])
|
||||
if len(text) > 1:
|
||||
srect |= lbbox
|
||||
bbox = +srect
|
||||
|
||||
if not bbox.is_empty:
|
||||
bboxes.append(bbox)
|
||||
|
||||
# Sort text bboxes by ascending background, top, then left coordinates
|
||||
bboxes.sort(key=lambda k: (in_bbox(k, path_bboxes), k.y0, k.x0))
|
||||
|
||||
# Extend bboxes to the right where possible
|
||||
bboxes = extend_right(
|
||||
bboxes, int(page.rect.width), path_bboxes, vert_bboxes, img_bboxes
|
||||
)
|
||||
|
||||
# immediately return of no text found
|
||||
if bboxes == []:
|
||||
return []
|
||||
|
||||
# --------------------------------------------------------------------
|
||||
# Join bboxes to establish some column structure
|
||||
# --------------------------------------------------------------------
|
||||
# the final block bboxes on page
|
||||
nblocks = [bboxes[0]] # pre-fill with first bbox
|
||||
bboxes = bboxes[1:] # remaining old bboxes
|
||||
|
||||
for i, bb in enumerate(bboxes): # iterate old bboxes
|
||||
check = False # indicates unwanted joins
|
||||
|
||||
# check if bb can extend one of the new blocks
|
||||
for j in range(len(nblocks)):
|
||||
nbb = nblocks[j] # a new block
|
||||
|
||||
# never join across columns
|
||||
if bb == None or nbb.x1 < bb.x0 or bb.x1 < nbb.x0:
|
||||
continue
|
||||
|
||||
# never join across different background colors
|
||||
if in_bbox(nbb, path_bboxes) != in_bbox(bb, path_bboxes):
|
||||
continue
|
||||
|
||||
temp = bb | nbb # temporary extension of new block
|
||||
check = can_extend(temp, nbb, nblocks)
|
||||
if check == True:
|
||||
break
|
||||
|
||||
if not check: # bb cannot be used to extend any of the new bboxes
|
||||
nblocks.append(bb) # so add it to the list
|
||||
j = len(nblocks) - 1 # index of it
|
||||
temp = nblocks[j] # new bbox added
|
||||
|
||||
# check if some remaining bbox is contained in temp
|
||||
check = can_extend(temp, bb, bboxes)
|
||||
if check == False:
|
||||
nblocks.append(bb)
|
||||
else:
|
||||
nblocks[j] = temp
|
||||
bboxes[i] = None
|
||||
|
||||
# do some elementary cleaning
|
||||
nblocks = clean_nblocks(nblocks)
|
||||
|
||||
# return identified text bboxes
|
||||
return nblocks
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
"""Only for debugging purposes, currently.
|
||||
|
||||
Draw red borders around the returned text bboxes and insert
|
||||
the bbox number.
|
||||
Then save the file under the name "input-blocks.pdf".
|
||||
"""
|
||||
|
||||
# get the file name
|
||||
filename = sys.argv[1]
|
||||
|
||||
# check if footer margin is given
|
||||
if len(sys.argv) > 2:
|
||||
footer_margin = int(sys.argv[2])
|
||||
else: # use default vaue
|
||||
footer_margin = 50
|
||||
|
||||
# check if header margin is given
|
||||
if len(sys.argv) > 3:
|
||||
header_margin = int(sys.argv[3])
|
||||
else: # use default vaue
|
||||
header_margin = 50
|
||||
|
||||
# open document
|
||||
doc = fitz.open(filename)
|
||||
|
||||
# iterate over the pages
|
||||
for page in doc:
|
||||
# remove any geometry issues
|
||||
page.wrap_contents()
|
||||
|
||||
# get the text bboxes
|
||||
bboxes = column_boxes(page, footer_margin=footer_margin, header_margin=header_margin)
|
||||
|
||||
# prepare a canvas to draw rectangles and text
|
||||
shape = page.new_shape()
|
||||
|
||||
# iterate over the bboxes
|
||||
for i, rect in enumerate(bboxes):
|
||||
shape.draw_rect(rect) # draw a border
|
||||
|
||||
# write sequence number
|
||||
shape.insert_text(rect.tl + (5, 15), str(i), color=fitz.pdfcolor["red"])
|
||||
|
||||
# finish drawing / text with color red
|
||||
shape.finish(color=fitz.pdfcolor["red"])
|
||||
shape.commit() # store to the page
|
||||
|
||||
# save document with text bboxes
|
||||
doc.ez_save(filename.replace(".pdf", "-blocks.pdf"))
|
||||
@@ -1,239 +0,0 @@
|
||||
import os
|
||||
import csv
|
||||
import json
|
||||
import pandas as pd
|
||||
from pandas import DataFrame as df
|
||||
from matplotlib import pyplot as plt
|
||||
from termcolor import cprint
|
||||
|
||||
"""
|
||||
Execute this script in the following way:
|
||||
|
||||
1. Make sure there are pdf_dic.json files under the directory code-clean/tmp/unittest/md/, such as the following:
|
||||
|
||||
code-clean/tmp/unittest/md/scihub/scihub_00500000/libgen.scimag00527000-00527999.zip_10.1002/app.25178/pdf_dic.json
|
||||
|
||||
2. Under the directory code-clean, execute the following command:
|
||||
|
||||
$ python -m libs.calc_span_stats
|
||||
|
||||
"""
|
||||
|
||||
|
||||
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 safe_get(dict_obj, key, default):
|
||||
val = dict_obj.get(key)
|
||||
if val is None:
|
||||
return default
|
||||
else:
|
||||
return val
|
||||
|
||||
|
||||
class SpanStatsCalc:
|
||||
"""Calculate statistics of span."""
|
||||
|
||||
def draw_charts(self, span_stats: pd.DataFrame, fig_num: int, save_path: str):
|
||||
"""Draw multiple figures in one figure."""
|
||||
# make a canvas
|
||||
fig = plt.figure(fig_num, figsize=(20, 20))
|
||||
|
||||
pass
|
||||
|
||||
def calc_stats_per_dict(self, pdf_dict) -> pd.DataFrame:
|
||||
"""Calculate statistics per pdf_dict."""
|
||||
span_stats = pd.DataFrame()
|
||||
|
||||
span_stats = []
|
||||
span_id = 0
|
||||
for page_id, blocks in pdf_dict.items():
|
||||
if page_id.startswith("page_"):
|
||||
if "para_blocks" in blocks.keys():
|
||||
for para_block in blocks["para_blocks"]:
|
||||
for line in para_block["lines"]:
|
||||
for span in line["spans"]:
|
||||
span_text = safe_get(span, "text", "")
|
||||
span_font_name = safe_get(span, "font", "")
|
||||
span_font_size = safe_get(span, "size", 0)
|
||||
span_font_color = safe_get(span, "color", "")
|
||||
span_font_flags = safe_get(span, "flags", 0)
|
||||
|
||||
span_font_flags_decoded = safe_get(span, "decomposed_flags", {})
|
||||
span_is_super_script = safe_get(span_font_flags_decoded, "is_superscript", False)
|
||||
span_is_italic = safe_get(span_font_flags_decoded, "is_italic", False)
|
||||
span_is_serifed = safe_get(span_font_flags_decoded, "is_serifed", False)
|
||||
span_is_sans_serifed = safe_get(span_font_flags_decoded, "is_sans_serifed", False)
|
||||
span_is_monospaced = safe_get(span_font_flags_decoded, "is_monospaced", False)
|
||||
span_is_proportional = safe_get(span_font_flags_decoded, "is_proportional", False)
|
||||
span_is_bold = safe_get(span_font_flags_decoded, "is_bold", False)
|
||||
|
||||
span_stats.append(
|
||||
{
|
||||
"span_id": span_id, # id of span
|
||||
"page_id": page_id, # page number of pdf
|
||||
"span_text": span_text, # text of span
|
||||
"span_font_name": span_font_name, # font name of span
|
||||
"span_font_size": span_font_size, # font size of span
|
||||
"span_font_color": span_font_color, # font color of span
|
||||
"span_font_flags": span_font_flags, # font flags of span
|
||||
"span_is_superscript": int(
|
||||
span_is_super_script
|
||||
), # indicate whether the span is super script or not
|
||||
"span_is_italic": int(span_is_italic), # indicate whether the span is italic or not
|
||||
"span_is_serifed": int(span_is_serifed), # indicate whether the span is serifed or not
|
||||
"span_is_sans_serifed": int(
|
||||
span_is_sans_serifed
|
||||
), # indicate whether the span is sans serifed or not
|
||||
"span_is_monospaced": int(
|
||||
span_is_monospaced
|
||||
), # indicate whether the span is monospaced or not
|
||||
"span_is_proportional": int(
|
||||
span_is_proportional
|
||||
), # indicate whether the span is proportional or not
|
||||
"span_is_bold": int(span_is_bold), # indicate whether the span is bold or not
|
||||
}
|
||||
)
|
||||
|
||||
span_id += 1
|
||||
|
||||
span_stats = pd.DataFrame(span_stats)
|
||||
# print(span_stats)
|
||||
|
||||
return span_stats
|
||||
|
||||
|
||||
def __find_pdf_dic_files(
|
||||
jf_name="pdf_dic.json",
|
||||
base_code_name="code-clean",
|
||||
tgt_base_dir_name="tmp",
|
||||
unittest_dir_name="unittest",
|
||||
md_dir_name="md",
|
||||
book_names=[
|
||||
"scihub",
|
||||
], # other possible values: "zlib", "arxiv" and so on
|
||||
):
|
||||
pdf_dict_files = []
|
||||
|
||||
curr_dir = os.path.dirname(__file__)
|
||||
|
||||
for i in range(len(curr_dir)):
|
||||
if curr_dir[i : i + len(base_code_name)] == base_code_name:
|
||||
base_code_dir_name = curr_dir[: i + len(base_code_name)]
|
||||
for book_name in book_names:
|
||||
search_dir_relative_name = os.path.join(tgt_base_dir_name, unittest_dir_name, md_dir_name, book_name)
|
||||
if os.path.exists(base_code_dir_name):
|
||||
search_dir_name = os.path.join(base_code_dir_name, search_dir_relative_name)
|
||||
for root, dirs, files in os.walk(search_dir_name):
|
||||
for file in files:
|
||||
if file == jf_name:
|
||||
pdf_dict_files.append(os.path.join(root, file))
|
||||
break
|
||||
|
||||
return pdf_dict_files
|
||||
|
||||
|
||||
def combine_span_texts(group_df, span_stats):
|
||||
combined_span_texts = []
|
||||
for _, row in group_df.iterrows():
|
||||
curr_span_id = row.name
|
||||
curr_span_text = row["span_text"]
|
||||
|
||||
pre_span_id = curr_span_id - 1
|
||||
pre_span_text = span_stats.at[pre_span_id, "span_text"] if pre_span_id in span_stats.index else ""
|
||||
|
||||
next_span_id = curr_span_id + 1
|
||||
next_span_text = span_stats.at[next_span_id, "span_text"] if next_span_id in span_stats.index else ""
|
||||
|
||||
# pointer_sign is a right arrow if the span is superscript, otherwise it is a down arrow
|
||||
pointer_sign = "→ → → "
|
||||
combined_text = "\n".join([pointer_sign + pre_span_text, pointer_sign + curr_span_text, pointer_sign + next_span_text])
|
||||
combined_span_texts.append(combined_text)
|
||||
|
||||
return "\n\n".join(combined_span_texts)
|
||||
|
||||
|
||||
# pd.set_option("display.max_colwidth", None) # 设置为 None 来显示完整的文本
|
||||
pd.set_option("display.max_rows", None) # 设置为 None 来显示更多的行
|
||||
|
||||
|
||||
def main():
|
||||
pdf_dict_files = __find_pdf_dic_files()
|
||||
# print(pdf_dict_files)
|
||||
|
||||
span_stats_calc = SpanStatsCalc()
|
||||
|
||||
for pdf_dict_file in pdf_dict_files:
|
||||
print("-" * 100)
|
||||
print_green_on_red(f"Processing {pdf_dict_file}")
|
||||
|
||||
with open(pdf_dict_file, "r", encoding="utf-8") as f:
|
||||
pdf_dict = json.load(f)
|
||||
|
||||
raw_df = span_stats_calc.calc_stats_per_dict(pdf_dict)
|
||||
save_path = pdf_dict_file.replace("pdf_dic.json", "span_stats_raw.csv")
|
||||
raw_df.to_csv(save_path, index=False)
|
||||
|
||||
filtered_df = raw_df[raw_df["span_is_superscript"] == 1]
|
||||
if filtered_df.empty:
|
||||
print("No superscript span found!")
|
||||
continue
|
||||
|
||||
filtered_grouped_df = filtered_df.groupby(["span_font_name", "span_font_size", "span_font_color"])
|
||||
|
||||
combined_span_texts = filtered_grouped_df.apply(combine_span_texts, span_stats=raw_df) # type: ignore
|
||||
|
||||
final_df = filtered_grouped_df.size().reset_index(name="count")
|
||||
final_df["span_texts"] = combined_span_texts.reset_index(level=[0, 1, 2], drop=True)
|
||||
|
||||
print(final_df)
|
||||
|
||||
final_df["span_texts"] = final_df["span_texts"].apply(lambda x: x.replace("\n", "\r\n"))
|
||||
|
||||
save_path = pdf_dict_file.replace("pdf_dic.json", "span_stats_final.csv")
|
||||
# 使用 UTF-8 编码并添加 BOM，确保所有字段被双引号包围
|
||||
final_df.to_csv(save_path, index=False, encoding="utf-8-sig", quoting=csv.QUOTE_ALL)
|
||||
|
||||
# 创建一个 2x2 的图表布局
|
||||
fig, axs = plt.subplots(2, 2, figsize=(15, 10))
|
||||
|
||||
# 按照 span_font_name 分类作图
|
||||
final_df.groupby("span_font_name")["count"].sum().plot(kind="bar", ax=axs[0, 0], title="By Font Name")
|
||||
|
||||
# 按照 span_font_size 分类作图
|
||||
final_df.groupby("span_font_size")["count"].sum().plot(kind="bar", ax=axs[0, 1], title="By Font Size")
|
||||
|
||||
# 按照 span_font_color 分类作图
|
||||
final_df.groupby("span_font_color")["count"].sum().plot(kind="bar", ax=axs[1, 0], title="By Font Color")
|
||||
|
||||
# 按照 span_font_name、span_font_size 和 span_font_color 共同分类作图
|
||||
grouped = final_df.groupby(["span_font_name", "span_font_size", "span_font_color"])
|
||||
grouped["count"].sum().unstack().plot(kind="bar", ax=axs[1, 1], title="Combined Grouping")
|
||||
|
||||
# 调整布局
|
||||
plt.tight_layout()
|
||||
|
||||
# 显示图表
|
||||
# plt.show()
|
||||
|
||||
# 保存图表到 PNG 文件
|
||||
save_path = pdf_dict_file.replace("pdf_dic.json", "span_stats_combined.png")
|
||||
plt.savefig(save_path)
|
||||
|
||||
# 清除画布
|
||||
plt.clf()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,21 +0,0 @@
|
||||
from collections import Counter
|
||||
|
||||
from magic_pdf.libs.language import detect_lang
|
||||
|
||||
def get_language_from_model(model_list: list):
|
||||
language_lst = []
|
||||
for ocr_page_info in model_list:
|
||||
page_text = ""
|
||||
layout_dets = ocr_page_info["layout_dets"]
|
||||
for layout_det in layout_dets:
|
||||
category_id = layout_det["category_id"]
|
||||
allow_category_id_list = [15]
|
||||
if category_id in allow_category_id_list:
|
||||
page_text += layout_det["text"]
|
||||
page_language = detect_lang(page_text)
|
||||
language_lst.append(page_language)
|
||||
# 统计text_language_list中每种语言的个数
|
||||
count_dict = Counter(language_lst)
|
||||
# 输出text_language_list中出现的次数最多的语言
|
||||
language = max(count_dict, key=count_dict.get)
|
||||
return language
|
||||
@@ -1,203 +0,0 @@
|
||||
import re
|
||||
from os import path
|
||||
|
||||
from collections import Counter
|
||||
|
||||
from loguru import logger
|
||||
|
||||
# from langdetect import detect
|
||||
import spacy
|
||||
import en_core_web_sm
|
||||
import zh_core_web_sm
|
||||
|
||||
from magic_pdf.libs.language import detect_lang
|
||||
|
||||
|
||||
class NLPModels:
|
||||
"""
|
||||
How to upload local models to s3:
|
||||
- config aws cli:
|
||||
doc\SETUP-CLI.md
|
||||
doc\setup_cli.sh
|
||||
app\config\__init__.py
|
||||
- $ cd {local_dir_storing_models}
|
||||
- $ ls models
|
||||
en_core_web_sm-3.7.1/
|
||||
zh_core_web_sm-3.7.0/
|
||||
- $ aws s3 sync models/ s3://llm-infra/models --profile=p_project_norm
|
||||
- $ aws s3 --profile=p_project_norm ls s3://llm-infra/models/
|
||||
PRE en_core_web_sm-3.7.1/
|
||||
PRE zh_core_web_sm-3.7.0/
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
# if OS is windows, set "TMP_DIR" to "D:/tmp"
|
||||
|
||||
home_dir = path.expanduser("~")
|
||||
self.default_local_path = path.join(home_dir, ".nlp_models")
|
||||
self.default_shared_path = "/share/pdf_processor/nlp_models"
|
||||
self.default_hdfs_path = "hdfs://pdf_processor/nlp_models"
|
||||
self.default_s3_path = "s3://llm-infra/models"
|
||||
self.nlp_models = self.nlp_models = {
|
||||
"en_core_web_sm": {
|
||||
"type": "spacy",
|
||||
"version": "3.7.1",
|
||||
},
|
||||
"en_core_web_md": {
|
||||
"type": "spacy",
|
||||
"version": "3.7.1",
|
||||
},
|
||||
"en_core_web_lg": {
|
||||
"type": "spacy",
|
||||
"version": "3.7.1",
|
||||
},
|
||||
"zh_core_web_sm": {
|
||||
"type": "spacy",
|
||||
"version": "3.7.0",
|
||||
},
|
||||
"zh_core_web_md": {
|
||||
"type": "spacy",
|
||||
"version": "3.7.0",
|
||||
},
|
||||
"zh_core_web_lg": {
|
||||
"type": "spacy",
|
||||
"version": "3.7.0",
|
||||
},
|
||||
}
|
||||
self.en_core_web_sm_model = en_core_web_sm.load()
|
||||
self.zh_core_web_sm_model = zh_core_web_sm.load()
|
||||
|
||||
def load_model(self, model_name, model_type, model_version):
|
||||
if (
|
||||
model_name in self.nlp_models
|
||||
and self.nlp_models[model_name]["type"] == model_type
|
||||
and self.nlp_models[model_name]["version"] == model_version
|
||||
):
|
||||
return spacy.load(model_name) if spacy.util.is_package(model_name) else None
|
||||
|
||||
else:
|
||||
logger.error(f"Unsupported model name or version: {model_name} {model_version}")
|
||||
return None
|
||||
|
||||
def detect_language(self, text, use_langdetect=False):
|
||||
if len(text) == 0:
|
||||
return None
|
||||
if use_langdetect:
|
||||
# print("use_langdetect")
|
||||
# print(detect_lang(text))
|
||||
# return detect_lang(text)
|
||||
if detect_lang(text) == "zh":
|
||||
return "zh"
|
||||
else:
|
||||
return "en"
|
||||
|
||||
if not use_langdetect:
|
||||
en_count = len(re.findall(r"[a-zA-Z]", text))
|
||||
cn_count = len(re.findall(r"[\u4e00-\u9fff]", text))
|
||||
|
||||
if en_count > cn_count:
|
||||
return "en"
|
||||
|
||||
if cn_count > en_count:
|
||||
return "zh"
|
||||
|
||||
def detect_entity_catgr_using_nlp(self, text, threshold=0.5):
|
||||
"""
|
||||
Detect entity categories using NLP models and return the most frequent entity types.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
text : str
|
||||
Text to be processed.
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
The most frequent entity type.
|
||||
"""
|
||||
lang = self.detect_language(text, use_langdetect=True)
|
||||
|
||||
if lang == "en":
|
||||
nlp_model = self.en_core_web_sm_model
|
||||
elif lang == "zh":
|
||||
nlp_model = self.zh_core_web_sm_model
|
||||
else:
|
||||
# logger.error(f"Unsupported language: {lang}")
|
||||
return {}
|
||||
|
||||
# Splitting text into smaller parts
|
||||
text_parts = re.split(r"[,;，；、\s & |]+", text)
|
||||
|
||||
text_parts = [part for part in text_parts if not re.match(r"[\d\W]+", part)] # Remove non-words
|
||||
text_combined = " ".join(text_parts)
|
||||
|
||||
try:
|
||||
doc = nlp_model(text_combined)
|
||||
entity_counts = Counter([ent.label_ for ent in doc.ents])
|
||||
word_counts_in_entities = Counter()
|
||||
|
||||
for ent in doc.ents:
|
||||
word_counts_in_entities[ent.label_] += len(ent.text.split())
|
||||
|
||||
total_words_in_entities = sum(word_counts_in_entities.values())
|
||||
total_words = len([token for token in doc if not token.is_punct])
|
||||
|
||||
if total_words_in_entities == 0 or total_words == 0:
|
||||
return None
|
||||
|
||||
entity_percentage = total_words_in_entities / total_words
|
||||
if entity_percentage < 0.5:
|
||||
return None
|
||||
|
||||
most_common_entity, word_count = word_counts_in_entities.most_common(1)[0]
|
||||
entity_percentage = word_count / total_words_in_entities
|
||||
|
||||
if entity_percentage >= threshold:
|
||||
return most_common_entity
|
||||
else:
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"Error in entity detection: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def __main__():
|
||||
nlpModel = NLPModels()
|
||||
|
||||
test_strings = [
|
||||
"张三",
|
||||
"张三, 李四，王五; 赵六",
|
||||
"John Doe",
|
||||
"Jane Smith",
|
||||
"Lee, John",
|
||||
"John Doe, Jane Smith; Alice Johnson，Bob Lee",
|
||||
"孙七, Michael Jordan；赵八",
|
||||
"David Smith Michael O'Connor; Kevin ßáçøñ",
|
||||
"李雷·韩梅梅, 张三·李四",
|
||||
"Charles Robert Darwin, Isaac Newton",
|
||||
"莱昂纳多·迪卡普里奥, 杰克·吉伦哈尔",
|
||||
"John Doe, Jane Smith; Alice Johnson",
|
||||
"张三, 李四，王五; 赵六",
|
||||
"Lei Wang, Jia Li, and Xiaojun Chen, LINKE YANG OU, and YUAN ZHANG",
|
||||
"Rachel Mills & William Barry & Susanne B. Haga",
|
||||
"Claire Chabut* and Jean-François Bussières",
|
||||
"1 Department of Chemistry, Northeastern University, Shenyang 110004, China 2 State Key Laboratory of Polymer Physics and Chemistry, Changchun Institute of Applied Chemistry, Chinese Academy of Sciences, Changchun 130022, China",
|
||||
"Changchun",
|
||||
"china",
|
||||
"Rongjun Song, 1,2 Baoyan Zhang, 1 Baotong Huang, 2 Tao Tang 2",
|
||||
"Synergistic Effect of Supported Nickel Catalyst with Intumescent Flame-Retardants on Flame Retardancy and Thermal Stability of Polypropylene",
|
||||
"Synergistic Effect of Supported Nickel Catalyst with",
|
||||
"Intumescent Flame-Retardants on Flame Retardancy",
|
||||
"and Thermal Stability of Polypropylene",
|
||||
]
|
||||
|
||||
for test in test_strings:
|
||||
print()
|
||||
print(f"Original String: {test}")
|
||||
|
||||
result = nlpModel.detect_entity_catgr_using_nlp(test)
|
||||
print(f"Detected entities: {result}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
__main__()
|
||||
@@ -1,33 +0,0 @@
|
||||
import math
|
||||
|
||||
|
||||
def __inc_dict_val(mp, key, val_inc:int):
|
||||
if mp.get(key):
|
||||
mp[key] = mp[key] + val_inc
|
||||
else:
|
||||
mp[key] = val_inc
|
||||
|
||||
|
||||
|
||||
def get_text_block_base_info(block):
|
||||
"""
|
||||
获取这个文本块里的字体的颜色、字号、字体
|
||||
按照正文字数最多的返回
|
||||
"""
|
||||
|
||||
counter = {}
|
||||
|
||||
for line in block['lines']:
|
||||
for span in line['spans']:
|
||||
color = span['color']
|
||||
size = round(span['size'], 2)
|
||||
font = span['font']
|
||||
|
||||
txt_len = len(span['text'])
|
||||
__inc_dict_val(counter, (color, size, font), txt_len)
|
||||
|
||||
|
||||
c, s, ft = max(counter, key=counter.get)
|
||||
|
||||
return c, s, ft
|
||||
|
||||
@@ -1,308 +0,0 @@
|
||||
from magic_pdf.libs.commons import fitz
|
||||
import os
|
||||
|
||||
|
||||
def draw_bbox_on_page(raw_pdf_doc: fitz.Document, paras_dict:dict, save_path: str):
|
||||
"""
|
||||
在page上画出bbox，保存到save_path
|
||||
"""
|
||||
# 检查文件是否存在
|
||||
is_new_pdf = False
|
||||
if os.path.exists(save_path):
|
||||
# 打开现有的 PDF 文件
|
||||
doc = fitz.open(save_path)
|
||||
else:
|
||||
# 创建一个新的空白 PDF 文件
|
||||
is_new_pdf = True
|
||||
doc = fitz.open('')
|
||||
|
||||
color_map = {
|
||||
'image': fitz.pdfcolor["yellow"],
|
||||
'text': fitz.pdfcolor['blue'],
|
||||
"table": fitz.pdfcolor['green']
|
||||
}
|
||||
|
||||
for k, v in paras_dict.items():
|
||||
page_idx = v['page_idx']
|
||||
width = raw_pdf_doc[page_idx].rect.width
|
||||
height = raw_pdf_doc[page_idx].rect.height
|
||||
new_page = doc.new_page(width=width, height=height)
|
||||
|
||||
shape = new_page.new_shape()
|
||||
for order, block in enumerate(v['preproc_blocks']):
|
||||
rect = fitz.Rect(block['bbox'])
|
||||
shape = new_page.new_shape()
|
||||
shape.draw_rect(rect)
|
||||
shape.finish(color=None, fill=color_map['text'], fill_opacity=0.2)
|
||||
shape.finish()
|
||||
shape.commit()
|
||||
|
||||
for img in v['images']:
|
||||
# 原始box画上去
|
||||
rect = fitz.Rect(img['bbox'])
|
||||
shape = new_page.new_shape()
|
||||
shape.draw_rect(rect)
|
||||
shape.finish(color=None, fill=fitz.pdfcolor['yellow'])
|
||||
shape.finish()
|
||||
shape.commit()
|
||||
|
||||
for img in v['image_backup']:
|
||||
# 原始box画上去
|
||||
rect = fitz.Rect(img['bbox'])
|
||||
shape = new_page.new_shape()
|
||||
shape.draw_rect(rect)
|
||||
shape.finish(color=fitz.pdfcolor['yellow'], fill=None)
|
||||
shape.finish()
|
||||
shape.commit()
|
||||
|
||||
for tb in v['droped_text_block']:
|
||||
# 原始box画上去
|
||||
rect = fitz.Rect(tb['bbox'])
|
||||
shape = new_page.new_shape()
|
||||
shape.draw_rect(rect)
|
||||
shape.finish(color=None, fill=fitz.pdfcolor['black'], fill_opacity=0.4)
|
||||
shape.finish()
|
||||
shape.commit()
|
||||
|
||||
# TODO table
|
||||
for tb in v['tables']:
|
||||
rect = fitz.Rect(tb['bbox'])
|
||||
shape = new_page.new_shape()
|
||||
shape.draw_rect(rect)
|
||||
shape.finish(color=None, fill=fitz.pdfcolor['green'], fill_opacity=0.2)
|
||||
shape.finish()
|
||||
shape.commit()
|
||||
|
||||
|
||||
parent_dir = os.path.dirname(save_path)
|
||||
if not os.path.exists(parent_dir):
|
||||
os.makedirs(parent_dir)
|
||||
|
||||
if is_new_pdf:
|
||||
doc.save(save_path)
|
||||
else:
|
||||
doc.saveIncr()
|
||||
doc.close()
|
||||
|
||||
|
||||
def debug_show_bbox(raw_pdf_doc: fitz.Document, page_idx: int, bboxes: list, droped_bboxes:list, expect_drop_bboxes:list, save_path: str, expected_page_id:int):
|
||||
"""
|
||||
以覆盖的方式写个临时的pdf，用于debug
|
||||
"""
|
||||
if page_idx!=expected_page_id:
|
||||
return
|
||||
|
||||
if os.path.exists(save_path):
|
||||
# 删除已经存在的文件
|
||||
os.remove(save_path)
|
||||
# 创建一个新的空白 PDF 文件
|
||||
doc = fitz.open('')
|
||||
|
||||
width = raw_pdf_doc[page_idx].rect.width
|
||||
height = raw_pdf_doc[page_idx].rect.height
|
||||
new_page = doc.new_page(width=width, height=height)
|
||||
|
||||
shape = new_page.new_shape()
|
||||
for bbox in bboxes:
|
||||
# 原始box画上去
|
||||
rect = fitz.Rect(*bbox[0:4])
|
||||
shape = new_page.new_shape()
|
||||
shape.draw_rect(rect)
|
||||
shape.finish(color=fitz.pdfcolor['red'], fill=fitz.pdfcolor['blue'], fill_opacity=0.2)
|
||||
shape.finish()
|
||||
shape.commit()
|
||||
|
||||
for bbox in droped_bboxes:
|
||||
# 原始box画上去
|
||||
rect = fitz.Rect(*bbox[0:4])
|
||||
shape = new_page.new_shape()
|
||||
shape.draw_rect(rect)
|
||||
shape.finish(color=None, fill=fitz.pdfcolor['yellow'], fill_opacity=0.2)
|
||||
shape.finish()
|
||||
shape.commit()
|
||||
|
||||
for bbox in expect_drop_bboxes:
|
||||
# 原始box画上去
|
||||
rect = fitz.Rect(*bbox[0:4])
|
||||
shape = new_page.new_shape()
|
||||
shape.draw_rect(rect)
|
||||
shape.finish(color=fitz.pdfcolor['red'], fill=None)
|
||||
shape.finish()
|
||||
shape.commit()
|
||||
|
||||
# shape.insert_textbox(fitz.Rect(200, 0, 600, 20), f"total bboxes: {len(bboxes)}", fontname="helv", fontsize=12,
|
||||
# color=(0, 0, 0))
|
||||
# shape.finish(color=fitz.pdfcolor['black'])
|
||||
# shape.commit()
|
||||
|
||||
parent_dir = os.path.dirname(save_path)
|
||||
if not os.path.exists(parent_dir):
|
||||
os.makedirs(parent_dir)
|
||||
|
||||
doc.save(save_path)
|
||||
doc.close()
|
||||
|
||||
|
||||
def debug_show_page(page, bboxes1: list,bboxes2: list,bboxes3: list,):
|
||||
save_path = "./tmp/debug.pdf"
|
||||
if os.path.exists(save_path):
|
||||
# 删除已经存在的文件
|
||||
os.remove(save_path)
|
||||
# 创建一个新的空白 PDF 文件
|
||||
doc = fitz.open('')
|
||||
|
||||
width = page.rect.width
|
||||
height = page.rect.height
|
||||
new_page = doc.new_page(width=width, height=height)
|
||||
|
||||
shape = new_page.new_shape()
|
||||
for bbox in bboxes1:
|
||||
# 原始box画上去
|
||||
rect = fitz.Rect(*bbox[0:4])
|
||||
shape = new_page.new_shape()
|
||||
shape.draw_rect(rect)
|
||||
shape.finish(color=fitz.pdfcolor['red'], fill=fitz.pdfcolor['blue'], fill_opacity=0.2)
|
||||
shape.finish()
|
||||
shape.commit()
|
||||
|
||||
for bbox in bboxes2:
|
||||
# 原始box画上去
|
||||
rect = fitz.Rect(*bbox[0:4])
|
||||
shape = new_page.new_shape()
|
||||
shape.draw_rect(rect)
|
||||
shape.finish(color=None, fill=fitz.pdfcolor['yellow'], fill_opacity=0.2)
|
||||
shape.finish()
|
||||
shape.commit()
|
||||
|
||||
for bbox in bboxes3:
|
||||
# 原始box画上去
|
||||
rect = fitz.Rect(*bbox[0:4])
|
||||
shape = new_page.new_shape()
|
||||
shape.draw_rect(rect)
|
||||
shape.finish(color=fitz.pdfcolor['red'], fill=None)
|
||||
shape.finish()
|
||||
shape.commit()
|
||||
|
||||
parent_dir = os.path.dirname(save_path)
|
||||
if not os.path.exists(parent_dir):
|
||||
os.makedirs(parent_dir)
|
||||
|
||||
doc.save(save_path)
|
||||
doc.close()
|
||||
|
||||
|
||||
|
||||
|
||||
def draw_layout_bbox_on_page(raw_pdf_doc: fitz.Document, paras_dict:dict, header, footer, pdf_path: str):
|
||||
"""
|
||||
在page上画出bbox，保存到save_path
|
||||
"""
|
||||
# 检查文件是否存在
|
||||
is_new_pdf = False
|
||||
if os.path.exists(pdf_path):
|
||||
# 打开现有的 PDF 文件
|
||||
doc = fitz.open(pdf_path)
|
||||
else:
|
||||
# 创建一个新的空白 PDF 文件
|
||||
is_new_pdf = True
|
||||
doc = fitz.open('')
|
||||
|
||||
for k, v in paras_dict.items():
|
||||
page_idx = v['page_idx']
|
||||
layouts = v['layout_bboxes']
|
||||
page = doc[page_idx]
|
||||
shape = page.new_shape()
|
||||
for order, layout in enumerate(layouts):
|
||||
border_offset = 1
|
||||
rect_box = layout['layout_bbox']
|
||||
layout_label = layout['layout_label']
|
||||
fill_color = fitz.pdfcolor['pink'] if layout_label=='U' else None
|
||||
rect_box = [rect_box[0]+1, rect_box[1]-border_offset, rect_box[2]-1, rect_box[3]+border_offset]
|
||||
rect = fitz.Rect(*rect_box)
|
||||
shape.draw_rect(rect)
|
||||
shape.finish(color=fitz.pdfcolor['red'], fill=fill_color, fill_opacity=0.4)
|
||||
"""
|
||||
draw order text on layout box
|
||||
"""
|
||||
font_size = 10
|
||||
shape.insert_text((rect_box[0] + 1, rect_box[1] + font_size), f"{order}", fontsize=font_size, color=(0, 0, 0))
|
||||
|
||||
"""画上footer header"""
|
||||
if header:
|
||||
shape.draw_rect(fitz.Rect(header))
|
||||
shape.finish(color=None, fill=fitz.pdfcolor['black'], fill_opacity=0.2)
|
||||
if footer:
|
||||
shape.draw_rect(fitz.Rect(footer))
|
||||
shape.finish(color=None, fill=fitz.pdfcolor['black'], fill_opacity=0.2)
|
||||
|
||||
shape.commit()
|
||||
|
||||
if is_new_pdf:
|
||||
doc.save(pdf_path)
|
||||
else:
|
||||
doc.saveIncr()
|
||||
doc.close()
|
||||
|
||||
|
||||
@DeprecationWarning
|
||||
def draw_layout_on_page(raw_pdf_doc: fitz.Document, page_idx: int, page_layout: list, pdf_path: str):
|
||||
"""
|
||||
把layout的box用红色边框花在pdf_path的page_idx上
|
||||
"""
|
||||
def draw(shape, layout, fill_color=fitz.pdfcolor['pink']):
|
||||
border_offset = 1
|
||||
rect_box = layout['layout_bbox']
|
||||
layout_label = layout['layout_label']
|
||||
sub_layout = layout['sub_layout']
|
||||
if len(sub_layout)==0:
|
||||
fill_color = fill_color if layout_label=='U' else None
|
||||
rect_box = [rect_box[0]+1, rect_box[1]-border_offset, rect_box[2]-1, rect_box[3]+border_offset]
|
||||
rect = fitz.Rect(*rect_box)
|
||||
shape.draw_rect(rect)
|
||||
shape.finish(color=fitz.pdfcolor['red'], fill=fill_color, fill_opacity=0.2)
|
||||
# if layout_label=='U':
|
||||
# bad_boxes = layout.get("bad_boxes", [])
|
||||
# for bad_box in bad_boxes:
|
||||
# rect = fitz.Rect(*bad_box)
|
||||
# shape.draw_rect(rect)
|
||||
# shape.finish(color=fitz.pdfcolor['red'], fill=fitz.pdfcolor['red'], fill_opacity=0.2)
|
||||
# else:
|
||||
# rect = fitz.Rect(*rect_box)
|
||||
# shape.draw_rect(rect)
|
||||
# shape.finish(color=fitz.pdfcolor['blue'])
|
||||
|
||||
for sub_layout in sub_layout:
|
||||
draw(shape, sub_layout)
|
||||
shape.commit()
|
||||
|
||||
|
||||
# 检查文件是否存在
|
||||
is_new_pdf = False
|
||||
if os.path.exists(pdf_path):
|
||||
# 打开现有的 PDF 文件
|
||||
doc = fitz.open(pdf_path)
|
||||
else:
|
||||
# 创建一个新的空白 PDF 文件
|
||||
is_new_pdf = True
|
||||
doc = fitz.open('')
|
||||
|
||||
page = doc[page_idx]
|
||||
shape = page.new_shape()
|
||||
for order, layout in enumerate(page_layout):
|
||||
draw(shape, layout, fitz.pdfcolor['yellow'])
|
||||
|
||||
# shape.insert_textbox(fitz.Rect(200, 0, 600, 20), f"total bboxes: {len(layout)}", fontname="helv", fontsize=12,
|
||||
# color=(0, 0, 0))
|
||||
# shape.finish(color=fitz.pdfcolor['black'])
|
||||
# shape.commit()
|
||||
|
||||
parent_dir = os.path.dirname(pdf_path)
|
||||
if not os.path.exists(parent_dir):
|
||||
os.makedirs(parent_dir)
|
||||
|
||||
if is_new_pdf:
|
||||
doc.save(pdf_path)
|
||||
else:
|
||||
doc.saveIncr()
|
||||
doc.close()
|
||||
|
||||
@@ -1,562 +0,0 @@
|
||||
import os
|
||||
import unicodedata
|
||||
|
||||
from magic_pdf.para.commons import *
|
||||
|
||||
|
||||
if sys.version_info[0] >= 3:
|
||||
sys.stdout.reconfigure(encoding="utf-8") # type: ignore
|
||||
|
||||
|
||||
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_type1, font_type2, 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_type1, list):
|
||||
font_type1 = font_type1[0] if font_type1 else ""
|
||||
if isinstance(font_type2, list):
|
||||
font_type2 = font_type2[0] if font_type2 else ""
|
||||
|
||||
if font_type1 == font_type2:
|
||||
return True
|
||||
|
||||
# Find the length of the common prefix
|
||||
common_prefix_length = len(os.path.commonprefix([font_type1, font_type2]))
|
||||
|
||||
# Calculate the minimum prefix length based on the ratio
|
||||
min_prefix_length = int(min(len(font_type1), len(font_type2)) * prefix_length_ratio)
|
||||
|
||||
return common_prefix_length >= min_prefix_length
|
||||
|
||||
def __is_same_block_font(self, block1, block2):
|
||||
"""
|
||||
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(block1, "block_font_type", "")
|
||||
block_1_font_size = safe_get(block1, "block_font_size", 0)
|
||||
block_1_avg_char_width = safe_get(block1, "avg_char_width", 0)
|
||||
|
||||
block_2_font_type = safe_get(block2, "block_font_type", "")
|
||||
block_2_font_size = safe_get(block2, "block_font_size", 0)
|
||||
block_2_avg_char_width = safe_get(block2, "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(block1, "text", "")
|
||||
block_2_text = safe_get(block2, "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_condtion = 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_condtion
|
||||
)
|
||||
|
||||
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 __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]))
|
||||
)
|
||||
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, block1, block2):
|
||||
"""
|
||||
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(block1, block2)
|
||||
|
||||
def _is_para_continued(self, para1, para2):
|
||||
"""
|
||||
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(para1, para2)
|
||||
is_para_puncs_consistent = self._is_para_puncs_consistent(para1, para2)
|
||||
|
||||
return is_para_font_consistent and is_para_puncs_consistent
|
||||
|
||||
def _are_boundaries_of_block_consistent(self, block1, block2):
|
||||
"""
|
||||
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_block1 = block1["lines"][-1]
|
||||
first_line_of_block2 = block2["lines"][0]
|
||||
|
||||
spans_of_last_line_of_block1 = last_line_of_block1["spans"]
|
||||
spans_of_first_line_of_block2 = first_line_of_block2["spans"]
|
||||
|
||||
font_type_of_last_line_of_block1 = spans_of_last_line_of_block1[0]["font"].lower()
|
||||
font_size_of_last_line_of_block1 = spans_of_last_line_of_block1[0]["size"]
|
||||
font_color_of_last_line_of_block1 = spans_of_last_line_of_block1[0]["color"]
|
||||
font_flags_of_last_line_of_block1 = spans_of_last_line_of_block1[0]["flags"]
|
||||
|
||||
font_type_of_first_line_of_block2 = spans_of_first_line_of_block2[0]["font"].lower()
|
||||
font_size_of_first_line_of_block2 = spans_of_first_line_of_block2[0]["size"]
|
||||
font_color_of_first_line_of_block2 = spans_of_first_line_of_block2[0]["color"]
|
||||
font_flags_of_first_line_of_block2 = spans_of_first_line_of_block2[0]["flags"]
|
||||
|
||||
return (
|
||||
self.__is_similar_font_type(font_type_of_last_line_of_block1, font_type_of_first_line_of_block2)
|
||||
and abs(font_size_of_last_line_of_block1 - font_size_of_first_line_of_block2) < 1
|
||||
# and font_color_of_last_line_of_block1 == font_color_of_first_line_of_block2
|
||||
and font_flags_of_last_line_of_block1 == font_flags_of_first_line_of_block2
|
||||
)
|
||||
|
||||
def _get_last_paragraph(self, block):
|
||||
"""
|
||||
Retrieves the last paragraph from a block.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
block : dict
|
||||
The block from which to retrieve the paragraph.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict
|
||||
The last paragraph of the block.
|
||||
"""
|
||||
if block["paras"]:
|
||||
last_para_key = list(block["paras"].keys())[-1]
|
||||
return block["paras"][last_para_key]
|
||||
else:
|
||||
return None
|
||||
|
||||
def _get_first_paragraph(self, block):
|
||||
"""
|
||||
Retrieves the first paragraph from a block.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
block : dict
|
||||
The block from which to retrieve the paragraph.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict
|
||||
The first paragraph of the block.
|
||||
"""
|
||||
if block["paras"]:
|
||||
first_para_key = list(block["paras"].keys())[0]
|
||||
return block["paras"][first_para_key]
|
||||
else:
|
||||
return None
|
||||
|
||||
def should_merge_next_para(self, curr_para, next_para):
|
||||
if self._is_para_continued(curr_para, next_para):
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
def batch_tag_paras(self, pdf_dict):
|
||||
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):
|
||||
for block in para_blocks:
|
||||
if block.get("block_id") == block_id:
|
||||
return block
|
||||
return None
|
||||
|
||||
def batch_merge_paras(self, pdf_dict):
|
||||
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()):
|
||||
# 跳过标题段落
|
||||
if curr_para.get("is_para_title"):
|
||||
continue
|
||||
|
||||
while curr_para.get("merge_next_para"):
|
||||
next_para_location = curr_para.get("next_para_location")
|
||||
if not next_para_location:
|
||||
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
|
||||
@@ -1,480 +0,0 @@
|
||||
from magic_pdf.para.commons import *
|
||||
|
||||
|
||||
if sys.version_info[0] >= 3:
|
||||
sys.stdout.reconfigure(encoding="utf-8") # type: ignore
|
||||
|
||||
|
||||
|
||||
class BlockTerminationProcessor:
|
||||
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_types = [span["font"] for line in lines for span in line["spans"]]
|
||||
# most_common_font_type = max(set(font_types), key=font_types.count) if font_types 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
|
||||
@@ -1,222 +0,0 @@
|
||||
import sys
|
||||
|
||||
from magic_pdf.libs.commons import fitz
|
||||
from termcolor import cprint
|
||||
|
||||
|
||||
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
|
||||
@@ -1,246 +0,0 @@
|
||||
import math
|
||||
|
||||
from collections import defaultdict
|
||||
from magic_pdf.para.commons import *
|
||||
|
||||
if sys.version_info[0] >= 3:
|
||||
sys.stdout.reconfigure(encoding="utf-8") # type: ignore
|
||||
|
||||
|
||||
class HeaderFooterProcessor:
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def get_most_common_bboxes(self, 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_header(self, 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
|
||||
"""
|
||||
|
||||
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
|
||||
|
||||
# Traverse all blocks in the document
|
||||
single_preproc_blocks = 0
|
||||
total_blocks = 0
|
||||
single_preproc_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_preproc_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_preproc_blocks / total_blocks > 0.5: # 50% of the blocks are single-line
|
||||
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 = self.get_most_common_bboxes(all_bboxes, page_height, position="top") if all_bboxes else []
|
||||
common_footer_bboxes = self.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
|
||||
|
||||
|
||||
class NonHorizontalTextProcessor:
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def detect_non_horizontal_texts(self, result_dict):
|
||||
"""
|
||||
This function detects watermarks and vertical margin notes in the document.
|
||||
|
||||
Watermarks are identified by finding blocks with the same coordinates and frequently occurring identical texts across multiple pages.
|
||||
If these conditions are met, the blocks are highly likely to be watermarks, as opposed to headers or footers, which can change from page to page.
|
||||
If the direction of these blocks is not horizontal, they are definitely considered to be watermarks.
|
||||
|
||||
Vertical margin notes are identified by finding blocks with the same coordinates and frequently occurring identical texts across multiple pages.
|
||||
If these conditions are met, the blocks are highly likely to be vertical margin notes, which typically appear on the left and right sides of the page.
|
||||
If the direction of these blocks is vertical, they are definitely considered to be vertical margin notes.
|
||||
|
||||
|
||||
Parameters
|
||||
----------
|
||||
result_dict : dict
|
||||
The result dictionary.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result_dict : dict
|
||||
The updated result dictionary.
|
||||
"""
|
||||
# Dictionary to store information about potential watermarks
|
||||
potential_watermarks = {}
|
||||
potential_margin_notes = {}
|
||||
|
||||
for page_id, page_content in result_dict.items():
|
||||
if page_id.startswith("page_"):
|
||||
for block_id, block_data in page_content.items():
|
||||
if block_id.startswith("block_"):
|
||||
if "dir" in block_data:
|
||||
coordinates_text = (block_data["bbox"], block_data["text"]) # Tuple of coordinates and text
|
||||
|
||||
angle = math.atan2(block_data["dir"][1], block_data["dir"][0])
|
||||
angle = abs(math.degrees(angle))
|
||||
|
||||
if angle > 5 and angle < 85: # Check if direction is watermarks
|
||||
if coordinates_text in potential_watermarks:
|
||||
potential_watermarks[coordinates_text] += 1
|
||||
else:
|
||||
potential_watermarks[coordinates_text] = 1
|
||||
|
||||
if angle > 85 and angle < 105: # Check if direction is vertical
|
||||
if coordinates_text in potential_margin_notes:
|
||||
potential_margin_notes[coordinates_text] += 1 # Increment count
|
||||
else:
|
||||
potential_margin_notes[coordinates_text] = 1 # Initialize count
|
||||
|
||||
# Identify watermarks by finding entries with counts higher than a threshold (e.g., appearing on more than half of the pages)
|
||||
watermark_threshold = len(result_dict) // 2
|
||||
watermarks = {k: v for k, v in potential_watermarks.items() if v > watermark_threshold}
|
||||
|
||||
# Identify margin notes by finding entries with counts higher than a threshold (e.g., appearing on more than half of the pages)
|
||||
margin_note_threshold = len(result_dict) // 2
|
||||
margin_notes = {k: v for k, v in potential_margin_notes.items() if v > margin_note_threshold}
|
||||
|
||||
# Add watermark information to the result dictionary
|
||||
for page_id, blocks in result_dict.items():
|
||||
if page_id.startswith("page_"):
|
||||
for block_id, block_data in blocks.items():
|
||||
coordinates_text = (block_data["bbox"], block_data["text"])
|
||||
if coordinates_text in watermarks:
|
||||
block_data["is_watermark"] = 1
|
||||
else:
|
||||
block_data["is_watermark"] = 0
|
||||
|
||||
if coordinates_text in margin_notes:
|
||||
block_data["is_vertical_margin_note"] = 1
|
||||
else:
|
||||
block_data["is_vertical_margin_note"] = 0
|
||||
|
||||
return result_dict
|
||||
|
||||
|
||||
class NoiseRemover:
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def skip_data_noises(self, result_dict):
|
||||
"""
|
||||
This function skips the data noises, including overlap blocks, header, footer, watermark, vertical margin note, title
|
||||
"""
|
||||
filtered_result_dict = {}
|
||||
for page_id, blocks in result_dict.items():
|
||||
if page_id.startswith("page_"):
|
||||
filtered_blocks = {}
|
||||
for block_id, block in blocks.items():
|
||||
if block_id.startswith("block_"):
|
||||
if any(
|
||||
block.get(key, 0)
|
||||
for key in [
|
||||
"is_overlap",
|
||||
"is_header",
|
||||
"is_footer",
|
||||
"is_watermark",
|
||||
"is_vertical_margin_note",
|
||||
"is_block_title",
|
||||
]
|
||||
):
|
||||
continue
|
||||
filtered_blocks[block_id] = block
|
||||
if filtered_blocks:
|
||||
filtered_result_dict[page_id] = filtered_blocks
|
||||
|
||||
return filtered_result_dict
|
||||
@@ -1,121 +0,0 @@
|
||||
from magic_pdf.libs.commons import fitz
|
||||
|
||||
from magic_pdf.para.commons import *
|
||||
|
||||
|
||||
if sys.version_info[0] >= 3:
|
||||
sys.stdout.reconfigure(encoding="utf-8") # type: ignore
|
||||
|
||||
|
||||
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):
|
||||
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()
|
||||
@@ -1,198 +0,0 @@
|
||||
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
|
||||
@@ -1,40 +0,0 @@
|
||||
import math
|
||||
from magic_pdf.para.commons import *
|
||||
|
||||
|
||||
if sys.version_info[0] >= 3:
|
||||
sys.stdout.reconfigure(encoding="utf-8") # type: ignore
|
||||
|
||||
|
||||
class LayoutFilterProcessor:
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def batch_process_blocks(self, pdf_dict):
|
||||
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]
|
||||
|
||||
# Use math.ceil function to enlarge each value of x0, y0, x1, y1 of each layout_bbox
|
||||
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
|
||||
@@ -1,807 +0,0 @@
|
||||
import numpy as np
|
||||
from loguru import logger
|
||||
from sklearn.cluster import DBSCAN
|
||||
|
||||
from magic_pdf.config.ocr_content_type import 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
|
||||
|
||||
|
||||
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):
|
||||
"""探测是否包含了列表，并且把列表的行分开.
|
||||
|
||||
这样的段落特点是，顶格字母大写/数字，紧跟着几行缩进的。缩进的行首字母含小写的。
|
||||
"""
|
||||
|
||||
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
|
||||
else:
|
||||
total_lines = len(lines)
|
||||
line_fea_encode = []
|
||||
"""
|
||||
对每一行进行特征编码，编码规则如下：
|
||||
1. 如果行顶格，且大写字母开头或者数字开头，编码为1
|
||||
2. 如果顶格，其他非大写开头编码为4
|
||||
3. 如果非顶格，首字符大写，编码为2
|
||||
4. 如果非顶格，首字符非大写编码为3
|
||||
"""
|
||||
for l in lines: # noqa: E741
|
||||
first_char = __get_span_text(l['spans'][0])[0]
|
||||
layout_left = __find_layout_bbox_by_line(l['bbox'], new_layout_bboxes)[0]
|
||||
if l['bbox'][0] == layout_left:
|
||||
if first_char.isupper() or first_char.isdigit():
|
||||
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_patterns(line_fea_encode)
|
||||
if len(list_indice) > 0:
|
||||
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:
|
||||
logger.info(f'列表行的第{i}行不是顶格的')
|
||||
break
|
||||
else:
|
||||
logger.info(f'列表行的第{start}到第{end}行是列表')
|
||||
|
||||
return split_indices(total_lines, list_indice), list_start_idx
|
||||
|
||||
|
||||
def __valign_lines(blocks, layout_bboxes):
|
||||
"""在一个layoutbox内对齐行的左侧和右侧。 扫描行的左侧和右侧，如果x0,
|
||||
x1差距不超过一个阈值，就强行对齐到所处layout的左右两侧（和layout有一段距离）。
|
||||
3是个经验值，TODO，计算得来，可以设置为1.5个正文字符。"""
|
||||
|
||||
min_distance = 3
|
||||
min_sample = 2
|
||||
new_layout_bboxes = []
|
||||
|
||||
for layout_box in layout_bboxes:
|
||||
blocks_in_layoutbox = [
|
||||
b for b in blocks if is_in_layout(b['bbox'], layout_box['layout_bbox'])
|
||||
]
|
||||
if len(blocks_in_layoutbox) == 0:
|
||||
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:
|
||||
block['bbox'] = [
|
||||
min([line['bbox'][0] for line in block['lines']]),
|
||||
min([line['bbox'][1] for line in block['lines']]),
|
||||
max([line['bbox'][2] for line in block['lines']]),
|
||||
max([line['bbox'][3] for line in block['lines']]),
|
||||
]
|
||||
|
||||
"""新计算layout的bbox，因为block的bbox变了。"""
|
||||
layout_x0 = min([block['bbox'][0] for block in blocks_in_layoutbox])
|
||||
layout_y0 = min([block['bbox'][1] for block in blocks_in_layoutbox])
|
||||
layout_x1 = max([block['bbox'][2] for block in blocks_in_layoutbox])
|
||||
layout_y1 = max([block['bbox'][3] for block in blocks_in_layoutbox])
|
||||
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 = [b for b in blocks if is_in_layout(b['bbox'], lb)]
|
||||
if len(blocks_in_layoutbox) == 0:
|
||||
continue
|
||||
|
||||
for block in blocks_in_layoutbox:
|
||||
for line in block['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, lang='en'):
|
||||
"""每个layout内的行进行聚合."""
|
||||
# 因为只是一个block一行目前, 一个block就是一个段落
|
||||
lines_group = []
|
||||
|
||||
for lyout in layout_bboxes:
|
||||
lines = [
|
||||
line
|
||||
for block in blocks
|
||||
if is_in_layout(block['bbox'], lyout['layout_bbox'])
|
||||
for line in block['lines']
|
||||
]
|
||||
lines_group.append(lines)
|
||||
|
||||
return lines_group
|
||||
|
||||
|
||||
def __split_para_in_layoutbox(lines_group, new_layout_bbox, lang='en', char_avg_len=10):
|
||||
"""
|
||||
lines_group 进行行分段——layout内部进行分段。lines_group内每个元素是一个Layoutbox内的所有行。
|
||||
1. 先计算每个group的左右边界。
|
||||
2. 然后根据行末尾特征进行分段。
|
||||
末尾特征：以句号等结束符结尾。并且距离右侧边界有一定距离。
|
||||
且下一行开头不留空白。
|
||||
|
||||
"""
|
||||
list_info = [] # 这个layout最后是不是列表,记录每一个layout里是不是列表开头，列表结尾
|
||||
layout_paras = []
|
||||
right_tail_distance = 1.5 * char_avg_len
|
||||
|
||||
for lines in lines_group:
|
||||
paras = []
|
||||
total_lines = len(lines)
|
||||
if total_lines == 0:
|
||||
continue # 0行无需处理
|
||||
if total_lines == 1: # 1行无法分段。
|
||||
layout_paras.append([lines])
|
||||
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分成几个部分
|
||||
|
||||
"""
|
||||
|
||||
layout_right = __find_layout_bbox_by_line(lines[0]['bbox'], new_layout_bbox)[2]
|
||||
layout_left = __find_layout_bbox_by_line(lines[0]['bbox'], new_layout_bbox)[0]
|
||||
para = [] # 元素是line
|
||||
layout_list_info = [
|
||||
False,
|
||||
False,
|
||||
] # 这个layout最后是不是列表,记录每一个layout里是不是列表开头，列表结尾
|
||||
for content_type, start, end in text_segments:
|
||||
if content_type == 'list':
|
||||
for i, line in enumerate(lines[start : end + 1]):
|
||||
line_x0 = line['bbox'][0]
|
||||
if line_x0 == layout_left: # 列表开头
|
||||
if len(para) > 0:
|
||||
paras.append(para)
|
||||
para = []
|
||||
para.append(line)
|
||||
else:
|
||||
para.append(line)
|
||||
if len(para) > 0:
|
||||
paras.append(para)
|
||||
para = []
|
||||
if start == 0:
|
||||
layout_list_info[0] = True
|
||||
if end == total_lines - 1:
|
||||
layout_list_info[1] = True
|
||||
else: # 是普通文本
|
||||
for i, line in enumerate(lines[start : end + 1]):
|
||||
# 如果i有下一行，那么就要根据下一行位置综合判断是否要分段。如果i之后没有行，那么只需要判断i行自己的结尾特征。
|
||||
cur_line_type = line['spans'][-1]['type']
|
||||
next_line = lines[i + 1] if i < total_lines - 1 else None
|
||||
|
||||
if cur_line_type in [TEXT, INLINE_EQUATION]:
|
||||
if line['bbox'][2] < layout_right - right_tail_distance:
|
||||
para.append(line)
|
||||
paras.append(para)
|
||||
para = []
|
||||
elif (
|
||||
line['bbox'][2] >= layout_right - right_tail_distance
|
||||
and next_line
|
||||
and next_line['bbox'][0] == layout_left
|
||||
): # 现在这行到了行尾沾满，下一行存在且顶格。
|
||||
para.append(line)
|
||||
else:
|
||||
para.append(line)
|
||||
paras.append(para)
|
||||
para = []
|
||||
else: # 其他，图片、表格、行间公式，各自占一段
|
||||
if len(para) > 0: # 先把之前的段落加入到结果中
|
||||
paras.append(para)
|
||||
para = []
|
||||
paras.append(
|
||||
[line]
|
||||
) # 再把当前行加入到结果中。当前行为行间公式、图、表等。
|
||||
para = []
|
||||
|
||||
if len(para) > 0:
|
||||
paras.append(para)
|
||||
para = []
|
||||
|
||||
list_info.append(layout_list_info)
|
||||
layout_paras.append(paras)
|
||||
paras = []
|
||||
|
||||
return layout_paras, list_info
|
||||
|
||||
|
||||
def __connect_list_inter_layout(
|
||||
layout_paras, new_layout_bbox, layout_list_info, page_num, lang
|
||||
):
|
||||
"""如果上个layout的最后一个段落是列表，下一个layout的第一个段落也是列表，那么将他们连接起来。 TODO
|
||||
因为没有区分列表和段落，所以这个方法暂时不实现。
|
||||
根据layout_list_info判断是不是列表。，下个layout的第一个段如果不是列表，那么看他们是否有几行都有相同的缩进。"""
|
||||
if (
|
||||
len(layout_paras) == 0 or len(layout_list_info) == 0
|
||||
): # 0的时候最后的return 会出错
|
||||
return layout_paras, [False, False]
|
||||
|
||||
for i in range(1, len(layout_paras)):
|
||||
pre_layout_list_info = layout_list_info[i - 1]
|
||||
next_layout_list_info = layout_list_info[i]
|
||||
pre_last_para = layout_paras[i - 1][-1]
|
||||
next_paras = layout_paras[i]
|
||||
|
||||
if (
|
||||
pre_layout_list_info[1] and not next_layout_list_info[0]
|
||||
): # 前一个是列表结尾，后一个是非列表开头，此时检测是否有相同的缩进
|
||||
logger.info(f'连接page {page_num} 内的list')
|
||||
# 向layout_paras[i] 寻找开头具有相同缩进的连续的行
|
||||
may_list_lines = []
|
||||
for j in range(len(next_paras)):
|
||||
line = next_paras[j]
|
||||
if len(line) == 1: # 只可能是一行，多行情况再需要分析了
|
||||
if (
|
||||
line[0]['bbox'][0]
|
||||
> __find_layout_bbox_by_line(line[0]['bbox'], new_layout_bbox)[
|
||||
0
|
||||
]
|
||||
):
|
||||
may_list_lines.append(line[0])
|
||||
else:
|
||||
break
|
||||
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)
|
||||
layout_paras[i] = layout_paras[i][len(may_list_lines) :]
|
||||
|
||||
return layout_paras, [
|
||||
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 (
|
||||
pre_page_list_info[1] and not next_page_list_info[0]
|
||||
): # 前一个是列表结尾，后一个是非列表开头，此时检测是否有相同的缩进
|
||||
logger.info(f'连接page {page_num} 内的list')
|
||||
# 向layout_paras[i] 寻找开头具有相同缩进的连续的行
|
||||
may_list_lines = []
|
||||
for j in range(len(next_page_paras[0])):
|
||||
line = next_page_paras[0][j]
|
||||
if len(line) == 1: # 只可能是一行，多行情况再需要分析了
|
||||
if (
|
||||
line[0]['bbox'][0]
|
||||
> __find_layout_bbox_by_line(
|
||||
line[0]['bbox'], next_page_layout_bbox
|
||||
)[0]
|
||||
):
|
||||
may_list_lines.append(line[0])
|
||||
else:
|
||||
break
|
||||
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)
|
||||
next_page_paras[0] = next_page_paras[0][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(layout_paras, new_layout_bbox, lang):
|
||||
"""
|
||||
layout之间进行分段。
|
||||
主要是计算前一个layOut的最后一行和后一个layout的第一行是否可以连接。
|
||||
连接的条件需要同时满足：
|
||||
1. 上一个layout的最后一行沾满整个行。并且没有结尾符号。
|
||||
2. 下一行开头不留空白。
|
||||
|
||||
"""
|
||||
connected_layout_paras = []
|
||||
if len(layout_paras) == 0:
|
||||
return connected_layout_paras
|
||||
|
||||
connected_layout_paras.append(layout_paras[0])
|
||||
for i in range(1, len(layout_paras)):
|
||||
try:
|
||||
if (
|
||||
len(layout_paras[i]) == 0 or len(layout_paras[i - 1]) == 0
|
||||
): # TODO 考虑连接问题，
|
||||
continue
|
||||
pre_last_line = layout_paras[i - 1][-1][-1]
|
||||
next_first_line = layout_paras[i][0][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_paras.append(layout_paras[i])
|
||||
continue
|
||||
|
||||
pre_x2_max = __find_layout_bbox_by_line(pre_last_line['bbox'], new_layout_bbox)[
|
||||
2
|
||||
]
|
||||
next_x0_min = __find_layout_bbox_by_line(
|
||||
next_first_line['bbox'], new_layout_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的段落连接。"""
|
||||
connected_layout_paras[-1][-1].extend(layout_paras[i][0])
|
||||
layout_paras[i].pop(
|
||||
0
|
||||
) # 删除后一个layout的第一个段落， 因为他已经被合并到前一个layout的最后一个段落了。
|
||||
if len(layout_paras[i]) == 0:
|
||||
layout_paras.pop(i)
|
||||
else:
|
||||
connected_layout_paras.append(layout_paras[i])
|
||||
else:
|
||||
"""连接段落条件不成立，将前一个layout的段落加入到结果中。"""
|
||||
connected_layout_paras.append(layout_paras[i])
|
||||
|
||||
return connected_layout_paras
|
||||
|
||||
|
||||
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_para = pre_page_paras[-1][-1]
|
||||
next_first_para = next_page_paras[0][0]
|
||||
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 = __find_layout_bbox_by_line(
|
||||
pre_last_line['bbox'], pre_page_layout_bbox
|
||||
)[2]
|
||||
next_x0_min = __find_layout_bbox_by_line(
|
||||
next_first_line['bbox'], next_page_layout_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的段落连接。"""
|
||||
pre_last_para.extend(next_first_para)
|
||||
next_page_paras[0].pop(
|
||||
0
|
||||
) # 删除后一个页面的第一个段落， 因为他已经被合并到前一个页面的最后一个段落了。
|
||||
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, debug_mode
|
||||
):
|
||||
"""
|
||||
找出来中间对齐的连续单行文本，如果连续行高度相同，那么合并为一个段落。
|
||||
一个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
|
||||
)
|
||||
|
||||
"""找出来连续的单行文本，如果连续行高度相同，那么合并为一个段落。"""
|
||||
consecutive_single_line_indices = find_consecutive_true_regions(
|
||||
single_line_paras_tag
|
||||
)
|
||||
if len(consecutive_single_line_indices) > 0:
|
||||
index_offset = 0
|
||||
"""检查这些行是否是高度相同的，居中的"""
|
||||
for start, end in consecutive_single_line_indices:
|
||||
start += index_offset
|
||||
end += index_offset
|
||||
line_hi = np.array(
|
||||
[
|
||||
line[0]['bbox'][3] - line[0]['bbox'][1]
|
||||
for line in layout_para[start : end + 1]
|
||||
]
|
||||
)
|
||||
first_line_text = ''.join(
|
||||
[__get_span_text(span) for span in layout_para[start][0]['spans']]
|
||||
)
|
||||
if 'Table' in first_line_text or 'Figure' in first_line_text:
|
||||
pass
|
||||
if debug_mode:
|
||||
logger.debug(line_hi.std())
|
||||
|
||||
if line_hi.std() < 2:
|
||||
"""行高度相同，那么判断是否居中."""
|
||||
all_left_x0 = [
|
||||
line[0]['bbox'][0] for line in layout_para[start : end + 1]
|
||||
]
|
||||
all_right_x1 = [
|
||||
line[0]['bbox'][2] for line 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 = [l[0] for l in layout_para[start : end + 1]] # noqa: E741
|
||||
para_text = ''.join(
|
||||
[
|
||||
__get_span_text(span)
|
||||
for line in merge_para
|
||||
for span in line['spans']
|
||||
]
|
||||
)
|
||||
if debug_mode:
|
||||
logger.debug(para_text)
|
||||
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. 图、表，目前独占一行，不考虑分段。
|
||||
"""
|
||||
if page_num == 343:
|
||||
pass
|
||||
lines_group = __group_line_by_layout(blocks, layout_bboxes, lang) # block内分段
|
||||
layout_paras, layout_list_info = __split_para_in_layoutbox(
|
||||
lines_group, new_layout_bbox, lang
|
||||
) # layout内分段
|
||||
layout_paras2, page_list_info = __connect_list_inter_layout(
|
||||
layout_paras, new_layout_bbox, layout_list_info, page_num, lang
|
||||
) # layout之间连接列表段落
|
||||
connected_layout_paras = __connect_para_inter_layoutbox(
|
||||
layout_paras2, new_layout_bbox, lang
|
||||
) # layout间链接段落
|
||||
|
||||
return connected_layout_paras, page_list_info
|
||||
|
||||
|
||||
def para_split(pdf_info_dict, debug_mode, lang='en'):
|
||||
"""根据line和layout情况进行分段."""
|
||||
new_layout_of_pages = [] # 数组的数组，每个元素是一个页面的layoutS
|
||||
all_page_list_info = [] # 保存每个页面开头和结尾是否是列表
|
||||
for page_num, page in pdf_info_dict.items():
|
||||
blocks = 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_mode:
|
||||
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_mode:
|
||||
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, debug_mode=debug_mode
|
||||
)
|
||||
__merge_signle_list_text(page_paras, new_layout_bbox, page_num, lang)
|
||||
@@ -1,959 +0,0 @@
|
||||
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
|
||||
@@ -1,207 +0,0 @@
|
||||
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
|
||||
|
||||
@@ -1,268 +0,0 @@
|
||||
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
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -1,60 +0,0 @@
|
||||
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
|
||||
@@ -1,153 +0,0 @@
|
||||
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
|
||||
@@ -1,161 +0,0 @@
|
||||
"""
|
||||
去掉正文的引文引用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
|
||||
@@ -1,134 +0,0 @@
|
||||
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
|
||||
@@ -1,64 +0,0 @@
|
||||
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
|
||||
|
||||
@@ -1,284 +0,0 @@
|
||||
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
|
||||
@@ -1,170 +0,0 @@
|
||||
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 []
|
||||
|
||||
|
||||
|
||||
@@ -1,64 +0,0 @@
|
||||
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
|
||||
|
||||
@@ -1,647 +0,0 @@
|
||||
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
|
||||
|
||||
@@ -1,64 +0,0 @@
|
||||
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
|
||||
|
||||
@@ -1,62 +0,0 @@
|
||||
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
|
||||
|
||||
@@ -1,550 +0,0 @@
|
||||
"""对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
|
||||
@@ -1,244 +0,0 @@
|
||||
|
||||
|
||||
|
||||
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
|
||||
@@ -1,270 +0,0 @@
|
||||
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
|
||||
@@ -1,23 +0,0 @@
|
||||
import collections
|
||||
|
||||
|
||||
def get_main_text_font(pdf_docs):
|
||||
font_names = collections.Counter()
|
||||
for page in pdf_docs:
|
||||
blocks = page.get_text('dict')['blocks']
|
||||
if blocks is not None:
|
||||
for block in blocks:
|
||||
lines = block.get('lines')
|
||||
if lines is not None:
|
||||
for line in lines:
|
||||
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[font] += count
|
||||
main_text_font = font_names.most_common(1)[0][0]
|
||||
return main_text_font
|
||||
|
||||
@@ -1,133 +0,0 @@
|
||||
import fitz
|
||||
|
||||
from magic_pdf.layout.layout_sort import get_bboxes_layout
|
||||
from magic_pdf.libs.boxbase import _is_part_overlap, _is_in
|
||||
from magic_pdf.libs.coordinate_transform import get_scale_ratio
|
||||
|
||||
|
||||
def get_center_point(bbox):
|
||||
"""
|
||||
根据边界框坐标信息，计算出该边界框的中心点坐标。
|
||||
Args:
|
||||
bbox (list): 边界框坐标信息，包含四个元素，分别为左上角x坐标、左上角y坐标、右下角x坐标、右下角y坐标。
|
||||
Returns:
|
||||
list: 中心点坐标信息，包含两个元素，分别为x坐标和y坐标。
|
||||
"""
|
||||
return [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2]
|
||||
|
||||
|
||||
def get_area(bbox):
|
||||
"""
|
||||
根据边界框坐标信息，计算出该边界框的面积。
|
||||
Args:
|
||||
bbox (list): 边界框坐标信息，包含四个元素，分别为左上角x坐标、左上角y坐标、右下角x坐标、右下角y坐标。
|
||||
Returns:
|
||||
float: 该边界框的面积。
|
||||
"""
|
||||
return (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
|
||||
|
||||
|
||||
def adjust_layouts(layout_bboxes, page_boundry, page_id):
|
||||
# 遍历所有布局框
|
||||
for i in range(len(layout_bboxes)):
|
||||
# 遍历当前布局框之后的布局框
|
||||
for j in range(i + 1, len(layout_bboxes)):
|
||||
# 判断两个布局框是否重叠
|
||||
if _is_part_overlap(layout_bboxes[i], layout_bboxes[j]):
|
||||
# 计算每个布局框的中心点坐标和面积
|
||||
area_i = get_area(layout_bboxes[i])
|
||||
area_j = get_area(layout_bboxes[j])
|
||||
|
||||
# 较大布局框和较小布局框的赋值
|
||||
if area_i > area_j:
|
||||
larger_layout, smaller_layout = layout_bboxes[i], layout_bboxes[j]
|
||||
else:
|
||||
larger_layout, smaller_layout = layout_bboxes[j], layout_bboxes[i]
|
||||
|
||||
center_large = get_center_point(larger_layout)
|
||||
center_small = get_center_point(smaller_layout)
|
||||
# 计算横向和纵向的距离差
|
||||
distance_x = center_large[0] - center_small[0]
|
||||
distance_y = center_large[1] - center_small[1]
|
||||
|
||||
# 根据距离差判断重叠方向并修正边界
|
||||
if abs(distance_x) > abs(distance_y): # 左右重叠
|
||||
if distance_x > 0 and larger_layout[0] < smaller_layout[2]:
|
||||
larger_layout[0] = smaller_layout[2]+1
|
||||
if distance_x < 0 and larger_layout[2] > smaller_layout[0]:
|
||||
larger_layout[2] = smaller_layout[0]-1
|
||||
else: # 上下重叠
|
||||
if distance_y > 0 and larger_layout[1] < smaller_layout[3]:
|
||||
larger_layout[1] = smaller_layout[3]+1
|
||||
if distance_y < 0 and larger_layout[3] > smaller_layout[1]:
|
||||
larger_layout[3] = smaller_layout[1]-1
|
||||
# 排序调整布局边界框列表
|
||||
new_bboxes = []
|
||||
for layout_bbox in layout_bboxes:
|
||||
new_bboxes.append([layout_bbox[0], layout_bbox[1], layout_bbox[2], layout_bbox[3], None, None, None, None, None, None, None, None, None])
|
||||
|
||||
layout_bboxes, layout_tree = get_bboxes_layout(new_bboxes, page_boundry, page_id)
|
||||
|
||||
# 返回排序调整后的布局边界框列表
|
||||
return layout_bboxes, layout_tree
|
||||
|
||||
|
||||
def layout_detect(layout_info, page: fitz.Page, ocr_page_info):
|
||||
"""
|
||||
对输入的布局信息进行解析，提取出每个子布局的边界框，并对所有子布局进行排序调整。
|
||||
|
||||
Args:
|
||||
layout_info (list): 包含子布局信息的列表，每个子布局信息为字典类型，包含'poly'字段，表示子布局的边界框坐标信息。
|
||||
|
||||
Returns:
|
||||
list: 经过排序调整后的所有子布局边界框信息的列表，每个边界框信息为字典类型，包含'layout_bbox'字段，表示边界框的坐标信息。
|
||||
|
||||
"""
|
||||
page_id = ocr_page_info['page_info']['page_no']-1
|
||||
horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(ocr_page_info, page)
|
||||
# 初始化布局边界框列表
|
||||
layout_bboxes = []
|
||||
# 遍历每个子布局
|
||||
for sub_layout in layout_info:
|
||||
# 提取子布局的边界框坐标信息
|
||||
x0, y0, _, _, x1, y1, _, _ = sub_layout['poly']
|
||||
bbox = [int(x0 / horizontal_scale_ratio), int(y0 / vertical_scale_ratio),
|
||||
int(x1 / horizontal_scale_ratio), int(y1 / vertical_scale_ratio)]
|
||||
|
||||
# 将子布局的边界框添加到列表中
|
||||
layout_bboxes.append(bbox)
|
||||
|
||||
# 初始化新的布局边界框列表
|
||||
new_layout_bboxes = []
|
||||
# 遍历每个布局边界框
|
||||
for i in range(len(layout_bboxes)):
|
||||
# 初始化标记变量，用于判断当前边界框是否需要保留
|
||||
keep = True
|
||||
# 获取当前边界框的坐标信息
|
||||
box_i = layout_bboxes[i]
|
||||
|
||||
# 遍历其他边界框
|
||||
for j in range(len(layout_bboxes)):
|
||||
# 排除当前边界框自身
|
||||
if i != j:
|
||||
# 获取其他边界框的坐标信息
|
||||
box_j = layout_bboxes[j]
|
||||
# 检测box_i是否被box_j包含
|
||||
if _is_in(box_i, box_j):
|
||||
# 如果当前边界框被其他边界框包含，则标记为不需要保留
|
||||
keep = False
|
||||
# 跳出内层循环
|
||||
break
|
||||
|
||||
# 如果当前边界框需要保留，则添加到新的布局边界框列表中
|
||||
if keep:
|
||||
new_layout_bboxes.append(layout_bboxes[i])
|
||||
|
||||
# 对新的布局边界框列表进行排序调整
|
||||
page_width = page.rect.width
|
||||
page_height = page.rect.height
|
||||
page_boundry = [0, 0, page_width, page_height]
|
||||
layout_bboxes, layout_tree = adjust_layouts(new_layout_bboxes, page_boundry, page_id)
|
||||
|
||||
# 返回排序调整后的布局边界框列表
|
||||
return layout_bboxes, layout_tree
|
||||
@@ -1,78 +0,0 @@
|
||||
from magic_pdf.config.drop_reason import DropReason
|
||||
from magic_pdf.libs.boxbase import _is_in, _is_in_or_part_overlap
|
||||
from magic_pdf.libs.commons import fitz
|
||||
|
||||
|
||||
def __area(box):
|
||||
return (box[2] - box[0]) * (box[3] - box[1])
|
||||
|
||||
|
||||
def __is_contain_color_background_rect(
|
||||
page: fitz.Page, text_blocks, image_bboxes
|
||||
) -> bool:
|
||||
"""检查page是包含有颜色背景的矩形."""
|
||||
color_bg_rect = []
|
||||
p_width, p_height = page.rect.width, page.rect.height
|
||||
|
||||
# 先找到最大的带背景矩形
|
||||
blocks = page.get_cdrawings()
|
||||
for block in blocks:
|
||||
if 'fill' in block and block['fill']: # 过滤掉透明的
|
||||
fill = list(block['fill'])
|
||||
fill[0], fill[1], fill[2] = int(fill[0]), int(fill[1]), int(fill[2])
|
||||
if fill == (1.0, 1.0, 1.0):
|
||||
continue
|
||||
rect = block['rect']
|
||||
# 过滤掉特别小的矩形
|
||||
if __area(rect) < 10 * 10:
|
||||
continue
|
||||
# 为了防止是svg图片上的色块，这里过滤掉这类
|
||||
|
||||
if any(
|
||||
[_is_in_or_part_overlap(rect, img_bbox) for img_bbox in image_bboxes]
|
||||
):
|
||||
continue
|
||||
color_bg_rect.append(rect)
|
||||
|
||||
# 找到最大的背景矩形
|
||||
if len(color_bg_rect) > 0:
|
||||
max_rect = max(color_bg_rect, key=lambda x: __area(x))
|
||||
max_rect_int = (
|
||||
int(max_rect[0]),
|
||||
int(max_rect[1]),
|
||||
int(max_rect[2]),
|
||||
int(max_rect[3]),
|
||||
)
|
||||
# 判断最大的背景矩形是否包含超过3行文字，或者50个字 TODO
|
||||
if (
|
||||
max_rect[2] - max_rect[0] > 0.2 * p_width
|
||||
and max_rect[3] - max_rect[1] > 0.1 * p_height
|
||||
): # 宽度符合
|
||||
# 看是否有文本块落入到这个矩形中
|
||||
for text_block in text_blocks:
|
||||
box = text_block['bbox']
|
||||
box_int = (int(box[0]), int(box[1]), int(box[2]), int(box[3]))
|
||||
if _is_in(box_int, max_rect_int):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def __is_table_overlap_text_block(text_blocks, table_bbox):
|
||||
"""检查table_bbox是否覆盖了text_blocks里的文本块 TODO."""
|
||||
for text_block in text_blocks:
|
||||
box = text_block['bbox']
|
||||
if _is_in_or_part_overlap(table_bbox, box):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def pdf_filter(page: fitz.Page, text_blocks, table_bboxes, image_bboxes) -> tuple:
|
||||
"""return:(True|False, err_msg) True, 如果pdf符合要求 False, 如果pdf不符合要求."""
|
||||
if __is_contain_color_background_rect(page, text_blocks, image_bboxes):
|
||||
return False, {
|
||||
'_need_drop': True,
|
||||
'_drop_reason': DropReason.COLOR_BACKGROUND_TEXT_BOX,
|
||||
}
|
||||
|
||||
return True, None
|
||||
@@ -1,101 +0,0 @@
|
||||
from loguru import logger
|
||||
|
||||
from magic_pdf.config.drop_tag import COLOR_BG_HEADER_TXT_BLOCK
|
||||
from magic_pdf.libs.boxbase import (_is_in, _is_in_or_part_overlap,
|
||||
calculate_overlap_area_2_minbox_area_ratio)
|
||||
|
||||
|
||||
def __area(box):
|
||||
return (box[2] - box[0]) * (box[3] - box[1])
|
||||
|
||||
|
||||
def rectangle_position_determination(rect, p_width):
|
||||
"""判断矩形是否在页面中轴线附近。
|
||||
|
||||
Args:
|
||||
rect (list): 矩形坐标，格式为[x1, y1, x2, y2]。
|
||||
p_width (int): 页面宽度。
|
||||
|
||||
Returns:
|
||||
bool: 若矩形在页面中轴线附近则返回True，否则返回False。
|
||||
"""
|
||||
# 页面中轴线x坐标
|
||||
x_axis = p_width / 2
|
||||
# 矩形是否跨越中轴线
|
||||
is_span = rect[0] < x_axis and rect[2] > x_axis
|
||||
if is_span:
|
||||
return True
|
||||
else:
|
||||
# 矩形与中轴线的距离，只算近的那一边
|
||||
distance = rect[0] - x_axis if rect[0] > x_axis else x_axis - rect[2]
|
||||
# 判断矩形与中轴线的距离是否小于页面宽度的20%
|
||||
if distance < p_width * 0.2:
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
def remove_colored_strip_textblock(remain_text_blocks, page):
|
||||
"""根据页面中特定颜色和大小过滤文本块，将符合条件的文本块从remain_text_blocks中移除，并返回移除的文本块列表colored_str
|
||||
ip_textblock。
|
||||
|
||||
Args:
|
||||
remain_text_blocks (list): 剩余文本块列表。
|
||||
page (Page): 页面对象。
|
||||
|
||||
Returns:
|
||||
tuple: 剩余文本块列表和移除的文本块列表。
|
||||
"""
|
||||
colored_strip_textblocks = [] # 先构造一个空的返回
|
||||
if len(remain_text_blocks) > 0:
|
||||
p_width, p_height = page.rect.width, page.rect.height
|
||||
blocks = page.get_cdrawings()
|
||||
colored_strip_bg_rect = []
|
||||
for block in blocks:
|
||||
is_filled = (
|
||||
'fill' in block and block['fill'] and block['fill'] != (1.0, 1.0, 1.0)
|
||||
) # 过滤掉透明的
|
||||
rect = block['rect']
|
||||
area_is_large_enough = __area(rect) > 100 # 过滤掉特别小的矩形
|
||||
rectangle_position_determination_result = rectangle_position_determination(
|
||||
rect, p_width
|
||||
)
|
||||
in_upper_half_page = (
|
||||
rect[3] < p_height * 0.3
|
||||
) # 找到位于页面上半部分的矩形，下边界小于页面高度的30%
|
||||
aspect_ratio_exceeds_4 = (rect[2] - rect[0]) > (
|
||||
rect[3] - rect[1]
|
||||
) * 4 # 找到长宽比超过4的矩形
|
||||
|
||||
if (
|
||||
is_filled
|
||||
and area_is_large_enough
|
||||
and rectangle_position_determination_result
|
||||
and in_upper_half_page
|
||||
and aspect_ratio_exceeds_4
|
||||
):
|
||||
colored_strip_bg_rect.append(rect)
|
||||
|
||||
if len(colored_strip_bg_rect) > 0:
|
||||
for colored_strip_block_bbox in colored_strip_bg_rect:
|
||||
for text_block in remain_text_blocks:
|
||||
text_bbox = text_block['bbox']
|
||||
if _is_in(text_bbox, colored_strip_block_bbox) or (
|
||||
_is_in_or_part_overlap(text_bbox, colored_strip_block_bbox)
|
||||
and calculate_overlap_area_2_minbox_area_ratio(
|
||||
text_bbox, colored_strip_block_bbox
|
||||
)
|
||||
> 0.6
|
||||
):
|
||||
logger.info(
|
||||
f'remove_colored_strip_textblock: {text_bbox}, {colored_strip_block_bbox}'
|
||||
)
|
||||
text_block['tag'] = COLOR_BG_HEADER_TXT_BLOCK
|
||||
colored_strip_textblocks.append(text_block)
|
||||
|
||||
if len(colored_strip_textblocks) > 0:
|
||||
for colored_strip_textblock in colored_strip_textblocks:
|
||||
if colored_strip_textblock in remain_text_blocks:
|
||||
remain_text_blocks.remove(colored_strip_textblock)
|
||||
|
||||
return remain_text_blocks, colored_strip_textblocks
|
||||
@@ -1,114 +0,0 @@
|
||||
import re
|
||||
|
||||
from magic_pdf.config.drop_tag import CONTENT_IN_FOOT_OR_HEADER, PAGE_NO
|
||||
from magic_pdf.libs.boxbase import _is_in_or_part_overlap
|
||||
|
||||
|
||||
def remove_headder_footer_one_page(text_raw_blocks, image_bboxes, table_bboxes, header_bboxs, footer_bboxs,
|
||||
page_no_bboxs, page_w, page_h):
|
||||
"""删除页眉页脚，页码 从line级别进行删除，删除之后观察这个text-block是否是空的，如果是空的，则移动到remove_list中."""
|
||||
header = []
|
||||
footer = []
|
||||
if len(header) == 0:
|
||||
model_header = header_bboxs
|
||||
if model_header:
|
||||
x0 = min([x for x, _, _, _ in model_header])
|
||||
y0 = min([y for _, y, _, _ in model_header])
|
||||
x1 = max([x1 for _, _, x1, _ in model_header])
|
||||
y1 = max([y1 for _, _, _, y1 in model_header])
|
||||
header = [x0, y0, x1, y1]
|
||||
if len(footer) == 0:
|
||||
model_footer = footer_bboxs
|
||||
if model_footer:
|
||||
x0 = min([x for x, _, _, _ in model_footer])
|
||||
y0 = min([y for _, y, _, _ in model_footer])
|
||||
x1 = max([x1 for _, _, x1, _ in model_footer])
|
||||
y1 = max([y1 for _, _, _, y1 in model_footer])
|
||||
footer = [x0, y0, x1, y1]
|
||||
|
||||
header_y0 = 0 if len(header) == 0 else header[3]
|
||||
footer_y0 = page_h if len(footer) == 0 else footer[1]
|
||||
if page_no_bboxs:
|
||||
top_part = [b for b in page_no_bboxs if b[3] < page_h / 2]
|
||||
btn_part = [b for b in page_no_bboxs if b[1] > page_h / 2]
|
||||
|
||||
top_max_y0 = max([b[1] for b in top_part]) if top_part else 0
|
||||
btn_min_y1 = min([b[3] for b in btn_part]) if btn_part else page_h
|
||||
|
||||
header_y0 = max(header_y0, top_max_y0)
|
||||
footer_y0 = min(footer_y0, btn_min_y1)
|
||||
|
||||
content_boundry = [0, header_y0, page_w, footer_y0]
|
||||
|
||||
header = [0, 0, page_w, header_y0]
|
||||
footer = [0, footer_y0, page_w, page_h]
|
||||
|
||||
"""以上计算出来了页眉页脚的边界，下面开始进行删除"""
|
||||
text_block_to_remove = []
|
||||
# 首先检查每个textblock
|
||||
for blk in text_raw_blocks:
|
||||
if len(blk['lines']) > 0:
|
||||
for line in blk['lines']:
|
||||
line_del = []
|
||||
for span in line['spans']:
|
||||
span_del = []
|
||||
if span['bbox'][3] < header_y0:
|
||||
span_del.append(span)
|
||||
elif _is_in_or_part_overlap(span['bbox'], header) or _is_in_or_part_overlap(span['bbox'], footer):
|
||||
span_del.append(span)
|
||||
for span in span_del:
|
||||
line['spans'].remove(span)
|
||||
if not line['spans']:
|
||||
line_del.append(line)
|
||||
|
||||
for line in line_del:
|
||||
blk['lines'].remove(line)
|
||||
else:
|
||||
# if not blk['lines']:
|
||||
blk['tag'] = CONTENT_IN_FOOT_OR_HEADER
|
||||
text_block_to_remove.append(blk)
|
||||
|
||||
"""有的时候由于pageNo太小了，总是会有一点和content_boundry重叠一点，被放入正文，因此对于pageNo，进行span粒度的删除"""
|
||||
page_no_block_2_remove = []
|
||||
if page_no_bboxs:
|
||||
for pagenobox in page_no_bboxs:
|
||||
for block in text_raw_blocks:
|
||||
if _is_in_or_part_overlap(pagenobox, block['bbox']): # 在span级别删除页码
|
||||
for line in block['lines']:
|
||||
for span in line['spans']:
|
||||
if _is_in_or_part_overlap(pagenobox, span['bbox']):
|
||||
# span['text'] = ''
|
||||
span['tag'] = PAGE_NO
|
||||
# 检查这个block是否只有这一个span，如果是，那么就把这个block也删除
|
||||
if len(line['spans']) == 1 and len(block['lines']) == 1:
|
||||
page_no_block_2_remove.append(block)
|
||||
else:
|
||||
# 测试最后一个是不是页码：规则是，最后一个block仅有1个line,一个span,且text是数字，空格，符号组成，不含字母,并且包含数字
|
||||
if len(text_raw_blocks) > 0:
|
||||
text_raw_blocks.sort(key=lambda x: x['bbox'][1], reverse=True)
|
||||
last_block = text_raw_blocks[0]
|
||||
if len(last_block['lines']) == 1:
|
||||
last_line = last_block['lines'][0]
|
||||
if len(last_line['spans']) == 1:
|
||||
last_span = last_line['spans'][0]
|
||||
if last_span['text'].strip() and not re.search('[a-zA-Z]', last_span['text']) and re.search('[0-9]',
|
||||
last_span[
|
||||
'text']):
|
||||
last_span['tag'] = PAGE_NO
|
||||
page_no_block_2_remove.append(last_block)
|
||||
|
||||
for b in page_no_block_2_remove:
|
||||
text_block_to_remove.append(b)
|
||||
|
||||
for blk in text_block_to_remove:
|
||||
if blk in text_raw_blocks:
|
||||
text_raw_blocks.remove(blk)
|
||||
|
||||
text_block_remain = text_raw_blocks
|
||||
image_bbox_to_remove = [bbox for bbox in image_bboxes if not _is_in_or_part_overlap(bbox, content_boundry)]
|
||||
|
||||
image_bbox_remain = [bbox for bbox in image_bboxes if _is_in_or_part_overlap(bbox, content_boundry)]
|
||||
table_bbox_to_remove = [bbox for bbox in table_bboxes if not _is_in_or_part_overlap(bbox, content_boundry)]
|
||||
table_bbox_remain = [bbox for bbox in table_bboxes if _is_in_or_part_overlap(bbox, content_boundry)]
|
||||
|
||||
return image_bbox_remain, table_bbox_remain, text_block_remain, text_block_to_remove, image_bbox_to_remove, table_bbox_to_remove
|
||||
@@ -1,236 +0,0 @@
|
||||
import math
|
||||
import re
|
||||
|
||||
from magic_pdf.config.drop_tag import (EMPTY_SIDE_BLOCK, ROTATE_TEXT,
|
||||
VERTICAL_TEXT)
|
||||
from magic_pdf.libs.boxbase import is_vbox_on_side
|
||||
|
||||
|
||||
def detect_non_horizontal_texts(result_dict):
|
||||
"""This function detects watermarks and vertical margin notes in the
|
||||
document.
|
||||
|
||||
Watermarks are identified by finding blocks with the same coordinates and frequently occurring identical texts across multiple pages.
|
||||
If these conditions are met, the blocks are highly likely to be watermarks, as opposed to headers or footers, which can change from page to page.
|
||||
If the direction of these blocks is not horizontal, they are definitely considered to be watermarks.
|
||||
|
||||
Vertical margin notes are identified by finding blocks with the same coordinates and frequently occurring identical texts across multiple pages.
|
||||
If these conditions are met, the blocks are highly likely to be vertical margin notes, which typically appear on the left and right sides of the page. # noqa: E501
|
||||
If the direction of these blocks is vertical, they are definitely considered to be vertical margin notes.
|
||||
|
||||
|
||||
Parameters
|
||||
----------
|
||||
result_dict : dict
|
||||
The result dictionary.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result_dict : dict
|
||||
The updated result dictionary.
|
||||
"""
|
||||
# Dictionary to store information about potential watermarks
|
||||
potential_watermarks = {}
|
||||
potential_margin_notes = {}
|
||||
|
||||
for page_id, page_content in result_dict.items():
|
||||
if page_id.startswith('page_'):
|
||||
for block_id, block_data in page_content.items():
|
||||
if block_id.startswith('block_'):
|
||||
if 'dir' in block_data:
|
||||
coordinates_text = (
|
||||
block_data['bbox'],
|
||||
block_data['text'],
|
||||
) # Tuple of coordinates and text
|
||||
|
||||
angle = math.atan2(block_data['dir'][1], block_data['dir'][0])
|
||||
angle = abs(math.degrees(angle))
|
||||
|
||||
if angle > 5 and angle < 85: # Check if direction is watermarks
|
||||
if coordinates_text in potential_watermarks:
|
||||
potential_watermarks[coordinates_text] += 1
|
||||
else:
|
||||
potential_watermarks[coordinates_text] = 1
|
||||
|
||||
if angle > 85 and angle < 105: # Check if direction is vertical
|
||||
if coordinates_text in potential_margin_notes:
|
||||
potential_margin_notes[coordinates_text] += (
|
||||
1 # Increment count
|
||||
)
|
||||
else:
|
||||
potential_margin_notes[coordinates_text] = (
|
||||
1 # Initialize count
|
||||
)
|
||||
|
||||
# Identify watermarks by finding entries with counts higher than a threshold (e.g., appearing on more than half of the pages)
|
||||
watermark_threshold = len(result_dict) // 2
|
||||
watermarks = {
|
||||
k: v for k, v in potential_watermarks.items() if v > watermark_threshold
|
||||
}
|
||||
|
||||
# Identify margin notes by finding entries with counts higher than a threshold (e.g., appearing on more than half of the pages)
|
||||
margin_note_threshold = len(result_dict) // 2
|
||||
margin_notes = {
|
||||
k: v for k, v in potential_margin_notes.items() if v > margin_note_threshold
|
||||
}
|
||||
|
||||
# Add watermark information to the result dictionary
|
||||
for page_id, blocks in result_dict.items():
|
||||
if page_id.startswith('page_'):
|
||||
for block_id, block_data in blocks.items():
|
||||
coordinates_text = (block_data['bbox'], block_data['text'])
|
||||
if coordinates_text in watermarks:
|
||||
block_data['is_watermark'] = 1
|
||||
else:
|
||||
block_data['is_watermark'] = 0
|
||||
|
||||
if coordinates_text in margin_notes:
|
||||
block_data['is_vertical_margin_note'] = 1
|
||||
else:
|
||||
block_data['is_vertical_margin_note'] = 0
|
||||
|
||||
return result_dict
|
||||
|
||||
|
||||
"""
|
||||
1. 当一个block里全部文字都不是dir=(1,0)，这个block整体去掉
|
||||
2. 当一个block里全部文字都是dir=(1,0)，但是每行只有一个字，这个block整体去掉。这个block必须出现在页面的四周，否则不去掉
|
||||
"""
|
||||
|
||||
|
||||
def __is_a_word(sentence):
|
||||
# 如果输入是中文并且长度为1，则返回True
|
||||
if re.fullmatch(r'[\u4e00-\u9fa5]', sentence):
|
||||
return True
|
||||
# 判断是否为单个英文单词或字符（包括ASCII标点）
|
||||
elif re.fullmatch(r'[a-zA-Z0-9]+', sentence) and len(sentence) <= 2:
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
def __get_text_color(num):
|
||||
"""获取字体的颜色RGB值."""
|
||||
blue = num & 255
|
||||
green = (num >> 8) & 255
|
||||
red = (num >> 16) & 255
|
||||
return red, green, blue
|
||||
|
||||
|
||||
def __is_empty_side_box(text_block):
|
||||
"""是否是边缘上的空白没有任何内容的block."""
|
||||
for line in text_block['lines']:
|
||||
for span in line['spans']:
|
||||
font_color = span['color']
|
||||
r, g, b = __get_text_color(font_color)
|
||||
if len(span['text'].strip()) > 0 and (r, g, b) != (255, 255, 255):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def remove_rotate_side_textblock(pymu_text_block, page_width, page_height):
|
||||
"""返回删除了垂直，水印，旋转的textblock 删除的内容打上tag返回."""
|
||||
removed_text_block = []
|
||||
|
||||
for i, block in enumerate(
|
||||
pymu_text_block
|
||||
): # 格式参考test/assets/papre/pymu_textblocks.json
|
||||
lines = block['lines']
|
||||
block_bbox = block['bbox']
|
||||
if not is_vbox_on_side(
|
||||
block_bbox, page_width, page_height, 0.2
|
||||
): # 保证这些box必须在页面的两边
|
||||
continue
|
||||
|
||||
if (
|
||||
all(
|
||||
[
|
||||
__is_a_word(line['spans'][0]['text'])
|
||||
for line in lines
|
||||
if len(line['spans']) > 0
|
||||
]
|
||||
)
|
||||
and len(lines) > 1
|
||||
and all([len(line['spans']) == 1 for line in lines])
|
||||
):
|
||||
is_box_valign = (
|
||||
(
|
||||
len(
|
||||
set(
|
||||
[
|
||||
int(line['spans'][0]['bbox'][0])
|
||||
for line in lines
|
||||
if len(line['spans']) > 0
|
||||
]
|
||||
)
|
||||
)
|
||||
== 1
|
||||
)
|
||||
and (
|
||||
len(
|
||||
[
|
||||
int(line['spans'][0]['bbox'][0])
|
||||
for line in lines
|
||||
if len(line['spans']) > 0
|
||||
]
|
||||
)
|
||||
> 1
|
||||
)
|
||||
) # 测试bbox在垂直方向是不是x0都相等，也就是在垂直方向排列.同时必须大于等于2个字
|
||||
|
||||
if is_box_valign:
|
||||
block['tag'] = VERTICAL_TEXT
|
||||
removed_text_block.append(block)
|
||||
continue
|
||||
|
||||
for line in lines:
|
||||
if line['dir'] != (1, 0):
|
||||
block['tag'] = ROTATE_TEXT
|
||||
removed_text_block.append(
|
||||
block
|
||||
) # 只要有一个line不是dir=(1,0)，就把整个block都删掉
|
||||
break
|
||||
|
||||
for block in removed_text_block:
|
||||
pymu_text_block.remove(block)
|
||||
|
||||
return pymu_text_block, removed_text_block
|
||||
|
||||
|
||||
def get_side_boundry(rotate_bbox, page_width, page_height):
|
||||
"""根据rotate_bbox，返回页面的左右正文边界."""
|
||||
left_x = 0
|
||||
right_x = page_width
|
||||
for x in rotate_bbox:
|
||||
box = x['bbox']
|
||||
if box[2] < page_width / 2:
|
||||
left_x = max(left_x, box[2])
|
||||
else:
|
||||
right_x = min(right_x, box[0])
|
||||
|
||||
return left_x + 1, right_x - 1
|
||||
|
||||
|
||||
def remove_side_blank_block(pymu_text_block, page_width, page_height):
|
||||
"""删除页面两侧的空白block."""
|
||||
removed_text_block = []
|
||||
|
||||
for i, block in enumerate(
|
||||
pymu_text_block
|
||||
): # 格式参考test/assets/papre/pymu_textblocks.json
|
||||
block_bbox = block['bbox']
|
||||
if not is_vbox_on_side(
|
||||
block_bbox, page_width, page_height, 0.2
|
||||
): # 保证这些box必须在页面的两边
|
||||
continue
|
||||
|
||||
if __is_empty_side_box(block):
|
||||
block['tag'] = EMPTY_SIDE_BLOCK
|
||||
removed_text_block.append(block)
|
||||
continue
|
||||
|
||||
for block in removed_text_block:
|
||||
pymu_text_block.remove(block)
|
||||
|
||||
return pymu_text_block, removed_text_block
|
||||
@@ -1,184 +0,0 @@
|
||||
"""
|
||||
从pdf里提取出来api给出的bbox,然后根据重叠情况做出取舍
|
||||
1. 首先去掉出现在图片上的bbox，图片包括表格和图片
|
||||
2. 然后去掉出现在文字blcok上的图片bbox
|
||||
"""
|
||||
|
||||
from magic_pdf.config.drop_tag import ON_IMAGE_TEXT, ON_TABLE_TEXT
|
||||
from magic_pdf.libs.boxbase import (_is_in, _is_in_or_part_overlap,
|
||||
_is_left_overlap)
|
||||
|
||||
|
||||
def resolve_bbox_overlap_conflict(images: list, tables: list, interline_equations: list, inline_equations: list,
|
||||
text_raw_blocks: list):
|
||||
"""
|
||||
text_raw_blocks结构是从pymupdf里直接取到的结构，具体样例参考test/assets/papre/pymu_textblocks.json
|
||||
当下采用一种粗暴的方式：
|
||||
1. 去掉图片上的公式
|
||||
2. 去掉table上的公式
|
||||
2. 图片和文字block部分重叠，首先丢弃图片
|
||||
3. 图片和图片重叠，修改图片的bbox，使得图片不重叠(暂时没这么做，先把图片都扔掉)
|
||||
4. 去掉文字bbox里位于图片、表格上的文字（一定要完全在图、表内部）
|
||||
5. 去掉表格上的文字
|
||||
"""
|
||||
text_block_removed = []
|
||||
images_backup = []
|
||||
|
||||
# 去掉位于图片上的文字block
|
||||
for image_box in images:
|
||||
for text_block in text_raw_blocks:
|
||||
text_bbox = text_block['bbox']
|
||||
if _is_in(text_bbox, image_box):
|
||||
text_block['tag'] = ON_IMAGE_TEXT
|
||||
text_block_removed.append(text_block)
|
||||
# 去掉table上的文字block
|
||||
for table_box in tables:
|
||||
for text_block in text_raw_blocks:
|
||||
text_bbox = text_block['bbox']
|
||||
if _is_in(text_bbox, table_box):
|
||||
text_block['tag'] = ON_TABLE_TEXT
|
||||
text_block_removed.append(text_block)
|
||||
|
||||
for text_block in text_block_removed:
|
||||
if text_block in text_raw_blocks:
|
||||
text_raw_blocks.remove(text_block)
|
||||
|
||||
# 第一步去掉在图片上出现的公式box
|
||||
temp = []
|
||||
for image_box in images:
|
||||
for eq1 in interline_equations:
|
||||
if _is_in_or_part_overlap(image_box, eq1[:4]):
|
||||
temp.append(eq1)
|
||||
for eq2 in inline_equations:
|
||||
if _is_in_or_part_overlap(image_box, eq2[:4]):
|
||||
temp.append(eq2)
|
||||
|
||||
for eq in temp:
|
||||
if eq in interline_equations:
|
||||
interline_equations.remove(eq)
|
||||
if eq in inline_equations:
|
||||
inline_equations.remove(eq)
|
||||
|
||||
# 第二步去掉在表格上出现的公式box
|
||||
temp = []
|
||||
for table_box in tables:
|
||||
for eq1 in interline_equations:
|
||||
if _is_in_or_part_overlap(table_box, eq1[:4]):
|
||||
temp.append(eq1)
|
||||
for eq2 in inline_equations:
|
||||
if _is_in_or_part_overlap(table_box, eq2[:4]):
|
||||
temp.append(eq2)
|
||||
|
||||
for eq in temp:
|
||||
if eq in interline_equations:
|
||||
interline_equations.remove(eq)
|
||||
if eq in inline_equations:
|
||||
inline_equations.remove(eq)
|
||||
|
||||
# 图片和文字重叠，丢掉图片
|
||||
for image_box in images:
|
||||
for text_block in text_raw_blocks:
|
||||
text_bbox = text_block['bbox']
|
||||
if _is_in_or_part_overlap(image_box, text_bbox):
|
||||
images_backup.append(image_box)
|
||||
break
|
||||
for image_box in images_backup:
|
||||
images.remove(image_box)
|
||||
|
||||
# 图片和图片重叠，两张都暂时不参与版面计算
|
||||
images_dup_index = []
|
||||
for i in range(len(images)):
|
||||
for j in range(i + 1, len(images)):
|
||||
if _is_in_or_part_overlap(images[i], images[j]):
|
||||
images_dup_index.append(i)
|
||||
images_dup_index.append(j)
|
||||
|
||||
dup_idx = set(images_dup_index)
|
||||
for img_id in dup_idx:
|
||||
images_backup.append(images[img_id])
|
||||
images[img_id] = None
|
||||
|
||||
images = [img for img in images if img is not None]
|
||||
|
||||
# 如果行间公式和文字block重叠，放到临时的数据里，防止这些文字box影响到layout计算。通过计算IOU合并行间公式和文字block
|
||||
# 对于这样的文本块删除，然后保留行间公式的大小不变。
|
||||
# 当计算完毕layout，这部分再合并回来
|
||||
text_block_removed_2 = []
|
||||
# for text_block in text_raw_blocks:
|
||||
# text_bbox = text_block["bbox"]
|
||||
# for eq in interline_equations:
|
||||
# ratio = calculate_overlap_area_2_minbox_area_ratio(text_bbox, eq[:4])
|
||||
# if ratio>0.05:
|
||||
# text_block['tag'] = "belong-to-interline-equation"
|
||||
# text_block_removed_2.append(text_block)
|
||||
# break
|
||||
|
||||
# for tb in text_block_removed_2:
|
||||
# if tb in text_raw_blocks:
|
||||
# text_raw_blocks.remove(tb)
|
||||
|
||||
# text_block_removed = text_block_removed + text_block_removed_2
|
||||
|
||||
return images, tables, interline_equations, inline_equations, text_raw_blocks, text_block_removed, images_backup, text_block_removed_2
|
||||
|
||||
|
||||
def check_text_block_horizontal_overlap(text_blocks: list, header, footer) -> bool:
|
||||
"""检查文本block之间的水平重叠情况，这种情况如果发生，那么这个pdf就不再继续处理了。 因为这种情况大概率发生了公式没有被检测出来。"""
|
||||
if len(text_blocks) == 0:
|
||||
return False
|
||||
|
||||
page_min_y = 0
|
||||
page_max_y = max(yy['bbox'][3] for yy in text_blocks)
|
||||
|
||||
def __max_y(lst: list):
|
||||
if len(lst) > 0:
|
||||
return max([item[1] for item in lst])
|
||||
return page_min_y
|
||||
|
||||
def __min_y(lst: list):
|
||||
if len(lst) > 0:
|
||||
return min([item[3] for item in lst])
|
||||
return page_max_y
|
||||
|
||||
clip_y0 = __max_y(header)
|
||||
clip_y1 = __min_y(footer)
|
||||
|
||||
txt_bboxes = []
|
||||
for text_block in text_blocks:
|
||||
bbox = text_block['bbox']
|
||||
if bbox[1] >= clip_y0 and bbox[3] <= clip_y1:
|
||||
txt_bboxes.append(bbox)
|
||||
|
||||
for i in range(len(txt_bboxes)):
|
||||
for j in range(i + 1, len(txt_bboxes)):
|
||||
if _is_left_overlap(txt_bboxes[i], txt_bboxes[j]) or _is_left_overlap(txt_bboxes[j], txt_bboxes[i]):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def check_useful_block_horizontal_overlap(useful_blocks: list) -> bool:
|
||||
"""检查文本block之间的水平重叠情况，这种情况如果发生，那么这个pdf就不再继续处理了。 因为这种情况大概率发生了公式没有被检测出来。"""
|
||||
if len(useful_blocks) == 0:
|
||||
return False
|
||||
|
||||
page_min_y = 0
|
||||
page_max_y = max(yy['bbox'][3] for yy in useful_blocks)
|
||||
|
||||
useful_bboxes = []
|
||||
for text_block in useful_blocks:
|
||||
bbox = text_block['bbox']
|
||||
if bbox[1] >= page_min_y and bbox[3] <= page_max_y:
|
||||
useful_bboxes.append(bbox)
|
||||
|
||||
for i in range(len(useful_bboxes)):
|
||||
for j in range(i + 1, len(useful_bboxes)):
|
||||
area_i = (useful_bboxes[i][2] - useful_bboxes[i][0]) * (useful_bboxes[i][3] - useful_bboxes[i][1])
|
||||
area_j = (useful_bboxes[j][2] - useful_bboxes[j][0]) * (useful_bboxes[j][3] - useful_bboxes[j][1])
|
||||
if _is_left_overlap(useful_bboxes[i], useful_bboxes[j]) or _is_left_overlap(useful_bboxes[j], useful_bboxes[i]):
|
||||
if area_i > area_j:
|
||||
return True, useful_bboxes[j], useful_bboxes[i]
|
||||
else:
|
||||
return True, useful_bboxes[i], useful_bboxes[j]
|
||||
|
||||
return False, None, None
|
||||
@@ -1,29 +0,0 @@
|
||||
def solve_inline_too_large_interval(pdf_info_dict: dict) -> dict: # text_block -> json中的preproc_block
|
||||
"""解决行内文本间距过大问题"""
|
||||
for i in range(len(pdf_info_dict)):
|
||||
|
||||
text_blocks = pdf_info_dict[f'page_{i}']['preproc_blocks']
|
||||
|
||||
for block in text_blocks:
|
||||
|
||||
x_pre_1, y_pre_1, x_pre_2, y_pre_2 = 0, 0, 0, 0
|
||||
|
||||
for line in block['lines']:
|
||||
|
||||
x_cur_1, y_cur_1, x_cur_2, y_cur_2 = line['bbox']
|
||||
# line_box = [x1, y1, x2, y2]
|
||||
if int(y_cur_1) == int(y_pre_1) and int(y_cur_2) == int(y_pre_2):
|
||||
# if len(line['spans']) == 1:
|
||||
line['spans'][0]['text'] = ' ' + line['spans'][0]['text']
|
||||
|
||||
x_pre_1, y_pre_1, x_pre_2, y_pre_2 = line['bbox']
|
||||
|
||||
return pdf_info_dict
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -1,12 +0,0 @@
|
||||
|
||||
"""
|
||||
统计处需要跨页、全局性的数据
|
||||
- 统计出字号从大到小
|
||||
- 正文区域占比最高的前5
|
||||
- 正文平均行间距
|
||||
- 正文平均字间距
|
||||
- 正文平均字符宽度
|
||||
- 正文平均字符高度
|
||||
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user