Skip to content

P
pdf-miner
Unverified Commit f5268688 authored by myhloli's avatar myhloli Committed by GitHub
Browse files
Options

Merge pull request #38 from papayalove/master 

更新了para_split_by_model
parents 07012dca f519f63d
Hide whitespace changes
Inline Side-by-side
Showing with 104 additions and 70 deletions
magic_pdf/para/para_split_by_model.py
View file @ f5268688
......@@ -3,7 +3,7 @@ import numpy as np
from loguru import logger
from magic_pdf.libs.boxbase import _is_in_or_part_overlap_with_area_ratio as is_in_layout
from magic_pdf.libs.ocr_content_type import ContentType
from magic_pdf.libs.ocr_content_type import ContentType, BlockType
from magic_pdf.model.magic_model import MagicModel
LINE_STOP_FLAG = ['.', '!', '?', '。', '！', '？', "：", ":", ")", "）", ";"]
......@@ -131,7 +131,7 @@ def __valign_lines(blocks, layout_bboxes):
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'])]
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:
continue
......@@ -192,12 +192,12 @@ def __align_text_in_layout(blocks, layout_bboxes):
"""
for layout in layout_bboxes:
lb = layout['layout_bbox']
blocks_in_layoutbox = [b for b in blocks if is_in_layout(b['bbox'], lb)]
blocks_in_layoutbox = [b for b in blocks if block["type"] == BlockType.Text and is_in_layout(b['bbox'], lb)]
if len(blocks_in_layoutbox) == 0:
continue
for block in blocks_in_layoutbox:
for line in block['lines']:
for line in block.get("lines", []):
x0, x1 = line['bbox'][0], line['bbox'][2]
if x0 < lb[0]:
line['bbox'][0] = lb[0]
......@@ -238,7 +238,7 @@ def __group_line_by_layout(blocks, layout_bboxes, lang="en"):
lines_group = []
blocks_group = []
for lyout in layout_bboxes:
lines = [line for block in blocks if is_in_layout(block['bbox'], lyout['layout_bbox']) for line in
lines = [line for block in blocks if block["type"] == BlockType.Text and is_in_layout(block['bbox'], lyout['layout_bbox']) for line in
block['lines']]
blocks = [block for block in blocks if is_in_layout(block['bbox'], lyout['layout_bbox'])]
......@@ -256,7 +256,7 @@ def __split_para_in_layoutbox2(lines_group, new_layout_bbox, lang="en", char_avg
def __split_para_in_layoutbox(lines_group, new_layout_bbox, text_blocks, lang="en", char_avg_len=10):
def __split_para_in_layoutbox(blocks_group, new_layout_bbox, text_blocks, lang="en", char_avg_len=10):
"""
lines_group 进行行分段——layout内部进行分段。lines_group内每个元素是一个Layoutbox内的所有行。
1. 先计算每个group的左右边界。
......@@ -266,28 +266,34 @@ def __split_para_in_layoutbox(lines_group, new_layout_bbox, text_blocks, lang="e
"""
list_info = [] # 这个layout最后是不是列表,记录每一个layout里是不是列表开头，列表结尾
layout_paras = []
right_tail_distance = 1.5 * char_avg_len
for blocks in blocks_group:
is_start_list = None
is_end_list = None
if len(blocks) == 0:
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
for lines in lines_group:
paras = []
lines = [line for block in blocks if
block["type"] == BlockType.Text for line in
block['lines']]
total_lines = len(lines)
if total_lines == 0:
continue # 0行无需处理
if total_lines == 1: # 1行无法分段。
layout_paras.append([lines])
if total_lines == 1:
list_info.append([False, False])
continue
"""在进入到真正的分段之前，要对文字块从统计维度进行对齐方式的探测，
对齐方式分为以下：
1. 左对齐的文本块(特点是左侧顶格，或者左侧不顶格但是右侧顶格的行数大于非顶格的行数，顶格的首字母有大写也有小写)
1) 右侧对齐的行，单独成一段
2) 中间对齐的行，按照字体/行高聚合成一段
2. 左对齐的列表块（其特点是左侧顶格的行数小于等于非顶格的行数，非定格首字母会有小写，顶格90%是大写。并且左侧顶格行数大于1，大于1是为了这种模式连续出现才能称之为列表）
这样的文本块，顶格的为一个段落开头，紧随其后非顶格的行属于这个段落。
对齐方式分为以下：
1. 左对齐的文本块(特点是左侧顶格，或者左侧不顶格但是右侧顶格的行数大于非顶格的行数，顶格的首字母有大写也有小写)
1) 右侧对齐的行，单独成一段
2) 中间对齐的行，按照字体/行高聚合成一段
2. 左对齐的列表块（其特点是左侧顶格的行数小于等于非顶格的行数，非定格首字母会有小写，顶格90%是大写。并且左侧顶格行数大于1，大于1是为了这种模式连续出现才能称之为列表）
这样的文本块，顶格的为一个段落开头，紧随其后非顶格的行属于这个段落。
"""
text_segments, list_start_line = __detect_list_lines(lines, new_layout_bbox, lang)
"""根据list_range，把lines分成几个部分
......@@ -299,16 +305,13 @@ def __split_para_in_layoutbox(lines_group, new_layout_bbox, text_blocks, lang="e
layout_list_info = [False, False] # 这个layout最后是不是列表,记录每一个layout里是不是列表开头，列表结尾
for content_type, start, end in text_segments:
if content_type == 'list':
if start == 0:
if start == 0 and is_start_list is None:
layout_list_info[0] = True
if end == total_lines - 1:
if end == total_lines - 1 and is_end_list is None:
layout_list_info[1] = True
# paras = __split_para_lines(lines, text_blocks)
list_info.append(layout_list_info)
# layout_paras.append(paras)
return list_info
......@@ -348,7 +351,7 @@ def __split_para_lines(lines: list, text_blocks: list) -> list:
def __connect_list_inter_layout(layout_paras, blocks_group, new_layout_bbox, layout_list_info, page_num, lang):
def __connect_list_inter_layout(blocks_group, new_layout_bbox, layout_list_info, page_num, lang):
"""
如果上个layout的最后一个段落是列表，下一个layout的第一个段落也是列表，那么将他们连接起来。 TODO 因为没有区分列表和段落，所以这个方法暂时不实现。
根据layout_list_info判断是不是列表。，下个layout的第一个段如果不是列表，那么看他们是否有几行都有相同的缩进。
......@@ -359,16 +362,16 @@ def __connect_list_inter_layout(layout_paras, blocks_group, new_layout_bbox, lay
for i in range(1, len(blocks_group)):
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]["lines"]
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]: # 前一个是列表结尾，后一个是非列表开头，此时检测是否有相同的缩进
if pre_layout_list_info[1] and not next_layout_list_info[0] and next_first_para["type"] == BlockType.Text: # 前一个是列表结尾，后一个是非列表开头，此时检测是否有相同的缩进
logger.info(f"连接page {page_num} 内的list")
# 向layout_paras[i] 寻找开头具有相同缩进的连续的行
may_list_lines = []
for j in range(len(next_paras)):
lines = next_paras[j]["lines"]
lines = next_paras[j].get("lines", [])
if len(lines) == 1: # 只可能是一行，多行情况再需要分析了
if lines[0]['bbox'][0] > __find_layout_bbox_by_line(lines[0]['bbox'], new_layout_bbox)[0]:
may_list_lines.append(lines[0])
......@@ -393,23 +396,25 @@ def __connect_list_inter_page(pre_page_paras, next_page_paras, pre_page_layout_b
"""
if len(pre_page_paras) == 0 or len(next_page_paras) == 0: # 0的时候最后的return 会出错
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]: # 前一个是列表结尾，后一个是非列表开头，此时检测是否有相同的缩进
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])
lines = next_page_paras[0][j]["lines"]
if len(lines) == 1: # 只可能是一行，多行情况再需要分析了
if lines[0]['bbox'][0] > __find_layout_bbox_by_line(lines[0]['bbox'], next_page_layout_bbox)[0]:
may_list_lines.append(lines[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)
#pre_page_paras[-1].append(may_list_lines)
pre_page_paras[-1][-1]["lines"].extend(may_list_lines)
next_page_paras[0] = next_page_paras[0][len(may_list_lines):]
return True
......@@ -438,14 +443,18 @@ def __connect_para_inter_layoutbox(layout_paras, blocks_group, new_layout_bbox,
connected_layout_paras = []
connected_layout_blocks = []
if len(blocks_group) == 0:
return connected_layout_paras
return connected_layout_blocks
connected_layout_paras.append(layout_paras[0])
#connected_layout_paras.append(layout_paras[0])
connected_layout_blocks.append(blocks_group[0])
for i in range(1, len(blocks_group)):
try:
if len(blocks_group[i]) == 0 or len(blocks_group[i - 1]) == 0: # TODO 考虑连接问题，
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
pre_last_line = blocks_group[i - 1][-1]["lines"][-1]
next_first_line = blocks_group[i][0]["lines"][0]
except Exception as e:
......@@ -456,7 +465,8 @@ def __connect_para_inter_layoutbox(layout_paras, blocks_group, new_layout_bbox,
next_first_line_text = ''.join([__get_span_text(span) for span in next_first_line['spans']])
next_first_line_type = next_first_line['spans'][0]['type']
if pre_last_line_type not in [TEXT, INLINE_EQUATION] or next_first_line_type not in [TEXT, INLINE_EQUATION]:
connected_layout_paras.append(layout_paras[i])
#connected_layout_paras.append(layout_paras[i])
connected_layout_blocks.append(blocks_group[i])
continue
pre_x2_max = __find_layout_bbox_by_line(pre_last_line['bbox'], new_layout_bbox)[2]
......@@ -467,17 +477,21 @@ def __connect_para_inter_layoutbox(layout_paras, blocks_group, new_layout_bbox,
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])
#connected_layout_paras[-1][-1].extend(layout_paras[i][0])
connected_layout_blocks[-1][-1]["lines"].extend(blocks_group[i][0]["lines"])
layout_paras[i].pop(0) # 删除后一个layout的第一个段落， 因为他已经被合并到前一个layout的最后一个段落了。
if len(layout_paras[i]) == 0:
layout_paras.pop(i)
else:
connected_layout_paras.append(layout_paras[i])
blocks_group[i][0]["lines"] = [] #删除后一个layout第一个段落中的lines，因为他已经被合并到前一个layout的最后一个段落了
blocks_group[i][0]["lines_deleted"] = True
# if len(layout_paras[i]) == 0:
# layout_paras.pop(i)
# else:
# connected_layout_paras.append(layout_paras[i])
connected_layout_blocks.append(blocks_group[i])
else:
"""连接段落条件不成立，将前一个layout的段落加入到结果中。"""
connected_layout_paras.append(layout_paras[i])
return connected_layout_paras
#connected_layout_paras.append(layout_paras[i])
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,
......@@ -492,8 +506,12 @@ def __connect_para_inter_page(pre_page_paras, next_page_paras, pre_page_layout_b
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_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
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']])
......@@ -515,7 +533,9 @@ def __connect_para_inter_page(pre_page_paras, next_page_paras, pre_page_layout_b
next_first_line['bbox'][0] == next_x0_min: # 前面一行沾满了整个行，并且没有结尾符号.下一行没有空白开头。
"""连接段落条件成立，将前一个layout的段落和后一个layout的段落连接。"""
pre_last_para.extend(next_first_para)
next_page_paras[0].pop(0) # 删除后一个页面的第一个段落， 因为他已经被合并到前一个页面的最后一个段落了。
#next_page_paras[0].pop(0) # 删除后一个页面的第一个段落， 因为他已经被合并到前一个页面的最后一个段落了。
next_page_paras[0][0]["lines"] = []
next_page_paras[0][0]["lines_deleted"] = True
return True
else:
return False
......@@ -556,18 +576,18 @@ def __connect_middle_align_text(page_paras, new_layout_bbox, page_num, lang, deb
layout_box = new_layout_bbox[layout_i]
single_line_paras_tag = []
for i in range(len(layout_para)):
single_line_paras_tag.append(len(layout_para[i]) == 1 and layout_para[i][0]['spans'][0]['type'] == TEXT)
#single_line_paras_tag.append(len(layout_para[i]) == 1 and layout_para[i][0]['spans'][0]['type'] == TEXT)
single_line_paras_tag.append(layout_para[i]['type'] == BlockType.Text and len(layout_para[i]["lines"]) == 1)
"""找出来连续的单行文本，如果连续行高度相同，那么合并为一个段落。"""
consecutive_single_line_indices = find_consecutive_true_regions(single_line_paras_tag)
if len(consecutive_single_line_indices) > 0:
index_offset = 0
#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']])
#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_mode:
......@@ -575,18 +595,23 @@ def __connect_middle_align_text(page_paras, new_layout_bbox, page_num, lang, deb
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]]
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 = [l[0] for l in layout_para[start:end + 1]]
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_mode:
logger.info(para_text)
layout_para[start:end + 1] = [merge_para]
index_offset -= end - start
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
#layout_para[start:end + 1] = [merge_para]
#index_offset -= end - start
return
......@@ -612,16 +637,16 @@ def __do_split_page(blocks, layout_bboxes, new_layout_bbox, text_blocks, page_nu
4. 图、表，目前独占一行，不考虑分段。
"""
lines_group, blocks_group = __group_line_by_layout(blocks, layout_bboxes, lang) # block内分段
layout_list_info = __split_para_in_layoutbox(lines_group, new_layout_bbox, text_blocks, lang) # layout内分段
blocks_group, page_list_info = __connect_list_inter_layout(layout_paras, blocks_group, new_layout_bbox, layout_list_info,
layout_list_info = __split_para_in_layoutbox(blocks_group, new_layout_bbox, text_blocks, 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_paras = __connect_para_inter_layoutbox(layout_paras2, blocks_group, new_layout_bbox, lang) # layout间链接段落
connected_layout_blocks = __connect_para_inter_layoutbox(blocks_group, new_layout_bbox, lang) # layout间链接段落
return connected_layout_paras, page_list_info
return connected_layout_blocks, page_list_info
def para_split(pdf_info_dict, debug_mode, magic_model: MagicModel, lang="en"):
def para_split_by_model(pdf_info_dict, debug_mode, magic_model: MagicModel, lang="en"):
new_layout_of_pages = [] # 数组的数组，每个元素是一个页面的layoutS
all_page_list_info = [] # 保存每个页面开头和结尾是否是列表
for page_num, page in pdf_info_dict.items():
......
magic_pdf/pdf_parse_by_ocr.py
View file @ f5268688
......@@ -11,6 +11,7 @@ from magic_pdf.libs.drop_tag import DropTag
from magic_pdf.libs.hash_utils import compute_md5
from magic_pdf.libs.ocr_content_type import ContentType
from magic_pdf.para.para_split import para_split
from magic_pdf.para.para_split_by_model import para_split_by_model
from magic_pdf.pre_proc.construct_page_dict import ocr_construct_page_component
from magic_pdf.pre_proc.detect_footer_by_model import parse_footers
from magic_pdf.pre_proc.detect_footnote import parse_footnotes_by_model
......@@ -208,7 +209,10 @@ def parse_pdf_by_ocr(
pdf_info_dict[f"page_{page_id}"] = page_info
"""分段"""
para_split(pdf_info_dict, debug_mode=debug_mode)
if debug_mode:
para_split_by_model(pdf_info_dict, debug_mode=debug_mode)
else:
para_split(pdf_info_dict, debug_mode=debug_mode)
"""dict转list"""
pdf_info_list = dict_to_list(pdf_info_dict)
......
magic_pdf/pdf_parse_by_ocr_v2.py
View file @ f5268688
......@@ -12,6 +12,8 @@ from magic_pdf.pre_proc.cut_image import ocr_cut_image_and_table
from magic_pdf.pre_proc.ocr_detect_all_bboxes import ocr_prepare_bboxes_for_layout_split
from magic_pdf.pre_proc.ocr_dict_merge import sort_blocks_by_layout, fill_spans_in_blocks, fix_block_spans
from magic_pdf.pre_proc.ocr_span_list_modify import remove_overlaps_min_spans, get_qa_need_list_v2
from magic_pdf.para.para_split import para_split
from magic_pdf.para.para_split_by_model import para_split_by_model
def parse_pdf_by_ocr(pdf_bytes,
......@@ -90,7 +92,10 @@ def parse_pdf_by_ocr(pdf_bytes,
pdf_info_dict[f"page_{page_id}"] = page_info
"""分段"""
pass
if debug_mode:
para_split_by_model(pdf_info_dict, debug_mode=debug_mode)
else:
para_split(pdf_info_dict, debug_mode=debug_mode)
"""dict转list"""
pdf_info_list = dict_to_list(pdf_info_dict)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment