para_split_v3.py 16.5 KB
Newer Older
1
2
import copy

3
4
from loguru import logger

5
6
from magic_pdf.config.constants import CROSS_PAGE, LINES_DELETED
from magic_pdf.config.ocr_content_type import BlockType, ContentType
7
from magic_pdf.libs.language import detect_lang
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

LINE_STOP_FLAG = (
    '.',
    '!',
    '?',
    '。',
    '!',
    '?',
    ')',
    ')',
    '"',
    '”',
    ':',
    ':',
    ';',
    ';',
)
25
26
27
28
LIST_END_FLAG = ('.', '。', ';', ';')


class ListLineTag:
29
30
    IS_LIST_START_LINE = 'is_list_start_line'
    IS_LIST_END_LINE = 'is_list_end_line'
31
32
33


def __process_blocks(blocks):
34
35
36
    # 对所有block预处理
    # 1.通过title和interline_equation将block分组
    # 2.bbox边界根据line信息重置
37
38
39
40
41
42
43
44
45

    result = []
    current_group = []

    for i in range(len(blocks)):
        current_block = blocks[i]

        # 如果当前块是 text 类型
        if current_block['type'] == 'text':
46
47
48
49
50
51
52
53
            current_block['bbox_fs'] = copy.deepcopy(current_block['bbox'])
            if 'lines' in current_block and len(current_block['lines']) > 0:
                current_block['bbox_fs'] = [
                    min([line['bbox'][0] for line in current_block['lines']]),
                    min([line['bbox'][1] for line in current_block['lines']]),
                    max([line['bbox'][2] for line in current_block['lines']]),
                    max([line['bbox'][3] for line in current_block['lines']]),
                ]
54
55
            current_group.append(current_block)

56
57
58
59
60
61
62
        # 检查下一个块是否存在
        if i + 1 < len(blocks):
            next_block = blocks[i + 1]
            # 如果下一个块不是 text 类型且是 title 或 interline_equation 类型
            if next_block['type'] in ['title', 'interline_equation']:
                result.append(current_group)
                current_group = []
63
64
65
66
67
68
69
70

    # 处理最后一个 group
    if current_group:
        result.append(current_group)

    return result


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

    # index block 是一种特殊的list block
    # 一个block如果是index block 应该同时满足以下特征
    # 1.block内有多个line 2.block 内有多个line两侧均顶格写 3.line的开头或者结尾均为数字
80
    if len(block['lines']) >= 2:
81
82
83
        first_line = block['lines'][0]
        line_height = first_line['bbox'][3] - first_line['bbox'][1]
        block_weight = block['bbox_fs'][2] - block['bbox_fs'][0]
84
        block_height = block['bbox_fs'][3] - block['bbox_fs'][1]
85
        page_weight, page_height = block['page_size']
86
87
88
89

        left_close_num = 0
        left_not_close_num = 0
        right_not_close_num = 0
90
        right_close_num = 0
91
        lines_text_list = []
92
93
        center_close_num = 0
        external_sides_not_close_num = 0
94
95
        multiple_para_flag = False
        last_line = block['lines'][-1]
96

97
98
99
100
101
102
        if page_weight == 0:
            block_weight_radio = 0
        else:
            block_weight_radio = block_weight / page_weight
        # logger.info(f"block_weight_radio: {block_weight_radio}")

103
        # 如果首行左边不顶格而右边顶格,末行左边顶格而右边不顶格 (第一行可能可以右边不顶格)
104
105
106
107
        if (
            first_line['bbox'][0] - block['bbox_fs'][0] > line_height / 2
            and abs(last_line['bbox'][0] - block['bbox_fs'][0]) < line_height / 2
            and block['bbox_fs'][2] - last_line['bbox'][2] > line_height
108
109
110
        ):
            multiple_para_flag = True

111
        for line in block['lines']:
112
113
            line_mid_x = (line['bbox'][0] + line['bbox'][2]) / 2
            block_mid_x = (block['bbox_fs'][0] + block['bbox_fs'][2]) / 2
114
            if (
115
116
                line['bbox'][0] - block['bbox_fs'][0] > 0.8 * line_height
                and block['bbox_fs'][2] - line['bbox'][2] > 0.8 * line_height
117
            ):
118
                external_sides_not_close_num += 1
119
            if abs(line_mid_x - block_mid_x) < line_height / 2:
120
121
                center_close_num += 1

122
            line_text = ''
123
124
125
126
127
128

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

hyastar's avatar
hyastar committed
129
            # 添加所有文本,包括空行,保持与block['lines']长度一致
130
            lines_text_list.append(line_text)
131
132
133
            block_text = ''.join(lines_text_list)
            block_lang = detect_lang(block_text)
            # logger.info(f"block_lang: {block_lang}")
134
135

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

141
            # 计算右侧是否顶格
142
            if abs(block['bbox_fs'][2] - line['bbox'][2]) < line_height:
143
                right_close_num += 1
144
            else:
145
146
                # 类中文没有超长单词的情况,可以用统一的阈值
                if block_lang in ['zh', 'ja', 'ko']:
147
148
                    closed_area = 0.26 * block_weight
                else:
149
150
151
152
153
154
                    # 右侧不顶格情况下是否有一段距离,拍脑袋用0.3block宽度做阈值
                    # block宽的阈值可以小些,block窄的阈值要大
                    if block_weight_radio >= 0.5:
                        closed_area = 0.26 * block_weight
                    else:
                        closed_area = 0.36 * block_weight
155
156
                if block['bbox_fs'][2] - line['bbox'][2] > closed_area:
                    right_not_close_num += 1
157

158
159
        # 判断lines_text_list中的元素是否有超过80%都以LIST_END_FLAG结尾
        line_end_flag = False
160
161
        # 判断lines_text_list中的元素是否有超过80%都以数字开头或都以数字结尾
        line_num_flag = False
162
163
164
        num_start_count = 0
        num_end_count = 0
        flag_end_count = 0
hyastar's avatar
hyastar committed
165

166
167
168
        if len(lines_text_list) > 0:
            for line_text in lines_text_list:
                if len(line_text) > 0:
169
                    if line_text[-1] in LIST_END_FLAG:
170
                        flag_end_count += 1
171
172
173
174
175
                    if line_text[0].isdigit():
                        num_start_count += 1
                    if line_text[-1].isdigit():
                        num_end_count += 1

176
177
178
179
            if (
                num_start_count / len(lines_text_list) >= 0.8
                or num_end_count / len(lines_text_list) >= 0.8
            ):
180
                line_num_flag = True
hyastar's avatar
hyastar committed
181
182
            if flag_end_count / len(lines_text_list) >= 0.8:
                line_end_flag = True
183

184
        # 有的目录右侧不贴边, 目前认为左边或者右边有一边全贴边,且符合数字规则极为index
185
186
187
188
        if (
            left_close_num / len(block['lines']) >= 0.8
            or right_close_num / len(block['lines']) >= 0.8
        ) and line_num_flag:
189
190
            for line in block['lines']:
                line[ListLineTag.IS_LIST_START_LINE] = True
191
            return BlockType.Index
192

193
        # 全部line都居中的特殊list识别,每行都需要换行,特征是多行,且大多数行都前后not_close,每line中点x坐标接近
194
195
        # 补充条件block的长宽比有要求
        elif (
196
197
198
199
            external_sides_not_close_num >= 2
            and center_close_num == len(block['lines'])
            and external_sides_not_close_num / len(block['lines']) >= 0.5
            and block_height / block_weight > 0.4
200
        ):
201
202
203
            for line in block['lines']:
                line[ListLineTag.IS_LIST_START_LINE] = True
            return BlockType.List
204

205
        elif (
206
207
208
209
            left_close_num >= 2
            and (right_not_close_num >= 2 or line_end_flag or left_not_close_num >= 2)
            and not multiple_para_flag
            # and block_weight_radio > 0.27
210
        ):
211
            # 处理一种特殊的没有缩进的list,所有行都贴左边,通过右边的空隙判断是否是item尾
212
            if left_close_num / len(block['lines']) > 0.8:
213
214
215
216
217
218
219
220
                # 这种是每个item只有一行,且左边都贴边的短item list
                if flag_end_count == 0 and right_close_num / len(block['lines']) < 0.5:
                    for line in block['lines']:
                        if abs(block['bbox_fs'][0] - line['bbox'][0]) < line_height / 2:
                            line[ListLineTag.IS_LIST_START_LINE] = True
                # 这种是大部分line item 都有结束标识符的情况,按结束标识符区分不同item
                elif line_end_flag:
                    for i, line in enumerate(block['lines']):
221
222
223
224
                        if (
                            len(lines_text_list[i]) > 0
                            and lines_text_list[i][-1] in LIST_END_FLAG
                        ):
225
226
                            line[ListLineTag.IS_LIST_END_LINE] = True
                            if i + 1 < len(block['lines']):
227
228
229
                                block['lines'][i + 1][
                                    ListLineTag.IS_LIST_START_LINE
                                ] = True
230
231
232
233
234
235
236
                # line item基本没有结束标识符,而且也没有缩进,按右侧空隙判断哪些是item end
                else:
                    line_start_flag = False
                    for i, line in enumerate(block['lines']):
                        if line_start_flag:
                            line[ListLineTag.IS_LIST_START_LINE] = True
                            line_start_flag = False
hyastar's avatar
hyastar committed
237

238
239
240
241
                        if (
                            abs(block['bbox_fs'][2] - line['bbox'][2])
                            > 0.1 * block_weight
                        ):
242
243
                            line[ListLineTag.IS_LIST_END_LINE] = True
                            line_start_flag = True
hyastar's avatar
hyastar committed
244
245
            # 一种有缩进的特殊有序list,start line 左侧不贴边且以数字开头,end line 以 IS_LIST_END_FLAG 结尾且数量和start line 一致
            elif num_start_count >= 2 and num_start_count == flag_end_count:
246
                for i, line in enumerate(block['lines']):
hyastar's avatar
hyastar committed
247
248
249
250
251
                    if len(lines_text_list[i]) > 0:
                        if lines_text_list[i][0].isdigit():
                            line[ListLineTag.IS_LIST_START_LINE] = True
                        if lines_text_list[i][-1] in LIST_END_FLAG:
                            line[ListLineTag.IS_LIST_END_LINE] = True
252
253
254
255
256
257
258
            else:
                # 正常有缩进的list处理
                for line in block['lines']:
                    if abs(block['bbox_fs'][0] - line['bbox'][0]) < line_height / 2:
                        line[ListLineTag.IS_LIST_START_LINE] = True
                    if abs(block['bbox_fs'][2] - line['bbox'][2]) > line_height:
                        line[ListLineTag.IS_LIST_END_LINE] = True
259
260

            return BlockType.List
261
        else:
262
            return BlockType.Text
263
    else:
264
        return BlockType.Text
265
266
267


def __merge_2_text_blocks(block1, block2):
268
269
270
    if len(block1['lines']) > 0:
        first_line = block1['lines'][0]
        line_height = first_line['bbox'][3] - first_line['bbox'][1]
271
272
273
274
        block1_weight = block1['bbox'][2] - block1['bbox'][0]
        block2_weight = block2['bbox'][2] - block2['bbox'][0]
        min_block_weight = min(block1_weight, block2_weight)
        if abs(block1['bbox_fs'][0] - first_line['bbox'][0]) < line_height / 2:
275
276
277
278
            last_line = block2['lines'][-1]
            if len(last_line['spans']) > 0:
                last_span = last_line['spans'][-1]
                line_height = last_line['bbox'][3] - last_line['bbox'][1]
279
280
281
282
                if len(first_line['spans']) > 0:
                    first_span = first_line['spans'][0]
                    if len(first_span['content']) > 0:
                        span_start_with_num = first_span['content'][0].isdigit()
283
                        span_start_with_big_char = first_span['content'][0].isupper()
284
                        if (
285
286
287
                            # 上一个block的最后一个line的右边界和block的右边界差距不超过line_height
                            abs(block2['bbox_fs'][2] - last_line['bbox'][2]) < line_height
                            # 上一个block的最后一个span不是以特定符号结尾
288
289
290
                            and not last_span['content'].endswith(LINE_STOP_FLAG)
                            # 两个block宽度差距超过2倍也不合并
                            and abs(block1_weight - block2_weight) < min_block_weight
291
                            # 下一个block的第一个字符是数字
292
                            and not span_start_with_num
293
294
                            # 下一个block的第一个字符是大写字母
                            and not span_start_with_big_char
295
296
297
298
299
300
301
302
                        ):
                            if block1['page_num'] != block2['page_num']:
                                for line in block1['lines']:
                                    for span in line['spans']:
                                        span[CROSS_PAGE] = True
                            block2['lines'].extend(block1['lines'])
                            block1['lines'] = []
                            block1[LINES_DELETED] = True
303
304
305
306

    return block1, block2


307
308
309
310
311
312
313
314
315
316
317
318
def __merge_2_list_blocks(block1, block2):
    if block1['page_num'] != block2['page_num']:
        for line in block1['lines']:
            for span in line['spans']:
                span[CROSS_PAGE] = True
    block2['lines'].extend(block1['lines'])
    block1['lines'] = []
    block1[LINES_DELETED] = True

    return block1, block2


319
320
321
322
323
324
325
326
327
def __is_list_group(text_blocks_group):
    # list group的特征是一个group内的所有block都满足以下条件
    # 1.每个block都不超过3行 2. 每个block 的左边界都比较接近(逻辑简单点先不加这个规则)
    for block in text_blocks_group:
        if len(block['lines']) > 3:
            return False
    return True


328
329
330
def __para_merge_page(blocks):
    page_text_blocks_groups = __process_blocks(blocks)
    for text_blocks_group in page_text_blocks_groups:
331
        if len(text_blocks_group) > 0:
332
            # 需要先在合并前对所有block判断是否为list or index block
333
            for block in text_blocks_group:
334
335
336
                block_type = __is_list_or_index_block(block)
                block['type'] = block_type
                # logger.info(f"{block['type']}:{block}")
337

338
        if len(text_blocks_group) > 1:
339
340
341
            # 在合并前判断这个group 是否是一个 list group
            is_list_group = __is_list_group(text_blocks_group)

342
            # 倒序遍历
343
            for i in range(len(text_blocks_group) - 1, -1, -1):
344
                current_block = text_blocks_group[i]
345

346
347
348
                # 检查是否有前一个块
                if i - 1 >= 0:
                    prev_block = text_blocks_group[i - 1]
349

350
351
352
353
354
                    if (
                        current_block['type'] == 'text'
                        and prev_block['type'] == 'text'
                        and not is_list_group
                    ):
355
                        __merge_2_text_blocks(current_block, prev_block)
356
                    elif (
357
358
359
360
361
                        current_block['type'] == BlockType.List
                        and prev_block['type'] == BlockType.List
                    ) or (
                        current_block['type'] == BlockType.Index
                        and prev_block['type'] == BlockType.Index
362
                    ):
363
                        __merge_2_list_blocks(current_block, prev_block)
364

365
366
367
368
        else:
            continue


369
def para_split(pdf_info_dict):
370
371
372
373
374
    all_blocks = []
    for page_num, page in pdf_info_dict.items():
        blocks = copy.deepcopy(page['preproc_blocks'])
        for block in blocks:
            block['page_num'] = page_num
375
            block['page_size'] = page['page_size']
376
377
378
379
380
381
382
383
384
385
386
        all_blocks.extend(blocks)

    __para_merge_page(all_blocks)
    for page_num, page in pdf_info_dict.items():
        page['para_blocks'] = []
        for block in all_blocks:
            if block['page_num'] == page_num:
                page['para_blocks'].append(block)


if __name__ == '__main__':
387
    input_blocks = []
388
    # 调用函数
389
390
    groups = __process_blocks(input_blocks)
    for group_index, group in enumerate(groups):
391
        print(f'Group {group_index}: {group}')