pdf_parse_by_model.py 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508
  1. import time
  2. # from anyio import Path
  3. from magic_pdf.libs.commons import fitz, get_delta_time, get_img_s3_client
  4. import json
  5. import os
  6. import math
  7. from loguru import logger
  8. from magic_pdf.layout.bbox_sort import (
  9. prepare_bboxes_for_layout_split,
  10. )
  11. from magic_pdf.layout.layout_sort import LAYOUT_UNPROC, get_bboxes_layout, get_columns_cnt_of_layout, sort_text_block
  12. from magic_pdf.libs.drop_reason import DropReason
  13. from magic_pdf.libs.markdown_utils import escape_special_markdown_char
  14. from magic_pdf.libs.safe_filename import sanitize_filename
  15. from magic_pdf.libs.vis_utils import draw_bbox_on_page, draw_layout_bbox_on_page
  16. from magic_pdf.pre_proc.detect_images import parse_images
  17. from magic_pdf.pre_proc.detect_tables import parse_tables # 获取tables的bbox
  18. from magic_pdf.pre_proc.detect_equation import parse_equations # 获取equations的bbox
  19. from magic_pdf.pre_proc.detect_header import parse_headers # 获取headers的bbox
  20. from magic_pdf.pre_proc.detect_page_number import parse_pageNos # 获取pageNos的bbox
  21. from magic_pdf.pre_proc.detect_footnote import parse_footnotes_by_model, parse_footnotes_by_rule # 获取footnotes的bbox
  22. from magic_pdf.pre_proc.detect_footer_by_model import parse_footers # 获取footers的bbox
  23. from magic_pdf.post_proc.detect_para import (
  24. ParaProcessPipeline,
  25. TitleDetectionException,
  26. TitleLevelException,
  27. ParaSplitException,
  28. ParaMergeException,
  29. DenseSingleLineBlockException,
  30. )
  31. from magic_pdf.pre_proc.main_text_font import get_main_text_font
  32. from magic_pdf.pre_proc.remove_colored_strip_bbox import remove_colored_strip_textblock
  33. from magic_pdf.pre_proc.remove_footer_header import remove_headder_footer_one_page
  34. '''
  35. from para.para_pipeline import ParaProcessPipeline
  36. from para.exceptions import (
  37. TitleDetectionException,
  38. TitleLevelException,
  39. ParaSplitException,
  40. ParaMergeException,
  41. DenseSingleLineBlockException,
  42. )
  43. '''
  44. from magic_pdf.libs.commons import read_file, join_path
  45. from magic_pdf.libs.pdf_image_tools import save_images_by_bboxes
  46. from magic_pdf.post_proc.remove_footnote import merge_footnote_blocks, remove_footnote_blocks
  47. from magic_pdf.pre_proc.citationmarker_remove import remove_citation_marker
  48. from magic_pdf.pre_proc.equations_replace import combine_chars_to_pymudict, remove_chars_in_text_blocks, replace_equations_in_textblock
  49. from magic_pdf.pre_proc.pdf_pre_filter import pdf_filter
  50. from magic_pdf.pre_proc.detect_footer_header_by_statistics import drop_footer_header
  51. from magic_pdf.pre_proc.construct_paras import construct_page_component
  52. from magic_pdf.pre_proc.fix_image import combine_images, fix_image_vertical, fix_seperated_image, include_img_title
  53. from magic_pdf.post_proc.pdf_post_filter import pdf_post_filter
  54. from magic_pdf.pre_proc.remove_rotate_bbox import get_side_boundry, remove_rotate_side_textblock, remove_side_blank_block
  55. from magic_pdf.pre_proc.resolve_bbox_conflict import check_text_block_horizontal_overlap, resolve_bbox_overlap_conflict
  56. from magic_pdf.pre_proc.fix_table import fix_table_text_block, fix_tables, include_table_title
  57. denseSingleLineBlockException_msg = DenseSingleLineBlockException().message
  58. titleDetectionException_msg = TitleDetectionException().message
  59. titleLevelException_msg = TitleLevelException().message
  60. paraSplitException_msg = ParaSplitException().message
  61. paraMergeException_msg = ParaMergeException().message
  62. def get_docx_model_output(pdf_model_output, pdf_model_s3_profile, page_id):
  63. if isinstance(pdf_model_output, str):
  64. model_output_json_path = join_path(pdf_model_output, f"page_{page_id + 1}.json") # 模型输出的页面编号从1开始的
  65. if os.path.exists(model_output_json_path):
  66. json_from_docx = read_file(model_output_json_path, pdf_model_s3_profile)
  67. model_output_json = json.loads(json_from_docx)
  68. else:
  69. try:
  70. model_output_json_path = join_path(pdf_model_output, "model.json")
  71. with open(model_output_json_path, "r", encoding="utf-8") as f:
  72. model_output_json = json.load(f)
  73. model_output_json = model_output_json["doc_layout_result"][page_id]
  74. except:
  75. s3_model_output_json_path = join_path(pdf_model_output, f"page_{page_id + 1}.json")
  76. s3_model_output_json_path = join_path(pdf_model_output, f"{page_id}.json")
  77. #s3_model_output_json_path = join_path(pdf_model_output, f"page_{page_id }.json")
  78. # logger.warning(f"model_output_json_path: {model_output_json_path} not found. try to load from s3: {s3_model_output_json_path}")
  79. s = read_file(s3_model_output_json_path, pdf_model_s3_profile)
  80. return json.loads(s)
  81. elif isinstance(pdf_model_output, list):
  82. model_output_json = pdf_model_output[page_id]
  83. return model_output_json
  84. def parse_pdf_by_model(
  85. s3_pdf_path,
  86. s3_pdf_profile,
  87. pdf_model_output,
  88. save_path,
  89. book_name,
  90. pdf_model_profile=None,
  91. image_s3_config=None,
  92. start_page_id=0,
  93. end_page_id=None,
  94. junk_img_bojids=[],
  95. debug_mode=False,
  96. ):
  97. pdf_bytes = read_file(s3_pdf_path, s3_pdf_profile)
  98. save_tmp_path = os.path.join(os.path.dirname(__file__), "../..", "tmp", "unittest")
  99. md_bookname_save_path = ""
  100. book_name = sanitize_filename(book_name)
  101. if debug_mode:
  102. save_path = join_path(save_tmp_path, "md")
  103. pdf_local_path = join_path(save_tmp_path, "download-pdfs", book_name)
  104. if not os.path.exists(os.path.dirname(pdf_local_path)):
  105. # 如果目录不存在,创建它
  106. os.makedirs(os.path.dirname(pdf_local_path))
  107. md_bookname_save_path = join_path(save_tmp_path, "md", book_name)
  108. if not os.path.exists(md_bookname_save_path):
  109. # 如果目录不存在,创建它
  110. os.makedirs(md_bookname_save_path)
  111. with open(pdf_local_path + ".pdf", "wb") as pdf_file:
  112. pdf_file.write(pdf_bytes)
  113. pdf_docs = fitz.open("pdf", pdf_bytes)
  114. pdf_info_dict = {}
  115. img_s3_client = get_img_s3_client(save_path, image_s3_config) # 更改函数名和参数,避免歧义
  116. # img_s3_client = "img_s3_client" #不创建这个对象,直接用字符串占位
  117. start_time = time.time()
  118. """通过统计pdf全篇文字,识别正文字体"""
  119. main_text_font = get_main_text_font(pdf_docs)
  120. end_page_id = end_page_id if end_page_id else len(pdf_docs) - 1
  121. for page_id in range(start_page_id, end_page_id + 1):
  122. page = pdf_docs[page_id]
  123. page_width = page.rect.width
  124. page_height = page.rect.height
  125. if debug_mode:
  126. time_now = time.time()
  127. logger.info(f"page_id: {page_id}, last_page_cost_time: {get_delta_time(start_time)}")
  128. start_time = time_now
  129. """
  130. # 通过一个规则,过滤掉单页超过1500非junkimg的pdf
  131. # 对单页面非重复id的img数量做统计,如果当前页超过1500则直接return need_drop
  132. """
  133. page_imgs = page.get_images()
  134. img_counts = 0
  135. for img in page_imgs:
  136. img_bojid = img[0]
  137. if img_bojid in junk_img_bojids: # 判断这个图片在不在junklist中
  138. continue # 如果在junklist就不用管了,跳过
  139. else:
  140. recs = page.get_image_rects(img, transform=True)
  141. if recs: # 如果这张图在当前页面有展示
  142. img_counts += 1
  143. if img_counts >= 1500: # 如果去除了junkimg的影响,单页img仍然超过1500的话,就排除当前pdf
  144. logger.warning(
  145. f"page_id: {page_id}, img_counts: {img_counts}, drop this pdf: {book_name}, drop_reason: {DropReason.HIGH_COMPUTATIONAL_lOAD_BY_IMGS}"
  146. )
  147. result = {"need_drop": True, "drop_reason": DropReason.HIGH_COMPUTATIONAL_lOAD_BY_IMGS}
  148. if not debug_mode:
  149. return result
  150. """
  151. ==================================================================================================================================
  152. 首先获取基本的block数据,对pdf进行分解,获取图片、表格、公式、text的bbox
  153. """
  154. # 解析pdf原始文本block
  155. text_raw_blocks = page.get_text(
  156. "dict",
  157. flags=fitz.TEXTFLAGS_TEXT,
  158. )["blocks"]
  159. model_output_json = get_docx_model_output(pdf_model_output, pdf_model_profile, page_id)
  160. # 解析图片
  161. image_bboxes = parse_images(page_id, page, model_output_json, junk_img_bojids)
  162. image_bboxes = fix_image_vertical(image_bboxes, text_raw_blocks) # 修正图片的位置
  163. image_bboxes = fix_seperated_image(image_bboxes) # 合并有边重合的图片
  164. image_bboxes = include_img_title(text_raw_blocks, image_bboxes) # 向图片上方和下方寻找title,使用规则进行匹配,暂时只支持英文规则
  165. """此时image_bboxes中可能出现这种情况,水平并列的2个图片,下方分别有各自的子标题,2个子标题下方又有大标题(形如Figxxx),会出现2个图片的bbox都包含了这个大标题,这种情况需要把图片合并"""
  166. image_bboxes = combine_images(image_bboxes) # 合并图片
  167. # 解析表格并对table_bboxes进行位置的微调,防止表格周围的文字被截断
  168. table_bboxes = parse_tables(page_id, page, model_output_json)
  169. table_bboxes = fix_tables(page, table_bboxes, include_table_title=True, scan_line_num=2) # 修正
  170. table_bboxes = fix_table_text_block(text_raw_blocks, table_bboxes) # 修正与text block的关系,某些table修正与pymupdf获取到的table内textblock没有完全包含,因此要进行一次修正。
  171. #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)
  172. table_bboxes = include_table_title(text_raw_blocks, table_bboxes) # 向table上方和下方寻找title,使用规则进行匹配,暂时只支持英文规则
  173. # 解析公式
  174. equations_inline_bboxes, equations_interline_bboxes = parse_equations(page_id, page, model_output_json)
  175. """
  176. ==================================================================================================================================
  177. 进入预处理-1阶段
  178. -------------------
  179. # # 解析标题
  180. # title_bboxs = parse_titles(page_id, page, model_output_json)
  181. # # 评估Layout是否规整、简单
  182. # isSimpleLayout_flag, fullColumn_cnt, subColumn_cnt, curPage_loss = evaluate_pdf_layout(page_id, page, model_output_json)
  183. 接下来开始进行预处理过程
  184. """
  185. """去掉每页的页码、页眉、页脚"""
  186. page_no_bboxs = parse_pageNos(page_id, page, model_output_json)
  187. header_bboxs = parse_headers(page_id, page, model_output_json)
  188. footer_bboxs = parse_footers(page_id, page, model_output_json)
  189. image_bboxes, table_bboxes, remain_text_blocks, removed_hdr_foot_txt_block, removed_hdr_foot_img_block, removed_hdr_foot_table = remove_headder_footer_one_page(text_raw_blocks, image_bboxes, table_bboxes, header_bboxs, footer_bboxs, page_no_bboxs, page_width, page_height)
  190. """去除页面上半部分长条色块内的文本块"""
  191. remain_text_blocks, removed_colored_narrow_strip_background_text_block = remove_colored_strip_textblock(remain_text_blocks, page)
  192. #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)
  193. """去掉旋转的文字:水印、垂直排列的文字"""
  194. remain_text_blocks, removed_non_horz_text_block = remove_rotate_side_textblock(
  195. remain_text_blocks, page_width, page_height
  196. ) # 去掉水印,非水平文字
  197. remain_text_blocks, removed_empty_side_block = remove_side_blank_block(remain_text_blocks, page_width, page_height) # 删除页面四周可能会留下的完全空白的textblock,这种block形成原因未知
  198. """出现在图片、表格上的文字块去掉,把层叠的图片单独分离出来,不参与layout的计算"""
  199. (
  200. image_bboxes,
  201. table_bboxes,
  202. equations_interline_bboxes,
  203. equations_inline_bboxes,
  204. remain_text_blocks,
  205. text_block_on_image_removed,
  206. images_overlap_backup,
  207. interline_eq_temp_text_block
  208. ) = resolve_bbox_overlap_conflict(
  209. image_bboxes, table_bboxes, equations_interline_bboxes, equations_inline_bboxes, remain_text_blocks
  210. )
  211. # """去掉footnote, 从文字和图片中"""
  212. # # 通过模型识别到的footnote
  213. # footnote_bboxes_by_model = parse_footnotes_by_model(page_id, page, model_output_json, md_bookname_save_path,
  214. # debug_mode=debug_mode)
  215. # # 通过规则识别到的footnote
  216. # footnote_bboxes_by_rule = parse_footnotes_by_rule(remain_text_blocks, page_height, page_id)
  217. """
  218. ==================================================================================================================================
  219. """
  220. if debug_mode: # debugmode截图到本地
  221. save_path = join_path(save_tmp_path, "md")
  222. # 把图、表、公式都进行截图,保存到存储上,返回图片路径作为内容
  223. image_info, image_backup_info, table_info, inline_eq_info, interline_eq_info = save_images_by_bboxes(
  224. book_name,
  225. page_id,
  226. page,
  227. save_path,
  228. image_bboxes,
  229. images_overlap_backup,
  230. table_bboxes,
  231. equations_inline_bboxes,
  232. equations_interline_bboxes,
  233. # 传入img_s3_client
  234. img_s3_client,
  235. ) # 只要表格和图片的截图
  236. """"以下进入到公式替换环节 """
  237. char_level_text_blocks = page.get_text("rawdict", flags=fitz.TEXTFLAGS_TEXT)['blocks']
  238. remain_text_blocks = combine_chars_to_pymudict(remain_text_blocks, char_level_text_blocks)# 合并chars
  239. remain_text_blocks = replace_equations_in_textblock(remain_text_blocks, inline_eq_info, interline_eq_info)
  240. remain_text_blocks = remove_citation_marker(remain_text_blocks) # 公式替换之后去角标,防止公式无法替换成功。但是这样也会带来个问题就是把角标当公式。各有优劣。
  241. remain_text_blocks = remove_chars_in_text_blocks(remain_text_blocks) # 减少中间态数据体积
  242. #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)
  243. """去掉footnote, 从文字和图片中(先去角标再去footnote试试)"""
  244. # 通过模型识别到的footnote
  245. footnote_bboxes_by_model = parse_footnotes_by_model(page_id, page, model_output_json, md_bookname_save_path, debug_mode=debug_mode)
  246. # 通过规则识别到的footnote
  247. footnote_bboxes_by_rule = parse_footnotes_by_rule(remain_text_blocks, page_height, page_id, main_text_font)
  248. """进入pdf过滤器,去掉一些不合理的pdf"""
  249. is_good_pdf, err = pdf_filter(page, remain_text_blocks, table_bboxes, image_bboxes)
  250. if not is_good_pdf:
  251. logger.warning(f"page_id: {page_id}, drop this pdf: {book_name}, reason: {err}")
  252. if not debug_mode:
  253. return err
  254. """
  255. ==================================================================================================================================
  256. 进行版面布局切分和过滤
  257. """
  258. """在切分之前,先检查一下bbox是否有左右重叠的情况,如果有,那么就认为这个pdf暂时没有能力处理好,这种左右重叠的情况大概率是由于pdf里的行间公式、表格没有被正确识别出来造成的 """
  259. is_text_block_horz_overlap = check_text_block_horizontal_overlap(remain_text_blocks, header_bboxs, footer_bboxs)
  260. if is_text_block_horz_overlap:
  261. # 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)
  262. logger.warning(f"page_id: {page_id}, drop this pdf: {book_name}, reason: {DropReason.TEXT_BLCOK_HOR_OVERLAP}")
  263. result = {"need_drop": True, "drop_reason": DropReason.TEXT_BLCOK_HOR_OVERLAP}
  264. if not debug_mode:
  265. return result
  266. """统一格式化成一个数据结构用于计算layout"""
  267. page_y0 = 0 if len(header_bboxs) == 0 else max([b[3] for b in header_bboxs])
  268. page_y1 = page_height if len(footer_bboxs) == 0 else min([b[1] for b in footer_bboxs])
  269. left_x, right_x = get_side_boundry(removed_non_horz_text_block, page_width, page_height)
  270. page_boundry = [math.floor(left_x), page_y0 + 1, math.ceil(right_x), page_y1 - 1]
  271. # 返回的是一个数组,每个元素[x0, y0, x1, y1, block_content, idx_x, idx_y], 初始时候idx_x, idx_y都是None. 对于图片、公式来说,block_content是图片的地址, 对于段落来说,block_content是段落的内容
  272. all_bboxes = prepare_bboxes_for_layout_split(
  273. image_info, image_backup_info, table_info, inline_eq_info, interline_eq_info, remain_text_blocks, page_boundry, page)
  274. #debug_show_bbox(pdf_docs, page_id, [], [], all_bboxes, join_path(save_path, book_name, f"{book_name}_debug.pdf"), 1)
  275. """page_y0, page_y1能够过滤掉页眉和页脚,不会算作layout内"""
  276. layout_bboxes, layout_tree = get_bboxes_layout(all_bboxes, page_boundry, page_id)
  277. if len(remain_text_blocks)>0 and len(all_bboxes)>0 and len(layout_bboxes)==0:
  278. logger.warning(f"page_id: {page_id}, drop this pdf: {book_name}, reason: {DropReason.CAN_NOT_DETECT_PAGE_LAYOUT}")
  279. result = {"need_drop": True, "drop_reason": DropReason.CAN_NOT_DETECT_PAGE_LAYOUT}
  280. if not debug_mode:
  281. return result
  282. """以下去掉复杂的布局和超过2列的布局"""
  283. if any([lay["layout_label"] == LAYOUT_UNPROC for lay in layout_bboxes]): # 复杂的布局
  284. logger.warning(f"page_id: {page_id}, drop this pdf: {book_name}, reason: {DropReason.COMPLICATED_LAYOUT}")
  285. result = {"need_drop": True, "drop_reason": DropReason.COMPLICATED_LAYOUT}
  286. if not debug_mode:
  287. return result
  288. layout_column_width = get_columns_cnt_of_layout(layout_tree)
  289. if layout_column_width > 2: # 去掉超过2列的布局pdf
  290. logger.warning(f"page_id: {page_id}, drop this pdf: {book_name}, reason: {DropReason.TOO_MANY_LAYOUT_COLUMNS}")
  291. result = {
  292. "need_drop": True,
  293. "drop_reason": DropReason.TOO_MANY_LAYOUT_COLUMNS,
  294. "extra_info": {"column_cnt": layout_column_width},
  295. }
  296. if not debug_mode:
  297. return result
  298. """
  299. ==================================================================================================================================
  300. 构造出下游需要的数据结构
  301. """
  302. remain_text_blocks = remain_text_blocks + interline_eq_temp_text_block # 把计算layout时候临时删除的行间公式再放回去,防止行间公式替换的时候丢失。
  303. removed_text_blocks = []
  304. removed_text_blocks.extend(removed_hdr_foot_txt_block)
  305. # removed_text_blocks.extend(removed_footnote_text_block)
  306. removed_text_blocks.extend(text_block_on_image_removed)
  307. removed_text_blocks.extend(removed_non_horz_text_block)
  308. removed_text_blocks.extend(removed_colored_narrow_strip_background_text_block)
  309. removed_images = []
  310. # removed_images.extend(footnote_imgs)
  311. removed_images.extend(removed_hdr_foot_img_block)
  312. images_backup = []
  313. images_backup.extend(image_backup_info)
  314. remain_text_blocks = escape_special_markdown_char(remain_text_blocks) # 转义span里的text
  315. sorted_text_remain_text_block = sort_text_block(remain_text_blocks, layout_bboxes)
  316. footnote_bboxes_tmp = []
  317. footnote_bboxes_tmp.extend(footnote_bboxes_by_model)
  318. footnote_bboxes_tmp.extend(footnote_bboxes_by_rule)
  319. page_info = construct_page_component(
  320. page_id,
  321. image_info,
  322. table_info,
  323. sorted_text_remain_text_block,
  324. layout_bboxes,
  325. inline_eq_info,
  326. interline_eq_info,
  327. page.get_text("dict", flags=fitz.TEXTFLAGS_TEXT)["blocks"],
  328. removed_text_blocks=removed_text_blocks,
  329. removed_image_blocks=removed_images,
  330. images_backup=images_backup,
  331. droped_table_block=[],
  332. table_backup=[],
  333. layout_tree=layout_tree,
  334. page_w=page.rect.width,
  335. page_h=page.rect.height,
  336. footnote_bboxes_tmp=footnote_bboxes_tmp
  337. )
  338. pdf_info_dict[f"page_{page_id}"] = page_info
  339. # end page for
  340. '''计算后处理阶段耗时'''
  341. start_time = time.time()
  342. """
  343. ==================================================================================================================================
  344. 去掉页眉和页脚,这里需要用到一定的统计量,所以放到最后
  345. 页眉和页脚主要从文本box和图片box中去除,位于页面的四周。
  346. 下面函数会直接修改pdf_info_dict,从文字块中、图片中删除属于页眉页脚的内容,删除内容做相对应记录
  347. """
  348. # 去页眉页脚
  349. header, footer = drop_footer_header(pdf_info_dict)
  350. """对单个layout内footnote和他下面的所有textbbox合并"""
  351. for page_key, page_info in pdf_info_dict.items():
  352. page_info = merge_footnote_blocks(page_info, main_text_font)
  353. page_info = remove_footnote_blocks(page_info)
  354. pdf_info_dict[page_key] = page_info
  355. """进入pdf后置过滤器,去掉一些不合理的pdf"""
  356. i = 0
  357. for page_info in pdf_info_dict.values():
  358. is_good_pdf, err = pdf_post_filter(page_info)
  359. if not is_good_pdf:
  360. logger.warning(f"page_id: {i}, drop this pdf: {book_name}, reason: {err}")
  361. if not debug_mode:
  362. return err
  363. i += 1
  364. if debug_mode:
  365. params_file_save_path = join_path(save_tmp_path, "md", book_name, "preproc_out.json")
  366. page_draw_rect_save_path = join_path(save_tmp_path, "md", book_name, "layout.pdf")
  367. # dir_path = os.path.dirname(page_draw_rect_save_path)
  368. # if not os.path.exists(dir_path):
  369. # # 如果目录不存在,创建它
  370. # os.makedirs(dir_path)
  371. with open(params_file_save_path, "w", encoding="utf-8") as f:
  372. json.dump(pdf_info_dict, f, ensure_ascii=False, indent=4)
  373. # 先检测本地 page_draw_rect_save_path 是否存在,如果存在则删除
  374. if os.path.exists(page_draw_rect_save_path):
  375. os.remove(page_draw_rect_save_path)
  376. # 绘制bbox和layout到pdf
  377. draw_bbox_on_page(pdf_docs, pdf_info_dict, page_draw_rect_save_path)
  378. draw_layout_bbox_on_page(pdf_docs, pdf_info_dict, header, footer, page_draw_rect_save_path)
  379. if debug_mode:
  380. # 打印后处理阶段耗时
  381. logger.info(f"post_processing_time: {get_delta_time(start_time)}")
  382. """
  383. ==================================================================================================================================
  384. 进入段落处理-2阶段
  385. """
  386. start_time = time.time()
  387. para_process_pipeline = ParaProcessPipeline()
  388. def _deal_with_text_exception(error_info):
  389. logger.warning(f"page_id: {page_id}, drop this pdf: {book_name}, reason: {error_info}")
  390. if error_info == denseSingleLineBlockException_msg:
  391. logger.warning(f"Drop this pdf: {book_name}, reason: {DropReason.DENSE_SINGLE_LINE_BLOCK}")
  392. result = {"need_drop": True, "drop_reason": DropReason.DENSE_SINGLE_LINE_BLOCK}
  393. return result
  394. if error_info == titleDetectionException_msg:
  395. logger.warning(f"Drop this pdf: {book_name}, reason: {DropReason.TITLE_DETECTION_FAILED}")
  396. result = {"need_drop": True, "drop_reason": DropReason.TITLE_DETECTION_FAILED}
  397. return result
  398. elif error_info == titleLevelException_msg:
  399. logger.warning(f"Drop this pdf: {book_name}, reason: {DropReason.TITLE_LEVEL_FAILED}")
  400. result = {"need_drop": True, "drop_reason": DropReason.TITLE_LEVEL_FAILED}
  401. return result
  402. elif error_info == paraSplitException_msg:
  403. logger.warning(f"Drop this pdf: {book_name}, reason: {DropReason.PARA_SPLIT_FAILED}")
  404. result = {"need_drop": True, "drop_reason": DropReason.PARA_SPLIT_FAILED}
  405. return result
  406. elif error_info == paraMergeException_msg:
  407. logger.warning(f"Drop this pdf: {book_name}, reason: {DropReason.PARA_MERGE_FAILED}")
  408. result = {"need_drop": True, "drop_reason": DropReason.PARA_MERGE_FAILED}
  409. return result
  410. if debug_mode:
  411. input_pdf_file = f"{pdf_local_path}.pdf"
  412. output_dir = f"{save_path}/{book_name}"
  413. output_pdf_file = f"{output_dir}/pdf_annos.pdf"
  414. """
  415. Call the para_process_pipeline function to process the pdf_info_dict.
  416. Parameters:
  417. para_debug_mode: str or None
  418. If para_debug_mode is None, the para_process_pipeline will not keep any intermediate results.
  419. 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.
  420. If para_debug_mode is "full", the para_process_pipeline will keep all the intermediate results generated during each step.
  421. """
  422. pdf_info_dict, error_info = para_process_pipeline.para_process_pipeline(
  423. pdf_info_dict,
  424. para_debug_mode="simple",
  425. input_pdf_path=input_pdf_file,
  426. output_pdf_path=output_pdf_file,
  427. )
  428. # 打印段落处理阶段耗时
  429. logger.info(f"para_process_time: {get_delta_time(start_time)}")
  430. # debug的时候不return drop信息
  431. if error_info is not None:
  432. _deal_with_text_exception(error_info)
  433. return pdf_info_dict
  434. else:
  435. pdf_info_dict, error_info = para_process_pipeline.para_process_pipeline(pdf_info_dict)
  436. if error_info is not None:
  437. return _deal_with_text_exception(error_info)
  438. return pdf_info_dict