magic-pdf 0.5.4__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (121) hide show
  1. magic_pdf/__init__.py +0 -0
  2. magic_pdf/cli/__init__.py +0 -0
  3. magic_pdf/cli/magicpdf.py +294 -0
  4. magic_pdf/dict2md/__init__.py +0 -0
  5. magic_pdf/dict2md/mkcontent.py +397 -0
  6. magic_pdf/dict2md/ocr_mkcontent.py +356 -0
  7. magic_pdf/filter/__init__.py +0 -0
  8. magic_pdf/filter/pdf_classify_by_type.py +381 -0
  9. magic_pdf/filter/pdf_meta_scan.py +368 -0
  10. magic_pdf/layout/__init__.py +0 -0
  11. magic_pdf/layout/bbox_sort.py +681 -0
  12. magic_pdf/layout/layout_det_utils.py +182 -0
  13. magic_pdf/layout/layout_sort.py +732 -0
  14. magic_pdf/layout/layout_spiler_recog.py +101 -0
  15. magic_pdf/layout/mcol_sort.py +336 -0
  16. magic_pdf/libs/Constants.py +11 -0
  17. magic_pdf/libs/MakeContentConfig.py +10 -0
  18. magic_pdf/libs/ModelBlockTypeEnum.py +9 -0
  19. magic_pdf/libs/__init__.py +0 -0
  20. magic_pdf/libs/boxbase.py +408 -0
  21. magic_pdf/libs/calc_span_stats.py +239 -0
  22. magic_pdf/libs/commons.py +204 -0
  23. magic_pdf/libs/config_reader.py +63 -0
  24. magic_pdf/libs/convert_utils.py +5 -0
  25. magic_pdf/libs/coordinate_transform.py +9 -0
  26. magic_pdf/libs/detect_language_from_model.py +21 -0
  27. magic_pdf/libs/draw_bbox.py +227 -0
  28. magic_pdf/libs/drop_reason.py +27 -0
  29. magic_pdf/libs/drop_tag.py +19 -0
  30. magic_pdf/libs/hash_utils.py +15 -0
  31. magic_pdf/libs/json_compressor.py +27 -0
  32. magic_pdf/libs/language.py +31 -0
  33. magic_pdf/libs/markdown_utils.py +31 -0
  34. magic_pdf/libs/math.py +9 -0
  35. magic_pdf/libs/nlp_utils.py +203 -0
  36. magic_pdf/libs/ocr_content_type.py +21 -0
  37. magic_pdf/libs/path_utils.py +23 -0
  38. magic_pdf/libs/pdf_image_tools.py +33 -0
  39. magic_pdf/libs/safe_filename.py +11 -0
  40. magic_pdf/libs/textbase.py +33 -0
  41. magic_pdf/libs/version.py +1 -0
  42. magic_pdf/libs/vis_utils.py +308 -0
  43. magic_pdf/model/__init__.py +0 -0
  44. magic_pdf/model/doc_analyze_by_360layout.py +8 -0
  45. magic_pdf/model/doc_analyze_by_pp_structurev2.py +125 -0
  46. magic_pdf/model/magic_model.py +632 -0
  47. magic_pdf/para/__init__.py +0 -0
  48. magic_pdf/para/block_continuation_processor.py +562 -0
  49. magic_pdf/para/block_termination_processor.py +480 -0
  50. magic_pdf/para/commons.py +222 -0
  51. magic_pdf/para/denoise.py +246 -0
  52. magic_pdf/para/draw.py +121 -0
  53. magic_pdf/para/exceptions.py +198 -0
  54. magic_pdf/para/layout_match_processor.py +40 -0
  55. magic_pdf/para/para_pipeline.py +297 -0
  56. magic_pdf/para/para_split.py +644 -0
  57. magic_pdf/para/para_split_v2.py +772 -0
  58. magic_pdf/para/raw_processor.py +207 -0
  59. magic_pdf/para/stats.py +268 -0
  60. magic_pdf/para/title_processor.py +1014 -0
  61. magic_pdf/pdf_parse_by_ocr.py +219 -0
  62. magic_pdf/pdf_parse_by_ocr_v2.py +17 -0
  63. magic_pdf/pdf_parse_by_txt.py +410 -0
  64. magic_pdf/pdf_parse_by_txt_v2.py +56 -0
  65. magic_pdf/pdf_parse_for_train.py +685 -0
  66. magic_pdf/pdf_parse_union_core.py +241 -0
  67. magic_pdf/pipe/AbsPipe.py +112 -0
  68. magic_pdf/pipe/OCRPipe.py +28 -0
  69. magic_pdf/pipe/TXTPipe.py +29 -0
  70. magic_pdf/pipe/UNIPipe.py +83 -0
  71. magic_pdf/pipe/__init__.py +0 -0
  72. magic_pdf/post_proc/__init__.py +0 -0
  73. magic_pdf/post_proc/detect_para.py +3472 -0
  74. magic_pdf/post_proc/pdf_post_filter.py +67 -0
  75. magic_pdf/post_proc/remove_footnote.py +153 -0
  76. magic_pdf/pre_proc/__init__.py +0 -0
  77. magic_pdf/pre_proc/citationmarker_remove.py +157 -0
  78. magic_pdf/pre_proc/construct_page_dict.py +72 -0
  79. magic_pdf/pre_proc/cut_image.py +71 -0
  80. magic_pdf/pre_proc/detect_equation.py +134 -0
  81. magic_pdf/pre_proc/detect_footer_by_model.py +64 -0
  82. magic_pdf/pre_proc/detect_footer_header_by_statistics.py +284 -0
  83. magic_pdf/pre_proc/detect_footnote.py +170 -0
  84. magic_pdf/pre_proc/detect_header.py +64 -0
  85. magic_pdf/pre_proc/detect_images.py +647 -0
  86. magic_pdf/pre_proc/detect_page_number.py +64 -0
  87. magic_pdf/pre_proc/detect_tables.py +62 -0
  88. magic_pdf/pre_proc/equations_replace.py +559 -0
  89. magic_pdf/pre_proc/fix_image.py +244 -0
  90. magic_pdf/pre_proc/fix_table.py +270 -0
  91. magic_pdf/pre_proc/main_text_font.py +23 -0
  92. magic_pdf/pre_proc/ocr_detect_all_bboxes.py +115 -0
  93. magic_pdf/pre_proc/ocr_detect_layout.py +133 -0
  94. magic_pdf/pre_proc/ocr_dict_merge.py +336 -0
  95. magic_pdf/pre_proc/ocr_span_list_modify.py +258 -0
  96. magic_pdf/pre_proc/pdf_pre_filter.py +74 -0
  97. magic_pdf/pre_proc/post_layout_split.py +0 -0
  98. magic_pdf/pre_proc/remove_bbox_overlap.py +98 -0
  99. magic_pdf/pre_proc/remove_colored_strip_bbox.py +79 -0
  100. magic_pdf/pre_proc/remove_footer_header.py +117 -0
  101. magic_pdf/pre_proc/remove_rotate_bbox.py +188 -0
  102. magic_pdf/pre_proc/resolve_bbox_conflict.py +191 -0
  103. magic_pdf/pre_proc/solve_line_alien.py +29 -0
  104. magic_pdf/pre_proc/statistics.py +12 -0
  105. magic_pdf/rw/AbsReaderWriter.py +34 -0
  106. magic_pdf/rw/DiskReaderWriter.py +66 -0
  107. magic_pdf/rw/S3ReaderWriter.py +107 -0
  108. magic_pdf/rw/__init__.py +0 -0
  109. magic_pdf/spark/__init__.py +0 -0
  110. magic_pdf/spark/spark_api.py +51 -0
  111. magic_pdf/train_utils/__init__.py +0 -0
  112. magic_pdf/train_utils/convert_to_train_format.py +65 -0
  113. magic_pdf/train_utils/extract_caption.py +59 -0
  114. magic_pdf/train_utils/remove_footer_header.py +159 -0
  115. magic_pdf/train_utils/vis_utils.py +327 -0
  116. magic_pdf/user_api.py +136 -0
  117. magic_pdf-0.5.4.dist-info/LICENSE.md +661 -0
  118. magic_pdf-0.5.4.dist-info/METADATA +24 -0
  119. magic_pdf-0.5.4.dist-info/RECORD +121 -0
  120. magic_pdf-0.5.4.dist-info/WHEEL +5 -0
  121. magic_pdf-0.5.4.dist-info/top_level.txt +1 -0
@@ -0,0 +1,685 @@
1
+ import time
2
+
3
+ # from anyio import Path
4
+
5
+ from magic_pdf.libs.commons import (
6
+ fitz,
7
+ get_delta_time,
8
+ get_img_s3_client,
9
+ get_docx_model_output,
10
+ )
11
+ import json
12
+ import os
13
+ from copy import deepcopy
14
+ import math
15
+ from loguru import logger
16
+ from magic_pdf.layout.bbox_sort import (
17
+ prepare_bboxes_for_layout_split,
18
+ )
19
+ from magic_pdf.layout.layout_sort import (
20
+ LAYOUT_UNPROC,
21
+ get_bboxes_layout,
22
+ get_columns_cnt_of_layout,
23
+ sort_text_block,
24
+ )
25
+ from magic_pdf.libs.drop_reason import DropReason
26
+ from magic_pdf.libs.markdown_utils import escape_special_markdown_char
27
+ from magic_pdf.libs.safe_filename import sanitize_filename
28
+ from magic_pdf.libs.vis_utils import draw_bbox_on_page, draw_layout_bbox_on_page
29
+ from magic_pdf.pre_proc.cut_image import txt_save_images_by_bboxes
30
+ from magic_pdf.pre_proc.detect_images import parse_images
31
+ from magic_pdf.pre_proc.detect_tables import parse_tables # 获取tables的bbox
32
+ from magic_pdf.pre_proc.detect_equation import parse_equations # 获取equations的bbox
33
+ from magic_pdf.pre_proc.detect_header import parse_headers # 获取headers的bbox
34
+ from magic_pdf.pre_proc.detect_page_number import parse_pageNos # 获取pageNos的bbox
35
+ from magic_pdf.pre_proc.detect_footnote import (
36
+ parse_footnotes_by_model,
37
+ parse_footnotes_by_rule,
38
+ ) # 获取footnotes的bbox
39
+ from magic_pdf.pre_proc.detect_footer_by_model import parse_footers # 获取footers的bbox
40
+
41
+ from magic_pdf.post_proc.detect_para import (
42
+ ParaProcessPipeline,
43
+ TitleDetectionException,
44
+ TitleLevelException,
45
+ ParaSplitException,
46
+ ParaMergeException,
47
+ DenseSingleLineBlockException,
48
+ )
49
+ from magic_pdf.pre_proc.main_text_font import get_main_text_font
50
+ from magic_pdf.pre_proc.remove_colored_strip_bbox import remove_colored_strip_textblock
51
+ from magic_pdf.pre_proc.remove_footer_header import remove_headder_footer_one_page
52
+ from magic_pdf.train_utils.extract_caption import extract_caption_bbox
53
+
54
+ """
55
+ from para.para_pipeline import ParaProcessPipeline
56
+ from para.exceptions import (
57
+ TitleDetectionException,
58
+ TitleLevelException,
59
+ ParaSplitException,
60
+ ParaMergeException,
61
+ DenseSingleLineBlockException,
62
+ )
63
+ """
64
+
65
+ from magic_pdf.libs.commons import read_file, join_path
66
+ from magic_pdf.post_proc.remove_footnote import (
67
+ merge_footnote_blocks,
68
+ remove_footnote_blocks,
69
+ )
70
+ from magic_pdf.pre_proc.citationmarker_remove import remove_citation_marker
71
+ from magic_pdf.pre_proc.equations_replace import (
72
+ combine_chars_to_pymudict,
73
+ remove_chars_in_text_blocks,
74
+ replace_equations_in_textblock,
75
+ )
76
+ from magic_pdf.pre_proc.pdf_pre_filter import pdf_filter
77
+ from magic_pdf.pre_proc.detect_footer_header_by_statistics import drop_footer_header
78
+ from magic_pdf.pre_proc.construct_page_dict import construct_page_component
79
+ from magic_pdf.pre_proc.fix_image import (
80
+ combine_images,
81
+ fix_image_vertical,
82
+ fix_seperated_image,
83
+ include_img_title,
84
+ )
85
+ from magic_pdf.post_proc.pdf_post_filter import pdf_post_filter
86
+ from magic_pdf.pre_proc.remove_rotate_bbox import (
87
+ get_side_boundry,
88
+ remove_rotate_side_textblock,
89
+ remove_side_blank_block,
90
+ )
91
+ from magic_pdf.pre_proc.resolve_bbox_conflict import (
92
+ check_text_block_horizontal_overlap,
93
+ resolve_bbox_overlap_conflict,
94
+ )
95
+ from magic_pdf.pre_proc.fix_table import (
96
+ fix_table_text_block,
97
+ fix_tables,
98
+ include_table_title,
99
+ )
100
+ from magic_pdf.pre_proc.solve_line_alien import solve_inline_too_large_interval
101
+
102
+ denseSingleLineBlockException_msg = DenseSingleLineBlockException().message
103
+ titleDetectionException_msg = TitleDetectionException().message
104
+ titleLevelException_msg = TitleLevelException().message
105
+ paraSplitException_msg = ParaSplitException().message
106
+ paraMergeException_msg = ParaMergeException().message
107
+
108
+
109
+ def parse_pdf_for_train(
110
+ s3_pdf_path,
111
+ s3_pdf_profile,
112
+ pdf_model_output,
113
+ save_path,
114
+ book_name,
115
+ image_s3_config=None,
116
+ start_page_id=0,
117
+ end_page_id=None,
118
+ junk_img_bojids=[],
119
+ debug_mode=False,
120
+ ):
121
+ pdf_bytes = read_file(s3_pdf_path, s3_pdf_profile)
122
+ save_tmp_path = os.path.join(os.path.dirname(__file__), "../..", "tmp", "unittest")
123
+ md_bookname_save_path = ""
124
+ book_name = sanitize_filename(book_name)
125
+ if debug_mode:
126
+ save_path = join_path(save_tmp_path, "md")
127
+ pdf_local_path = join_path(save_tmp_path, "download-pdfs", book_name)
128
+
129
+ if not os.path.exists(os.path.dirname(pdf_local_path)):
130
+ # 如果目录不存在,创建它
131
+ os.makedirs(os.path.dirname(pdf_local_path))
132
+
133
+ md_bookname_save_path = join_path(save_tmp_path, "md", book_name)
134
+ if not os.path.exists(md_bookname_save_path):
135
+ # 如果目录不存在,创建它
136
+ os.makedirs(md_bookname_save_path)
137
+
138
+ with open(pdf_local_path + ".pdf", "wb") as pdf_file:
139
+ pdf_file.write(pdf_bytes)
140
+
141
+ pdf_docs = fitz.open("pdf", pdf_bytes)
142
+ pdf_info_dict = {}
143
+ img_s3_client = get_img_s3_client(
144
+ save_path, image_s3_config
145
+ ) # 更改函数名和参数,避免歧义
146
+ # img_s3_client = "img_s3_client" #不创建这个对象,直接用字符串占位
147
+
148
+ start_time = time.time()
149
+
150
+ """通过统计pdf全篇文字,识别正文字体"""
151
+ main_text_font = get_main_text_font(pdf_docs)
152
+
153
+ end_page_id = end_page_id if end_page_id else len(pdf_docs) - 1
154
+ for page_id in range(start_page_id, end_page_id + 1):
155
+ page = pdf_docs[page_id]
156
+ page_width = page.rect.width
157
+ page_height = page.rect.height
158
+
159
+ if debug_mode:
160
+ time_now = time.time()
161
+ logger.info(
162
+ f"page_id: {page_id}, last_page_cost_time: {get_delta_time(start_time)}"
163
+ )
164
+ start_time = time_now
165
+ """
166
+ # 通过一个规则,过滤掉单页超过1500非junkimg的pdf
167
+ # 对单页面非重复id的img数量做统计,如果当前页超过1500则直接return need_drop
168
+ """
169
+ page_imgs = page.get_images()
170
+ img_counts = 0
171
+ for img in page_imgs:
172
+ img_bojid = img[0]
173
+ if img_bojid in junk_img_bojids: # 判断这个图片在不在junklist中
174
+ continue # 如果在junklist就不用管了,跳过
175
+ else:
176
+ recs = page.get_image_rects(img, transform=True)
177
+ if recs: # 如果这张图在当前页面有展示
178
+ img_counts += 1
179
+ if (
180
+ img_counts >= 1500
181
+ ): # 如果去除了junkimg的影响,单页img仍然超过1500的话,就排除当前pdf
182
+ logger.warning(
183
+ f"page_id: {page_id}, img_counts: {img_counts}, drop this pdf: {book_name}, drop_reason: {DropReason.HIGH_COMPUTATIONAL_lOAD_BY_IMGS}"
184
+ )
185
+ result = {
186
+ "_need_drop": True,
187
+ "_drop_reason": DropReason.HIGH_COMPUTATIONAL_lOAD_BY_IMGS,
188
+ }
189
+ if not debug_mode:
190
+ return result
191
+
192
+ """
193
+ ==================================================================================================================================
194
+ 首先获取基本的block数据,对pdf进行分解,获取图片、表格、公式、text的bbox
195
+ """
196
+ # 解析pdf原始文本block
197
+ text_raw_blocks = page.get_text(
198
+ "dict",
199
+ flags=fitz.TEXTFLAGS_TEXT,
200
+ )["blocks"]
201
+ model_output_json = get_docx_model_output(
202
+ pdf_model_output, page_id
203
+ )
204
+
205
+ # 解析图片
206
+ image_bboxes = parse_images(page_id, page, model_output_json, junk_img_bojids)
207
+ image_bboxes = fix_image_vertical(
208
+ image_bboxes, text_raw_blocks
209
+ ) # 修正图片的位置
210
+ image_bboxes = fix_seperated_image(image_bboxes) # 合并有边重合的图片
211
+
212
+ old_image_bboxes = deepcopy(image_bboxes)
213
+ image_bboxes = include_img_title(
214
+ text_raw_blocks, image_bboxes
215
+ ) # 向图片上方和下方寻找title,使用规则进行匹配,暂时只支持英文规则
216
+ """此时image_bboxes中可能出现这种情况,水平并列的2个图片,下方分别有各自的子标题,2个子标题下方又有大标题(形如Figxxx),会出现2个图片的bbox都包含了这个大标题,这种情况需要把图片合并"""
217
+ image_bboxes = combine_images(image_bboxes) # 合并图片
218
+
219
+ # 解析表格并对table_bboxes进行位置的微调,防止表格周围的文字被截断
220
+ table_bboxes = parse_tables(page_id, page, model_output_json)
221
+ table_bboxes = fix_tables(
222
+ page, table_bboxes, include_table_title=False, scan_line_num=2
223
+ ) # 修正
224
+ table_bboxes = fix_table_text_block(
225
+ text_raw_blocks, table_bboxes
226
+ ) # 修正与text block的关系,某些table修正与pymupdf获取到的table内textblock没有完全包含,因此要进行一次修正。
227
+ # debug_show_bbox(pdf_docs, page_id, table_bboxes, [], [b['bbox'] for b in text_raw_blocks], join_path(save_path, book_name, f"{book_name}_debug.pdf"), 7)
228
+
229
+ old_table_bboxes = deepcopy(table_bboxes)
230
+ table_bboxes = include_table_title(
231
+ text_raw_blocks, table_bboxes
232
+ ) # 向table上方和下方寻找title,使用规则进行匹配,暂时只支持英文规则
233
+
234
+ # 解析公式
235
+ equations_inline_bboxes, equations_interline_bboxes = parse_equations(
236
+ page_id, page, model_output_json
237
+ )
238
+
239
+ # get image box and caption !
240
+ image_bboxes_with_caption = extract_caption_bbox(image_bboxes, old_image_bboxes)
241
+
242
+ # get table box and caption !
243
+ table_bboxes_with_caption = extract_caption_bbox(table_bboxes, old_table_bboxes)
244
+
245
+ """
246
+ ==================================================================================================================================
247
+ 进入预处理-1阶段
248
+ -------------------
249
+ # # 解析标题
250
+ # title_bboxs = parse_titles(page_id, page, model_output_json)
251
+ # # 评估Layout是否规整、简单
252
+ # isSimpleLayout_flag, fullColumn_cnt, subColumn_cnt, curPage_loss = evaluate_pdf_layout(page_id, page, model_output_json)
253
+ 接下来开始进行预处理过程
254
+ """
255
+ # title_bboxs = parse_titles(page_id, page, model_output_json)
256
+
257
+ """去掉每页的页码、页眉、页脚"""
258
+ page_no_bboxs = parse_pageNos(page_id, page, model_output_json)
259
+ header_bboxs = parse_headers(page_id, page, model_output_json)
260
+ footer_bboxs = parse_footers(page_id, page, model_output_json)
261
+ (
262
+ image_bboxes,
263
+ table_bboxes,
264
+ remain_text_blocks,
265
+ removed_hdr_foot_txt_block,
266
+ removed_hdr_foot_img_block,
267
+ removed_hdr_foot_table,
268
+ ) = remove_headder_footer_one_page(
269
+ text_raw_blocks,
270
+ image_bboxes,
271
+ table_bboxes,
272
+ header_bboxs,
273
+ footer_bboxs,
274
+ page_no_bboxs,
275
+ page_width,
276
+ page_height,
277
+ )
278
+
279
+ """去除页面上半部分长条色块内的文本块"""
280
+ remain_text_blocks, removed_colored_narrow_strip_background_text_block = (
281
+ remove_colored_strip_textblock(remain_text_blocks, page)
282
+ )
283
+
284
+ # debug_show_bbox(pdf_docs, page_id, footnote_bboxes_by_model, [b['bbox'] for b in remain_text_blocks], header_bboxs, join_path(save_path, book_name, f"{book_name}_debug.pdf"), 7)
285
+
286
+ """去掉旋转的文字:水印、垂直排列的文字"""
287
+ remain_text_blocks, removed_non_horz_text_block = remove_rotate_side_textblock(
288
+ remain_text_blocks, page_width, page_height
289
+ ) # 去掉水印,非水平文字
290
+ remain_text_blocks, removed_empty_side_block = remove_side_blank_block(
291
+ remain_text_blocks, page_width, page_height
292
+ ) # 删除页面四周可能会留下的完全空白的textblock,这种block形成原因未知
293
+
294
+ """出现在图片、表格上的文字块去掉,把层叠的图片单独分离出来,不参与layout的计算"""
295
+ (
296
+ image_bboxes,
297
+ table_bboxes,
298
+ equations_interline_bboxes,
299
+ equations_inline_bboxes,
300
+ remain_text_blocks,
301
+ text_block_on_image_removed,
302
+ images_overlap_backup,
303
+ interline_eq_temp_text_block,
304
+ ) = resolve_bbox_overlap_conflict(
305
+ image_bboxes,
306
+ table_bboxes,
307
+ equations_interline_bboxes,
308
+ equations_inline_bboxes,
309
+ remain_text_blocks,
310
+ )
311
+
312
+ # """去掉footnote, 从文字和图片中"""
313
+ # # 通过模型识别到的footnote
314
+ # footnote_bboxes_by_model = parse_footnotes_by_model(page_id, page, model_output_json, md_bookname_save_path,
315
+ # debug_mode=debug_mode)
316
+ # # 通过规则识别到的footnote
317
+ # footnote_bboxes_by_rule = parse_footnotes_by_rule(remain_text_blocks, page_height, page_id)
318
+ """
319
+ ==================================================================================================================================
320
+ """
321
+ if debug_mode: # debugmode截图到本地
322
+ save_path = join_path(save_tmp_path, "md")
323
+
324
+ # 把图、表、公式都进行截图,保存到存储上,返回图片路径作为内容
325
+ image_info, image_backup_info, table_info, inline_eq_info, interline_eq_info = (
326
+ txt_save_images_by_bboxes(
327
+ book_name,
328
+ page_id,
329
+ page,
330
+ save_path,
331
+ image_bboxes,
332
+ images_overlap_backup,
333
+ table_bboxes,
334
+ equations_inline_bboxes,
335
+ equations_interline_bboxes,
336
+ # 传入img_s3_client
337
+ img_s3_client,
338
+ )
339
+ ) # 只要表格和图片的截图
340
+
341
+ """"以下进入到公式替换环节 """
342
+ char_level_text_blocks = page.get_text("rawdict", flags=fitz.TEXTFLAGS_TEXT)[
343
+ "blocks"
344
+ ]
345
+ remain_text_blocks = combine_chars_to_pymudict(
346
+ remain_text_blocks, char_level_text_blocks
347
+ ) # 合并chars
348
+ remain_text_blocks = replace_equations_in_textblock(
349
+ remain_text_blocks, inline_eq_info, interline_eq_info
350
+ )
351
+ remain_text_blocks = remove_citation_marker(
352
+ remain_text_blocks
353
+ ) # 公式替换之后去角标,防止公式无法替换成功。但是这样也会带来个问题就是把角标当公式。各有优劣。
354
+ remain_text_blocks = remove_chars_in_text_blocks(
355
+ remain_text_blocks
356
+ ) # 减少中间态数据体积
357
+ # debug_show_bbox(pdf_docs, page_id, [b['bbox'] for b in inline_eq_info], [b['bbox'] for b in interline_eq_info], [], join_path(save_path, book_name, f"{book_name}_debug.pdf"), 3)
358
+
359
+ """去掉footnote, 从文字和图片中(先去角标再去footnote试试)"""
360
+ # 通过模型识别到的footnote
361
+ footnote_bboxes_by_model = parse_footnotes_by_model(
362
+ page_id,
363
+ page,
364
+ model_output_json,
365
+ md_bookname_save_path,
366
+ debug_mode=debug_mode,
367
+ )
368
+ # 通过规则识别到的footnote
369
+ footnote_bboxes_by_rule = parse_footnotes_by_rule(
370
+ remain_text_blocks, page_height, page_id, main_text_font
371
+ )
372
+ """进入pdf过滤器,去掉一些不合理的pdf"""
373
+ is_good_pdf, err = pdf_filter(
374
+ page, remain_text_blocks, table_bboxes, image_bboxes
375
+ )
376
+ if not is_good_pdf:
377
+ logger.warning(
378
+ f"page_id: {page_id}, drop this pdf: {book_name}, reason: {err}"
379
+ )
380
+ if not debug_mode:
381
+ return err
382
+
383
+ """
384
+ ==================================================================================================================================
385
+ 进行版面布局切分和过滤
386
+ """
387
+ """在切分之前,先检查一下bbox是否有左右重叠的情况,如果有,那么就认为这个pdf暂时没有能力处理好,这种左右重叠的情况大概率是由于pdf里的行间公式、表格没有被正确识别出来造成的 """
388
+
389
+ is_text_block_horz_overlap = check_text_block_horizontal_overlap(
390
+ remain_text_blocks, header_bboxs, footer_bboxs
391
+ )
392
+
393
+ if is_text_block_horz_overlap:
394
+ # debug_show_bbox(pdf_docs, page_id, [b['bbox'] for b in remain_text_blocks], [], [], join_path(save_path, book_name, f"{book_name}_debug.pdf"), 0)
395
+ logger.warning(
396
+ f"page_id: {page_id}, drop this pdf: {book_name}, reason: {DropReason.TEXT_BLCOK_HOR_OVERLAP}"
397
+ )
398
+ result = {
399
+ "_need_drop": True,
400
+ "_drop_reason": DropReason.TEXT_BLCOK_HOR_OVERLAP,
401
+ }
402
+ if not debug_mode:
403
+ return result
404
+
405
+ """统一格式化成一个数据结构用于计算layout"""
406
+ page_y0 = 0 if len(header_bboxs) == 0 else max([b[3] for b in header_bboxs])
407
+ page_y1 = (
408
+ page_height if len(footer_bboxs) == 0 else min([b[1] for b in footer_bboxs])
409
+ )
410
+ left_x, right_x = get_side_boundry(
411
+ removed_non_horz_text_block, page_width, page_height
412
+ )
413
+ page_boundry = [
414
+ math.floor(left_x),
415
+ page_y0 + 1,
416
+ math.ceil(right_x),
417
+ page_y1 - 1,
418
+ ]
419
+ # 返回的是一个数组,每个元素[x0, y0, x1, y1, block_content, idx_x, idx_y], 初始时候idx_x, idx_y都是None. 对于图片、公式来说,block_content是图片的地址, 对于段落来说,block_content是段落的内容
420
+
421
+ all_bboxes = prepare_bboxes_for_layout_split(
422
+ image_info,
423
+ image_backup_info,
424
+ table_info,
425
+ inline_eq_info,
426
+ interline_eq_info,
427
+ remain_text_blocks,
428
+ page_boundry,
429
+ page,
430
+ )
431
+ # debug_show_bbox(pdf_docs, page_id, [], [], all_bboxes, join_path(save_path, book_name, f"{book_name}_debug.pdf"), 1)
432
+ """page_y0, page_y1能够过滤掉页眉和页脚,不会算作layout内"""
433
+ layout_bboxes, layout_tree = get_bboxes_layout(
434
+ all_bboxes, page_boundry, page_id
435
+ )
436
+
437
+ if (
438
+ len(remain_text_blocks) > 0
439
+ and len(all_bboxes) > 0
440
+ and len(layout_bboxes) == 0
441
+ ):
442
+ logger.warning(
443
+ f"page_id: {page_id}, drop this pdf: {book_name}, reason: {DropReason.CAN_NOT_DETECT_PAGE_LAYOUT}"
444
+ )
445
+ result = {
446
+ "_need_drop": True,
447
+ "_drop_reason": DropReason.CAN_NOT_DETECT_PAGE_LAYOUT,
448
+ }
449
+ if not debug_mode:
450
+ return result
451
+
452
+ """以下去掉复杂的布局和超过2列的布局"""
453
+ if any(
454
+ [lay["layout_label"] == LAYOUT_UNPROC for lay in layout_bboxes]
455
+ ): # 复杂的布局
456
+ logger.warning(
457
+ f"page_id: {page_id}, drop this pdf: {book_name}, reason: {DropReason.COMPLICATED_LAYOUT}"
458
+ )
459
+ result = {"_need_drop": True, "_drop_reason": DropReason.COMPLICATED_LAYOUT}
460
+ if not debug_mode:
461
+ return result
462
+
463
+ layout_column_width = get_columns_cnt_of_layout(layout_tree)
464
+ if layout_column_width > 2: # 去掉超过2列的布局pdf
465
+ logger.warning(
466
+ f"page_id: {page_id}, drop this pdf: {book_name}, reason: {DropReason.TOO_MANY_LAYOUT_COLUMNS}"
467
+ )
468
+ result = {
469
+ "_need_drop": True,
470
+ "_drop_reason": DropReason.TOO_MANY_LAYOUT_COLUMNS,
471
+ "extra_info": {"column_cnt": layout_column_width},
472
+ }
473
+ if not debug_mode:
474
+ return result
475
+
476
+ """
477
+ ==================================================================================================================================
478
+ 构造出下游需要的数据结构
479
+ """
480
+ remain_text_blocks = (
481
+ remain_text_blocks + interline_eq_temp_text_block
482
+ ) # 把计算layout时候临时删除的行间公式再放回去,防止行间公式替换的时候丢失。
483
+ removed_text_blocks = []
484
+ removed_text_blocks.extend(removed_hdr_foot_txt_block)
485
+ # removed_text_blocks.extend(removed_footnote_text_block)
486
+ removed_text_blocks.extend(text_block_on_image_removed)
487
+ removed_text_blocks.extend(removed_non_horz_text_block)
488
+ removed_text_blocks.extend(removed_colored_narrow_strip_background_text_block)
489
+
490
+ removed_images = []
491
+ # removed_images.extend(footnote_imgs)
492
+ removed_images.extend(removed_hdr_foot_img_block)
493
+
494
+ images_backup = []
495
+ images_backup.extend(image_backup_info)
496
+ remain_text_blocks = escape_special_markdown_char(
497
+ remain_text_blocks
498
+ ) # 转义span里的text
499
+ sorted_text_remain_text_block = sort_text_block(
500
+ remain_text_blocks, layout_bboxes
501
+ )
502
+
503
+ footnote_bboxes_tmp = []
504
+ footnote_bboxes_tmp.extend(footnote_bboxes_by_model)
505
+ footnote_bboxes_tmp.extend(footnote_bboxes_by_rule)
506
+
507
+ page_info = construct_page_component(
508
+ page_id,
509
+ image_info,
510
+ table_info,
511
+ sorted_text_remain_text_block,
512
+ layout_bboxes,
513
+ inline_eq_info,
514
+ interline_eq_info,
515
+ page.get_text("dict", flags=fitz.TEXTFLAGS_TEXT)["blocks"],
516
+ removed_text_blocks=removed_text_blocks,
517
+ removed_image_blocks=removed_images,
518
+ images_backup=images_backup,
519
+ droped_table_block=[],
520
+ table_backup=[],
521
+ layout_tree=layout_tree,
522
+ page_w=page.rect.width,
523
+ page_h=page.rect.height,
524
+ footnote_bboxes_tmp=footnote_bboxes_tmp,
525
+ )
526
+
527
+ page_info["image_bboxes_with_caption"] = image_bboxes_with_caption # add by xr
528
+ page_info["table_bboxes_with_caption"] = table_bboxes_with_caption
529
+
530
+ page_info["bak_page_no_bboxes"] = page_no_bboxs
531
+ page_info["bak_header_bboxes"] = header_bboxs
532
+ page_info["bak_footer_bboxes"] = footer_bboxs
533
+ page_info["bak_footer_note_bboxes"] = footnote_bboxes_tmp
534
+
535
+ pdf_info_dict[f"page_{page_id}"] = page_info
536
+
537
+ # end page for
538
+
539
+ """计算后处理阶段耗时"""
540
+ start_time = time.time()
541
+
542
+ """
543
+ ==================================================================================================================================
544
+ 去掉页眉和页脚,这里需要用到一定的统计量,所以放到最后
545
+ 页眉和页脚主要从文本box和图片box中去除,位于页面的四周。
546
+ 下面函数会直接修改pdf_info_dict,从文字块中、图片中删除属于页眉页脚的内容,删除内容做相对应记录
547
+ """
548
+ # 去页眉页脚
549
+ header, footer = drop_footer_header(
550
+ pdf_info_dict
551
+ ) # TODO: using header and footer boxes here !
552
+
553
+ """对单个layout内footnote和他下面的所有textbbox合并"""
554
+
555
+ for page_key, page_info in pdf_info_dict.items():
556
+ page_info = merge_footnote_blocks(page_info, main_text_font)
557
+ page_info = remove_footnote_blocks(page_info)
558
+ pdf_info_dict[page_key] = page_info
559
+
560
+ """进入pdf后置过滤器,去掉一些不合理的pdf"""
561
+
562
+ i = 0
563
+ for page_info in pdf_info_dict.values():
564
+ is_good_pdf, err = pdf_post_filter(page_info)
565
+ if not is_good_pdf:
566
+ logger.warning(f"page_id: {i}, drop this pdf: {book_name}, reason: {err}")
567
+ if not debug_mode:
568
+ return err
569
+ i += 1
570
+
571
+ if debug_mode:
572
+ params_file_save_path = join_path(
573
+ save_tmp_path, "md", book_name, "preproc_out.json"
574
+ )
575
+ page_draw_rect_save_path = join_path(
576
+ save_tmp_path, "md", book_name, "layout.pdf"
577
+ )
578
+ # dir_path = os.path.dirname(page_draw_rect_save_path)
579
+ # if not os.path.exists(dir_path):
580
+ # # 如果目录不存在,创建它
581
+ # os.makedirs(dir_path)
582
+
583
+ with open(params_file_save_path, "w", encoding="utf-8") as f:
584
+ json.dump(pdf_info_dict, f, ensure_ascii=False, indent=4)
585
+ # 先检测本地 page_draw_rect_save_path 是否存在,如果存在则删除
586
+ if os.path.exists(page_draw_rect_save_path):
587
+ os.remove(page_draw_rect_save_path)
588
+ # 绘制bbox和layout到pdf
589
+ draw_bbox_on_page(pdf_docs, pdf_info_dict, page_draw_rect_save_path)
590
+ draw_layout_bbox_on_page(
591
+ pdf_docs, pdf_info_dict, header, footer, page_draw_rect_save_path
592
+ )
593
+
594
+ if debug_mode:
595
+ # 打印后处理阶段耗时
596
+ logger.info(f"post_processing_time: {get_delta_time(start_time)}")
597
+
598
+ """
599
+ ==================================================================================================================================
600
+ 进入段落处理-2阶段
601
+ """
602
+
603
+ # 处理行内文字间距较大问题
604
+ pdf_info_dict = solve_inline_too_large_interval(pdf_info_dict)
605
+
606
+ start_time = time.time()
607
+
608
+ para_process_pipeline = ParaProcessPipeline()
609
+
610
+ def _deal_with_text_exception(error_info):
611
+ logger.warning(
612
+ f"page_id: {page_id}, drop this pdf: {book_name}, reason: {error_info}"
613
+ )
614
+ if error_info == denseSingleLineBlockException_msg:
615
+ logger.warning(
616
+ f"Drop this pdf: {book_name}, reason: {DropReason.DENSE_SINGLE_LINE_BLOCK}"
617
+ )
618
+ result = {
619
+ "_need_drop": True,
620
+ "_drop_reason": DropReason.DENSE_SINGLE_LINE_BLOCK,
621
+ }
622
+ return result
623
+ if error_info == titleDetectionException_msg:
624
+ logger.warning(
625
+ f"Drop this pdf: {book_name}, reason: {DropReason.TITLE_DETECTION_FAILED}"
626
+ )
627
+ result = {
628
+ "_need_drop": True,
629
+ "_drop_reason": DropReason.TITLE_DETECTION_FAILED,
630
+ }
631
+ return result
632
+ elif error_info == titleLevelException_msg:
633
+ logger.warning(
634
+ f"Drop this pdf: {book_name}, reason: {DropReason.TITLE_LEVEL_FAILED}"
635
+ )
636
+ result = {"_need_drop": True, "_drop_reason": DropReason.TITLE_LEVEL_FAILED}
637
+ return result
638
+ elif error_info == paraSplitException_msg:
639
+ logger.warning(
640
+ f"Drop this pdf: {book_name}, reason: {DropReason.PARA_SPLIT_FAILED}"
641
+ )
642
+ result = {"_need_drop": True, "_drop_reason": DropReason.PARA_SPLIT_FAILED}
643
+ return result
644
+ elif error_info == paraMergeException_msg:
645
+ logger.warning(
646
+ f"Drop this pdf: {book_name}, reason: {DropReason.PARA_MERGE_FAILED}"
647
+ )
648
+ result = {"_need_drop": True, "_drop_reason": DropReason.PARA_MERGE_FAILED}
649
+ return result
650
+
651
+ if debug_mode:
652
+ input_pdf_file = f"{pdf_local_path}.pdf"
653
+ output_dir = f"{save_path}/{book_name}"
654
+ output_pdf_file = f"{output_dir}/pdf_annos.pdf"
655
+
656
+ """
657
+ Call the para_process_pipeline function to process the pdf_info_dict.
658
+
659
+ Parameters:
660
+ para_debug_mode: str or None
661
+ If para_debug_mode is None, the para_process_pipeline will not keep any intermediate results.
662
+ If para_debug_mode is "simple", the para_process_pipeline will only keep the annos on the pdf and the final results as a json file.
663
+ If para_debug_mode is "full", the para_process_pipeline will keep all the intermediate results generated during each step.
664
+ """
665
+ pdf_info_dict, error_info = para_process_pipeline.para_process_pipeline(
666
+ pdf_info_dict,
667
+ para_debug_mode="simple",
668
+ input_pdf_path=input_pdf_file,
669
+ output_pdf_path=output_pdf_file,
670
+ )
671
+ # 打印段落处理阶段耗时
672
+ logger.info(f"para_process_time: {get_delta_time(start_time)}")
673
+
674
+ # debug的时候不return drop信息
675
+ if error_info is not None:
676
+ _deal_with_text_exception(error_info)
677
+ return pdf_info_dict
678
+ else:
679
+ pdf_info_dict, error_info = para_process_pipeline.para_process_pipeline(
680
+ pdf_info_dict
681
+ )
682
+ if error_info is not None:
683
+ return _deal_with_text_exception(error_info)
684
+
685
+ return pdf_info_dict