para_split_v2.py.bak 38.1 KB
Newer Older
1
import copy
2
import re
3

liukaiwen's avatar
liukaiwen committed
4
5
import numpy as np
from loguru import logger
6
7
8
9
10
11
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
liukaiwen's avatar
liukaiwen committed
12

13
LINE_STOP_FLAG = ['.', '!', '?', '。', '!', '?', ':', ':', ')', ')', ';']
liukaiwen's avatar
liukaiwen committed
14
15
16
INLINE_EQUATION = ContentType.InlineEquation
INTERLINE_EQUATION = ContentType.InterlineEquation
TEXT = ContentType.Text
17
debug_able = False
liukaiwen's avatar
liukaiwen committed
18
19
20
21
22
23
24
25
26
27
28


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):
29
    global debug_able
liukaiwen's avatar
liukaiwen committed
30
31
32
33
34
    """
    探测是否包含了列表,并且把列表的行分开.
    这样的段落特点是,顶格字母大写/数字,紧跟着几行缩进的。缩进的行首字母含小写的。
    """

liukaiwen's avatar
liukaiwen committed
35
36
37
38
39
    def find_repeating_patterns2(lst):
        indices = []
        ones_indices = []
        i = 0
        while i < len(lst):  # Loop through the entire list
40
41
42
            if (
                lst[i] == 1
            ):  # If we encounter a '1', we might be at the start of a pattern
liukaiwen's avatar
liukaiwen committed
43
44
45
46
47
48
49
50
                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
51
                if len(ones_in_this_interval) > 1 or (
52
53
54
55
                    start < len(lst) - 1
                    and ones_in_this_interval
                    and lst[start + 1] in [2, 3]
                ):
liukaiwen's avatar
liukaiwen committed
56
57
58
59
60
                    indices.append((start, i - 1))
                    ones_indices.append(ones_in_this_interval)
            else:
                i += 1
        return indices, ones_indices
61

liukaiwen's avatar
liukaiwen committed
62
63
64
65
66
67
68
69
70
71
72
73
    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
                # 验证下一个序列是否符合条件
74
75
76
77
78
79
                if (
                    i < len(lst) - 1
                    and lst[i] == 1
                    and lst[i + 1] in [2, 3]
                    and lst[i - 1] in [2, 3]
                ):
liukaiwen's avatar
liukaiwen committed
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
                    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
116
117
118
119
120
121
122
123
124
125
126
127

    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)
128
    for l in lines:  # noqa: E741
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
        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)
liukaiwen's avatar
liukaiwen committed
145
            else:
146
147
                if first_char.isupper():
                    line_fea_encode.append(2)
liukaiwen's avatar
liukaiwen committed
148
                else:
149
150
151
152
153
154
155
                    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:
156
            logger.info(f'发现了列表,列表行数:{list_indice}, {list_start_idx}')
157
158

    # TODO check一下这个特列表里缩进的行左侧是不是对齐的。
159

160
161
162
163
164
    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:
165
                        logger.info(f'列表行的第{i}行不是顶格的')
166
167
                    break
        else:
168
            if debug_able:
169
                logger.info(f'列表行的第{start}到第{end}行是列表')
liukaiwen's avatar
liukaiwen committed
170

171
    return split_indices(total_lines, list_indice), list_start_idx
liukaiwen's avatar
liukaiwen committed
172

173

174
def cluster_line_x(lines: list) -> dict:
175
    """对一个block内所有lines的bbox的x0聚类."""
176
177
178
179
180
    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_)
181
    # x1_lst = np.array([[line['bbox'][2], 0] for line in lines])
182
    x0_2_new_val = {}  # 存储旧值对应的新值映射
183
    min_x0 = round(lines[0]['bbox'][0])
184
185
186
187
188
189
    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])
190
191
192
        x0_2_new_val.update(
            {round(raw_val): round(x0_new_val) for raw_val in x0_raw_val}
        )
193
194
195
196
        if x0_new_val < min_x0:
            min_x0 = x0_new_val
    return x0_2_new_val, min_x0

197

198
199
200
201
202
203
204
def if_match_reference_list(text: str) -> bool:
    pattern = re.compile(r'^\d+\..*')
    if pattern.match(text):
        return True
    else:
        return False

liukaiwen's avatar
liukaiwen committed
205
206

def __valign_lines(blocks, layout_bboxes):
207
208
209
    """在一个layoutbox内对齐行的左侧和右侧。 扫描行的左侧和右侧,如果x0,
    x1差距不超过一个阈值,就强行对齐到所处layout的左右两侧(和layout有一段距离)。
    3是个经验值,TODO,计算得来,可以设置为1.5个正文字符。"""
liukaiwen's avatar
liukaiwen committed
210
211
212
213

    min_distance = 3
    min_sample = 2
    new_layout_bboxes = []
214
215
    # add bbox_fs for para split calculation
    for block in blocks:
216
        block['bbox_fs'] = copy.deepcopy(block['bbox'])
liukaiwen's avatar
liukaiwen committed
217
    for layout_box in layout_bboxes:
218
219
220
221
222
223
224
        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:
liukaiwen's avatar
liukaiwen committed
225
            new_layout_bboxes.append(layout_box['layout_bbox'])
liukaiwen's avatar
liukaiwen committed
226
227
            continue

228
229
230
231
232
233
234
235
236
237
238
239
240
241
        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']
            ]
        )
liukaiwen's avatar
liukaiwen committed
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
        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:
276
277
278
279
280
281
282
            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']]),
                ]
liukaiwen's avatar
liukaiwen committed
283
        """新计算layout的bbox,因为block的bbox变了。"""
284
285
286
287
        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])
liukaiwen's avatar
liukaiwen committed
288
289
290
291
292
293
        new_layout_bboxes.append([layout_x0, layout_y0, layout_x1, layout_y1])

    return new_layout_bboxes


def __align_text_in_layout(blocks, layout_bboxes):
294
    """由于ocr出来的line,有时候会在前后有一段空白,这个时候需要对文本进行对齐,超出的部分被layout左右侧截断。"""
liukaiwen's avatar
liukaiwen committed
295
296
    for layout in layout_bboxes:
        lb = layout['layout_bbox']
297
298
299
300
301
        blocks_in_layoutbox = [
            block
            for block in blocks
            if block['type'] == BlockType.Text and is_in_layout(block['bbox'], lb)
        ]
liukaiwen's avatar
liukaiwen committed
302
303
304
305
        if len(blocks_in_layoutbox) == 0:
            continue

        for block in blocks_in_layoutbox:
306
            for line in block.get('lines', []):
liukaiwen's avatar
liukaiwen committed
307
308
309
310
311
312
313
314
                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):
315
    """不分语言的,对文本进行预处理."""
liukaiwen's avatar
liukaiwen committed
316
317
318
319
320
321
322
323
    # __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):
324
    """对中文文本进行分段预处理."""
liukaiwen's avatar
liukaiwen committed
325
326
327
328
    pass


def __pre_proc_en_blocks(blocks, layout_bboxes):
329
    """对英文文本进行分段预处理."""
liukaiwen's avatar
liukaiwen committed
330
331
332
    pass


liukaiwen's avatar
liukaiwen committed
333
def __group_line_by_layout(blocks, layout_bboxes):
334
    """每个layout内的行进行聚合."""
liukaiwen's avatar
liukaiwen committed
335
336
337
    # 因为只是一个block一行目前, 一个block就是一个段落
    blocks_group = []
    for lyout in layout_bboxes:
338
339
340
341
342
        blocks_in_layout = [
            block
            for block in blocks
            if is_in_layout(block.get('bbox_fs', None), lyout['layout_bbox'])
        ]
liukaiwen's avatar
liukaiwen committed
343
        blocks_group.append(blocks_in_layout)
liukaiwen's avatar
liukaiwen committed
344
    return blocks_group
liukaiwen's avatar
liukaiwen committed
345
346


347
def __split_para_in_layoutbox(blocks_group, new_layout_bbox, lang='en'):
liukaiwen's avatar
liukaiwen committed
348
349
350
351
352
353
354
355
356
    """
    lines_group 进行行分段——layout内部进行分段。lines_group内每个元素是一个Layoutbox内的所有行。
    1. 先计算每个group的左右边界。
    2. 然后根据行末尾特征进行分段。
        末尾特征:以句号等结束符结尾。并且距离右侧边界有一定距离。
        且下一行开头不留空白。

    """
    list_info = []  # 这个layout最后是不是列表,记录每一个layout里是不是列表开头,列表结尾
liukaiwen's avatar
liukaiwen committed
357
358
359
360
    for blocks in blocks_group:
        is_start_list = None
        is_end_list = None
        if len(blocks) == 0:
liukaiwen's avatar
liukaiwen committed
361
            list_info.append([False, False])
liukaiwen's avatar
liukaiwen committed
362
            continue
363
        if blocks[0]['type'] != BlockType.Text and blocks[-1]['type'] != BlockType.Text:
liukaiwen's avatar
liukaiwen committed
364
365
            list_info.append([False, False])
            continue
366
        if blocks[0]['type'] != BlockType.Text:
liukaiwen's avatar
liukaiwen committed
367
            is_start_list = False
368
        if blocks[-1]['type'] != BlockType.Text:
liukaiwen's avatar
liukaiwen committed
369
            is_end_list = False
liukaiwen's avatar
liukaiwen committed
370

371
372
373
374
375
376
        lines = [
            line
            for block in blocks
            if block['type'] == BlockType.Text
            for line in block['lines']
        ]
liukaiwen's avatar
liukaiwen committed
377
        total_lines = len(lines)
liukaiwen's avatar
liukaiwen committed
378
        if total_lines == 1 or total_lines == 0:
liukaiwen's avatar
liukaiwen committed
379
380
381
            list_info.append([False, False])
            continue
        """在进入到真正的分段之前,要对文字块从统计维度进行对齐方式的探测,
liukaiwen's avatar
liukaiwen committed
382
383
384
385
386
387
                    对齐方式分为以下:
                    1. 左对齐的文本块(特点是左侧顶格,或者左侧不顶格但是右侧顶格的行数大于非顶格的行数,顶格的首字母有大写也有小写)
                        1) 右侧对齐的行,单独成一段
                        2) 中间对齐的行,按照字体/行高聚合成一段
                    2. 左对齐的列表块(其特点是左侧顶格的行数小于等于非顶格的行数,非定格首字母会有小写,顶格90%是大写。并且左侧顶格行数大于1,大于1是为了这种模式连续出现才能称之为列表)
                        这样的文本块,顶格的为一个段落开头,紧随其后非顶格的行属于这个段落。
liukaiwen's avatar
liukaiwen committed
388
        """
389
390
391
        text_segments, list_start_line = __detect_list_lines(
            lines, new_layout_bbox, lang
        )
liukaiwen's avatar
liukaiwen committed
392
393
394
        """根据list_range,把lines分成几个部分

        """
liukaiwen's avatar
liukaiwen committed
395
396
        for list_start in list_start_line:
            if len(list_start) > 1:
397
                for i in range(0, len(list_start)):
liukaiwen's avatar
liukaiwen committed
398
                    index = list_start[i] - 1
399
                    if index >= 0:
400
401
402
403
404
405
406
407
408
409
410
                        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里是不是列表开头,列表结尾
liukaiwen's avatar
liukaiwen committed
411
412
        for content_type, start, end in text_segments:
            if content_type == 'list':
liukaiwen's avatar
liukaiwen committed
413
                if start == 0 and is_start_list is None:
liukaiwen's avatar
liukaiwen committed
414
                    layout_list_info[0] = True
liukaiwen's avatar
liukaiwen committed
415
                if end == total_lines - 1 and is_end_list is None:
liukaiwen's avatar
liukaiwen committed
416
417
418
419
420
421
422
423
424
425
426
                    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:
427
        spans_types = [span['type'] for span in line]
liukaiwen's avatar
liukaiwen committed
428
429
430
431
432
433
434
435
436
437
438
439
        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:
440
        block_bbox = block['bbox']
liukaiwen's avatar
liukaiwen committed
441
442
        para = []
        for line in text_lines:
443
            bbox = line['bbox']
liukaiwen's avatar
liukaiwen committed
444
445
446
447
448
            if is_in_layout(bbox, block_bbox):
                para.append(line)
        if len(para) > 0:
            text_paras.append(para)
    paras = other_paras.extend(text_paras)
449
    paras_sorted = sorted(paras, key=lambda x: x[0]['bbox'][1])
liukaiwen's avatar
liukaiwen committed
450
451
452
    return paras_sorted


453
454
455
def __connect_list_inter_layout(
    blocks_group, new_layout_bbox, layout_list_info, page_num, lang
):
456
    global debug_able
liukaiwen's avatar
liukaiwen committed
457
458
459
460
461
462
463
464
    """
    如果上个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)):
465
        if len(blocks_group[i]) == 0 or len(blocks_group[i - 1]) == 0:
liukaiwen's avatar
liukaiwen committed
466
            continue
liukaiwen's avatar
liukaiwen committed
467
468
        pre_layout_list_info = layout_list_info[i - 1]
        next_layout_list_info = layout_list_info[i]
469
        pre_last_para = blocks_group[i - 1][-1].get('lines', [])
liukaiwen's avatar
liukaiwen committed
470
471
472
        next_paras = blocks_group[i]
        next_first_para = next_paras[0]

473
474
475
476
477
        if (
            pre_layout_list_info[1]
            and not next_layout_list_info[0]
            and next_first_para['type'] == BlockType.Text
        ):  # 前一个是列表结尾,后一个是非列表开头,此时检测是否有相同的缩进
478
            if debug_able:
479
                logger.info(f'连接page {page_num} 内的list')
liukaiwen's avatar
liukaiwen committed
480
481
            # 向layout_paras[i] 寻找开头具有相同缩进的连续的行
            may_list_lines = []
482
            lines = next_first_para.get('lines', [])
483
484

            for line in lines:
485
486
487
488
                if (
                    line['bbox'][0]
                    > __find_layout_bbox_by_line(line['bbox'], new_layout_bbox)[0]
                ):
489
                    may_list_lines.append(line)
liukaiwen's avatar
liukaiwen committed
490
491
492
                else:
                    break
            # 如果这些行的缩进是相等的,那么连到上一个layout的最后一个段落上。
493
494
495
496
            if (
                len(may_list_lines) > 0
                and len(set([x['bbox'][0] for x in may_list_lines])) == 1
            ):
liukaiwen's avatar
liukaiwen committed
497
                pre_last_para.extend(may_list_lines)
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
                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 会出错
liukaiwen's avatar
liukaiwen committed
524
        return False
liukaiwen's avatar
liukaiwen committed
525
526
    if len(pre_page_paras[-1]) == 0 or len(next_page_paras[0]) == 0:
        return False
527
528
529
530
    if (
        pre_page_paras[-1][-1]['type'] != BlockType.Text
        or next_page_paras[0][0]['type'] != BlockType.Text
    ):
liukaiwen's avatar
liukaiwen committed
531
        return False
532
533
534
    if (
        pre_page_list_info[1] and not next_page_list_info[0]
    ):  # 前一个是列表结尾,后一个是非列表开头,此时检测是否有相同的缩进
535
        if debug_able:
536
            logger.info(f'连接page {page_num} 内的list')
liukaiwen's avatar
liukaiwen committed
537
538
        # 向layout_paras[i] 寻找开头具有相同缩进的连续的行
        may_list_lines = []
539
        next_page_first_para = next_page_paras[0][0]
540
541
        if next_page_first_para['type'] == BlockType.Text:
            lines = next_page_first_para['lines']
542
            for line in lines:
543
544
545
546
                if (
                    line['bbox'][0]
                    > __find_layout_bbox_by_line(line['bbox'], next_page_layout_bbox)[0]
                ):
547
                    may_list_lines.append(line)
liukaiwen's avatar
liukaiwen committed
548
549
550
                else:
                    break
        # 如果这些行的缩进是相等的,那么连到上一个layout的最后一个段落上。
551
552
553
554
        if (
            len(may_list_lines) > 0
            and len(set([x['bbox'][0] for x in may_list_lines])) == 1
        ):
555
            # pre_page_paras[-1].append(may_list_lines)
liukaiwen's avatar
liukaiwen committed
556
557
            # 下一页合并到上一页最后一段,打一个cross_page的标签
            for line in may_list_lines:
558
559
560
561
562
563
                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) :
            ]
liukaiwen's avatar
liukaiwen committed
564
565
566
567
568
569
            return True

    return False


def __find_layout_bbox_by_line(line_bbox, layout_bboxes):
570
    """根据line找到所在的layout."""
liukaiwen's avatar
liukaiwen committed
571
572
573
574
575
576
    for layout in layout_bboxes:
        if is_in_layout(line_bbox, layout):
            return layout
    return None


liukaiwen's avatar
liukaiwen committed
577
def __connect_para_inter_layoutbox(blocks_group, new_layout_bbox):
liukaiwen's avatar
liukaiwen committed
578
579
580
581
582
583
584
585
586
587
    """
    layout之间进行分段。
    主要是计算前一个layOut的最后一行和后一个layout的第一行是否可以连接。
    连接的条件需要同时满足:
    1. 上一个layout的最后一行沾满整个行。并且没有结尾符号。
    2. 下一行开头不留空白。

    """
    connected_layout_blocks = []
    if len(blocks_group) == 0:
liukaiwen's avatar
liukaiwen committed
588
        return connected_layout_blocks
liukaiwen's avatar
liukaiwen committed
589
590
591
592

    connected_layout_blocks.append(blocks_group[0])
    for i in range(1, len(blocks_group)):
        try:
liukaiwen's avatar
liukaiwen committed
593
594
            if len(blocks_group[i]) == 0:
                continue
595
            if len(blocks_group[i - 1]) == 0:  # TODO 考虑连接问题,
liukaiwen's avatar
liukaiwen committed
596
                connected_layout_blocks.append(blocks_group[i])
liukaiwen's avatar
liukaiwen committed
597
                continue
liukaiwen's avatar
liukaiwen committed
598
            # text类型的段才需要考虑layout间的合并
599
600
601
602
            if (
                blocks_group[i - 1][-1]['type'] != BlockType.Text
                or blocks_group[i][0]['type'] != BlockType.Text
            ):
liukaiwen's avatar
liukaiwen committed
603
604
                connected_layout_blocks.append(blocks_group[i])
                continue
605
606
607
608
            if (
                len(blocks_group[i - 1][-1]['lines']) == 0
                or len(blocks_group[i][0]['lines']) == 0
            ):
609
610
                connected_layout_blocks.append(blocks_group[i])
                continue
611
612
613
614
            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')
liukaiwen's avatar
liukaiwen committed
615
            continue
616
617
618
        pre_last_line_text = ''.join(
            [__get_span_text(span) for span in pre_last_line['spans']]
        )
liukaiwen's avatar
liukaiwen committed
619
        pre_last_line_type = pre_last_line['spans'][-1]['type']
620
621
622
        next_first_line_text = ''.join(
            [__get_span_text(span) for span in next_first_line['spans']]
        )
liukaiwen's avatar
liukaiwen committed
623
        next_first_line_type = next_first_line['spans'][0]['type']
624
625
626
627
        if pre_last_line_type not in [
            TEXT,
            INLINE_EQUATION,
        ] or next_first_line_type not in [TEXT, INLINE_EQUATION]:
liukaiwen's avatar
liukaiwen committed
628
            connected_layout_blocks.append(blocks_group[i])
liukaiwen's avatar
liukaiwen committed
629
            continue
liukaiwen's avatar
liukaiwen committed
630
        pre_layout = __find_layout_bbox_by_line(pre_last_line['bbox'], new_layout_bbox)
631
632
633
        next_layout = __find_layout_bbox_by_line(
            next_first_line['bbox'], new_layout_bbox
        )
liukaiwen's avatar
liukaiwen committed
634

liukaiwen's avatar
liukaiwen committed
635
636
        pre_x2_max = pre_layout[2] if pre_layout else -1
        next_x0_min = next_layout[0] if next_layout else -1
liukaiwen's avatar
liukaiwen committed
637
638
639

        pre_last_line_text = pre_last_line_text.strip()
        next_first_line_text = next_first_line_text.strip()
640
641
642
643
644
645
        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
        ):  # 前面一行沾满了整个行,并且没有结尾符号.下一行没有空白开头。
liukaiwen's avatar
liukaiwen committed
646
            """连接段落条件成立,将前一个layout的段落和后一个layout的段落连接。"""
647
648
649
650
651
            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
liukaiwen's avatar
liukaiwen committed
652
653
654
655
656
            # 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])
liukaiwen's avatar
liukaiwen committed
657
658
        else:
            """连接段落条件不成立,将前一个layout的段落加入到结果中。"""
liukaiwen's avatar
liukaiwen committed
659
660
            connected_layout_blocks.append(blocks_group[i])
    return connected_layout_blocks
liukaiwen's avatar
liukaiwen committed
661
662


663
664
665
666
667
668
669
670
def __connect_para_inter_page(
    pre_page_paras,
    next_page_paras,
    pre_page_layout_bbox,
    next_page_layout_bbox,
    page_num,
    lang,
):
liukaiwen's avatar
liukaiwen committed
671
672
673
674
675
676
677
    """
    连接起来相邻两个页面的段落——前一个页面最后一个段落和后一个页面的第一个段落。
    是否可以连接的条件:
    1. 前一个页面的最后一个段落最后一行沾满整个行。并且没有结尾符号。
    2. 后一个页面的第一个段落第一行没有空白开头。
    """
    # 有的页面可能压根没有文字
678
679
680
681
682
683
    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里?
liukaiwen's avatar
liukaiwen committed
684
        return False
liukaiwen's avatar
liukaiwen committed
685
686
    pre_last_block = pre_page_paras[-1][-1]
    next_first_block = next_page_paras[0][0]
687
688
689
690
    if (
        pre_last_block['type'] != BlockType.Text
        or next_first_block['type'] != BlockType.Text
    ):
liukaiwen's avatar
liukaiwen committed
691
        return False
692
    if len(pre_last_block['lines']) == 0 or len(next_first_block['lines']) == 0:
liukaiwen's avatar
liukaiwen committed
693
        return False
694
695
    pre_last_para = pre_last_block['lines']
    next_first_para = next_first_block['lines']
liukaiwen's avatar
liukaiwen committed
696
697
    pre_last_line = pre_last_para[-1]
    next_first_line = next_first_para[0]
698
699
700
    pre_last_line_text = ''.join(
        [__get_span_text(span) for span in pre_last_line['spans']]
    )
liukaiwen's avatar
liukaiwen committed
701
    pre_last_line_type = pre_last_line['spans'][-1]['type']
702
703
704
    next_first_line_text = ''.join(
        [__get_span_text(span) for span in next_first_line['spans']]
    )
liukaiwen's avatar
liukaiwen committed
705
706
    next_first_line_type = next_first_line['spans'][0]['type']

707
708
709
710
711
712
713
    if pre_last_line_type not in [
        TEXT,
        INLINE_EQUATION,
    ] or next_first_line_type not in [
        TEXT,
        INLINE_EQUATION,
    ]:  # TODO,真的要做好,要考虑跨table, image, 行间的情况
liukaiwen's avatar
liukaiwen committed
714
715
716
        # 不是文本,不连接
        return False

717
718
719
    pre_x2_max_bbox = __find_layout_bbox_by_line(
        pre_last_line['bbox'], pre_page_layout_bbox
    )
liukaiwen's avatar
liukaiwen committed
720
721
    if not pre_x2_max_bbox:
        return False
722
723
724
    next_x0_min_bbox = __find_layout_bbox_by_line(
        next_first_line['bbox'], next_page_layout_bbox
    )
liukaiwen's avatar
liukaiwen committed
725
726
727
728
729
    if not next_x0_min_bbox:
        return False

    pre_x2_max = pre_x2_max_bbox[2]
    next_x0_min = next_x0_min_bbox[0]
liukaiwen's avatar
liukaiwen committed
730
731
732

    pre_last_line_text = pre_last_line_text.strip()
    next_first_line_text = next_first_line_text.strip()
733
734
735
736
737
    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
    ):  # 前面一行沾满了整个行,并且没有结尾符号.下一行没有空白开头。
liukaiwen's avatar
liukaiwen committed
738
        """连接段落条件成立,将前一个layout的段落和后一个layout的段落连接。"""
liukaiwen's avatar
liukaiwen committed
739
740
        # 下一页合并到上一页最后一段,打一个cross_page的标签
        for line in next_first_para:
741
742
            for span in line['spans']:
                span[CROSS_PAGE] = True  # noqa: F405
liukaiwen's avatar
liukaiwen committed
743
        pre_last_para.extend(next_first_para)
liukaiwen's avatar
liukaiwen committed
744

745
        # next_page_paras[0].pop(0)  # 删除后一个页面的第一个段落, 因为他已经被合并到前一个页面的最后一个段落了。
746
747
        next_page_paras[0][0]['lines'] = []
        next_page_paras[0][0][LINES_DELETED] = True  # noqa: F405
liukaiwen's avatar
liukaiwen committed
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
        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


776
777
def __connect_middle_align_text(page_paras, new_layout_bbox, page_num, lang):
    global debug_able
liukaiwen's avatar
liukaiwen committed
778
779
780
781
782
783
784
785
786
787
788
    """
    找出来中间对齐的连续单行文本,如果连续行高度相同,那么合并为一个段落。
    一个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)):
789
            # single_line_paras_tag.append(len(layout_para[i]) == 1 and layout_para[i][0]['spans'][0]['type'] == TEXT)
790
791
792
793
            single_line_paras_tag.append(
                layout_para[i]['type'] == BlockType.Text
                and len(layout_para[i]['lines']) == 1
            )
liukaiwen's avatar
liukaiwen committed
794
        """找出来连续的单行文本,如果连续行高度相同,那么合并为一个段落。"""
795
796
797
        consecutive_single_line_indices = find_consecutive_true_regions(
            single_line_paras_tag
        )
liukaiwen's avatar
liukaiwen committed
798
        if len(consecutive_single_line_indices) > 0:
799
            """检查这些行是否是高度相同的,居中的."""
liukaiwen's avatar
liukaiwen committed
800
            for start, end in consecutive_single_line_indices:
801
802
                # start += index_offset
                # end += index_offset
803
804
805
806
807
808
809
810
811
812
813
814
815
                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:
liukaiwen's avatar
liukaiwen committed
816
                    pass
817
                if debug_able:
liukaiwen's avatar
liukaiwen committed
818
819
820
                    logger.info(line_hi.std())

                if line_hi.std() < 2:
821
822
823
824
825
826
827
828
829
                    """行高度相同,那么判断是否居中."""
                    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]
                    ]
liukaiwen's avatar
liukaiwen committed
830
                    layout_center = (layout_box[0] + layout_box[2]) / 2
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
                    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']
                            ]
                        )
851
                        if debug_able:
liukaiwen's avatar
liukaiwen committed
852
                            logger.info(para_text)
853
                        layout_para[start]['lines'] = merge_para
854
                        for i_para in range(start + 1, end + 1):
855
856
                            layout_para[i_para]['lines'] = []
                            layout_para[i_para][LINES_DELETED] = True  # noqa: F405
857
                        # layout_para[start:end + 1] = [merge_para]
liukaiwen's avatar
liukaiwen committed
858

859
                        # index_offset -= end - start
liukaiwen's avatar
liukaiwen committed
860
861
862
863
864

    return


def __merge_signle_list_text(page_paras, new_layout_bbox, page_num, lang):
865
    """找出来连续的单行文本,如果首行顶格,接下来的几个单行段落缩进对齐,那么合并为一个段落。"""
liukaiwen's avatar
liukaiwen committed
866
867
868
869

    pass


liukaiwen's avatar
liukaiwen committed
870
def __do_split_page(blocks, layout_bboxes, new_layout_bbox, page_num, lang):
871
    """根据line和layout情况进行分段 先实现一个根据行末尾特征分段的简单方法。"""
liukaiwen's avatar
liukaiwen committed
872
873
874
875
876
877
878
    """
    算法思路:
    1. 扫描layout里每一行,找出来行尾距离layout有边界有一定距离的行。
    2. 从上述行中找到末尾是句号等可作为断行标志的行。
    3. 参照上述行尾特征进行分段。
    4. 图、表,目前独占一行,不考虑分段。
    """
liukaiwen's avatar
liukaiwen committed
879
    blocks_group = __group_line_by_layout(blocks, layout_bboxes)  # block内分段
880
881
882
883
884
885
886
887
888
    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间链接段落
liukaiwen's avatar
liukaiwen committed
889

liukaiwen's avatar
liukaiwen committed
890
    return connected_layout_blocks, page_list_info
liukaiwen's avatar
liukaiwen committed
891
892


893
def para_split(pdf_info_dict, debug_mode, lang='en'):
894
895
    global debug_able
    debug_able = debug_mode
liukaiwen's avatar
liukaiwen committed
896
897
898
    new_layout_of_pages = []  # 数组的数组,每个元素是一个页面的layoutS
    all_page_list_info = []  # 保存每个页面开头和结尾是否是列表
    for page_num, page in pdf_info_dict.items():
899
        blocks = copy.deepcopy(page['preproc_blocks'])
liukaiwen's avatar
liukaiwen committed
900
901
902
        layout_bboxes = page['layout_bboxes']
        new_layout_bbox = __common_pre_proc(blocks, layout_bboxes)
        new_layout_of_pages.append(new_layout_bbox)
903
904
905
        splited_blocks, page_list_info = __do_split_page(
            blocks, layout_bboxes, new_layout_bbox, page_num, lang
        )
liukaiwen's avatar
liukaiwen committed
906
907
908
909
910
911
912
913
914
915
916
917
918
        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]

919
920
921
922
923
924
925
926
        is_conn = __connect_para_inter_page(
            pre_page_paras,
            next_page_paras,
            pre_page_layout_bbox,
            next_page_layout_bbox,
            page_num,
            lang,
        )
927
        if debug_able:
liukaiwen's avatar
liukaiwen committed
928
            if is_conn:
929
930
931
932
933
934
935
936
937
938
939
940
                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,
        )
941
        if debug_able:
liukaiwen's avatar
liukaiwen committed
942
            if is_list_conn:
943
                logger.info(f'连接了第{page_num - 1}页和第{page_num}页的列表段落')
liukaiwen's avatar
liukaiwen committed
944
945
946
947
948
949
950
951

    """接下来可能会漏掉一些特别的一些可以合并的内容,对他们进行段落连接
    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]
952
        __connect_middle_align_text(page_paras, new_layout_bbox, page_num, lang)
liukaiwen's avatar
liukaiwen committed
953
954
955
956
957
958
        __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]
959
        page['para_blocks'] = page_blocks