pipeline.py 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720
  1. # coding=utf8
  2. import sys
  3. import time
  4. from urllib.parse import quote
  5. from magic_pdf.dict2md.ocr_mkcontent import ocr_mk_mm_markdown, ocr_mk_nlp_markdown_with_para, \
  6. ocr_mk_mm_markdown_with_para_and_pagination, ocr_mk_mm_markdown_with_para, ocr_mk_mm_standard_format, \
  7. make_standard_format_with_para
  8. from magic_pdf.libs.commons import (
  9. read_file,
  10. join_path,
  11. parse_bucket_key,
  12. formatted_time,
  13. s3_image_save_path,
  14. )
  15. from magic_pdf.libs.drop_reason import DropReason
  16. from magic_pdf.libs.json_compressor import JsonCompressor
  17. from magic_pdf.dict2md.mkcontent import mk_nlp_markdown, mk_universal_format
  18. from magic_pdf.pdf_parse_by_model import parse_pdf_by_model
  19. from magic_pdf.filter.pdf_classify_by_type import classify
  20. from magic_pdf.filter.pdf_meta_scan import pdf_meta_scan
  21. from loguru import logger
  22. from magic_pdf.pdf_parse_by_ocr import parse_pdf_by_ocr
  23. from magic_pdf.pdf_parse_for_train import parse_pdf_for_train
  24. from magic_pdf.train_utils.convert_to_train_format import convert_to_train_format
  25. from app.common.s3 import get_s3_config, get_s3_client
  26. def exception_handler(jso: dict, e):
  27. logger.exception(e)
  28. jso["need_drop"] = True
  29. jso["drop_reason"] = DropReason.Exception
  30. jso["exception"] = f"ERROR: {e}"
  31. return jso
  32. def get_data_type(jso: dict):
  33. data_type = jso.get("data_type")
  34. if data_type is None:
  35. data_type = jso.get("file_type")
  36. return data_type
  37. def get_bookid(jso: dict):
  38. book_id = jso.get("bookid")
  39. if book_id is None:
  40. book_id = jso.get("original_file_id")
  41. return book_id
  42. def get_data_source(jso: dict):
  43. data_source = jso.get("data_source")
  44. if data_source is None:
  45. data_source = jso.get("file_source")
  46. return data_source
  47. def meta_scan(jso: dict, doc_layout_check=True) -> dict:
  48. s3_pdf_path = jso.get("file_location")
  49. s3_config = get_s3_config(s3_pdf_path)
  50. if doc_layout_check:
  51. if (
  52. "doc_layout_result" not in jso
  53. ): # 检测json中是存在模型数据,如果没有则需要跳过该pdf
  54. jso["need_drop"] = True
  55. jso["drop_reason"] = DropReason.MISS_DOC_LAYOUT_RESULT
  56. return jso
  57. try:
  58. data_source = get_data_source(jso)
  59. file_id = jso.get("file_id")
  60. book_name = f"{data_source}/{file_id}"
  61. # 首页存在超量drawing问题
  62. # special_pdf_list = ['zlib/zlib_21822650']
  63. # if book_name in special_pdf_list:
  64. # jso['need_drop'] = True
  65. # jso['drop_reason'] = DropReason.SPECIAL_PDF
  66. # return jso
  67. start_time = time.time() # 记录开始时间
  68. logger.info(
  69. f"book_name is:{book_name},start_time is:{formatted_time(start_time)}",
  70. file=sys.stderr,
  71. )
  72. file_content = read_file(s3_pdf_path, s3_config)
  73. read_file_time = int(time.time() - start_time) # 计算执行时间
  74. start_time = time.time() # 记录开始时间
  75. res = pdf_meta_scan(s3_pdf_path, file_content)
  76. if res.get(
  77. "need_drop", False
  78. ): # 如果返回的字典里有need_drop,则提取drop_reason并跳过本次解析
  79. jso["need_drop"] = True
  80. jso["drop_reason"] = res["drop_reason"]
  81. else: # 正常返回
  82. jso["pdf_meta"] = res
  83. jso["content"] = ""
  84. jso["remark"] = ""
  85. jso["data_url"] = ""
  86. end_time = time.time() # 记录结束时间
  87. meta_scan_time = int(end_time - start_time) # 计算执行时间
  88. logger.info(
  89. f"book_name is:{book_name},end_time is:{formatted_time(end_time)},read_file_time is:{read_file_time},meta_scan_time is:{meta_scan_time}",
  90. file=sys.stderr,
  91. )
  92. jso["read_file_time"] = read_file_time
  93. jso["meta_scan_time"] = meta_scan_time
  94. except Exception as e:
  95. jso = exception_handler(jso, e)
  96. return jso
  97. def classify_by_type(jso: dict, debug_mode=False) -> dict:
  98. # 检测debug开关
  99. if debug_mode:
  100. pass
  101. else: # 如果debug没开,则检测是否有needdrop字段
  102. if jso.get("need_drop", False):
  103. return jso
  104. # 开始正式逻辑
  105. try:
  106. pdf_meta = jso.get("pdf_meta")
  107. data_source = get_data_source(jso)
  108. file_id = jso.get("file_id")
  109. book_name = f"{data_source}/{file_id}"
  110. total_page = pdf_meta["total_page"]
  111. page_width = pdf_meta["page_width_pts"]
  112. page_height = pdf_meta["page_height_pts"]
  113. img_sz_list = pdf_meta["image_info_per_page"]
  114. img_num_list = pdf_meta["imgs_per_page"]
  115. text_len_list = pdf_meta["text_len_per_page"]
  116. text_layout_list = pdf_meta["text_layout_per_page"]
  117. text_language = pdf_meta["text_language"]
  118. # allow_language = ['zh', 'en'] # 允许的语言,目前只允许简中和英文的
  119. # if text_language not in allow_language: # 如果语言不在允许的语言中,则drop
  120. # jso['need_drop'] = True
  121. # jso['drop_reason'] = DropReason.NOT_ALLOW_LANGUAGE
  122. # return jso
  123. pdf_path = pdf_meta["pdf_path"]
  124. is_encrypted = pdf_meta["is_encrypted"]
  125. is_needs_password = pdf_meta["is_needs_password"]
  126. if (
  127. is_encrypted or is_needs_password
  128. ): # 加密的,需要密码的,没有页面的,都不处理
  129. jso["need_drop"] = True
  130. jso["drop_reason"] = DropReason.ENCRYPTED
  131. else:
  132. start_time = time.time() # 记录开始时间
  133. is_text_pdf, results = classify(
  134. pdf_path,
  135. total_page,
  136. page_width,
  137. page_height,
  138. img_sz_list,
  139. text_len_list,
  140. img_num_list,
  141. text_layout_list,
  142. )
  143. classify_time = int(time.time() - start_time) # 计算执行时间
  144. if is_text_pdf:
  145. pdf_meta["is_text_pdf"] = is_text_pdf
  146. jso["pdf_meta"] = pdf_meta
  147. jso["classify_time"] = classify_time
  148. # print(json.dumps(pdf_meta, ensure_ascii=False))
  149. allow_language = ["zh", "en"] # 允许的语言,目前只允许简中和英文的
  150. if (
  151. text_language not in allow_language
  152. ): # 如果语言不在允许的语言中,则drop
  153. jso["need_drop"] = True
  154. jso["drop_reason"] = DropReason.NOT_ALLOW_LANGUAGE
  155. return jso
  156. else:
  157. # 先不drop
  158. pdf_meta["is_text_pdf"] = is_text_pdf
  159. jso["pdf_meta"] = pdf_meta
  160. jso["classify_time"] = classify_time
  161. jso["need_drop"] = True
  162. jso["drop_reason"] = DropReason.NOT_IS_TEXT_PDF
  163. extra_info = {"classify_rules": []}
  164. for condition, result in results.items():
  165. if not result:
  166. extra_info["classify_rules"].append(condition)
  167. jso["extra_info"] = extra_info
  168. except Exception as e:
  169. jso = exception_handler(jso, e)
  170. return jso
  171. def save_tables_to_s3(jso: dict, debug_mode=False) -> dict:
  172. if debug_mode:
  173. pass
  174. else: # 如果debug没开,则检测是否有needdrop字段
  175. if jso.get("need_drop", False):
  176. logger.info(
  177. f"book_name is:{get_data_source(jso)}/{jso['file_id']} need drop",
  178. file=sys.stderr,
  179. )
  180. jso["dropped"] = True
  181. return jso
  182. try:
  183. data_source = get_data_source(jso)
  184. file_id = jso.get("file_id")
  185. book_name = f"{data_source}/{file_id}"
  186. title = jso.get("title")
  187. url_encode_title = quote(title, safe="")
  188. if data_source != "scihub":
  189. return jso
  190. pdf_intermediate_dict = jso["pdf_intermediate_dict"]
  191. # 将 pdf_intermediate_dict 解压
  192. pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
  193. i = 0
  194. for page in pdf_intermediate_dict.values():
  195. if page.get("tables"):
  196. if len(page["tables"]) > 0:
  197. j = 0
  198. for table in page["tables"]:
  199. if debug_mode:
  200. image_path = join_path(
  201. "s3://mllm-raw-media/pdf2md_img/",
  202. book_name,
  203. table["image_path"],
  204. )
  205. else:
  206. image_path = join_path(
  207. "s3://mllm-raw-media/pdf2md_img/", table["image_path"]
  208. )
  209. if image_path.endswith(".jpg"):
  210. j += 1
  211. s3_client = get_s3_client(image_path)
  212. bucket_name, bucket_key = parse_bucket_key(image_path)
  213. # 通过s3_client获取图片到内存
  214. image_bytes = s3_client.get_object(
  215. Bucket=bucket_name, Key=bucket_key
  216. )["Body"].read()
  217. # 保存图片到新的位置
  218. if debug_mode:
  219. new_image_path = join_path(
  220. "s3://mllm-raw-media/pdf2md_img/table_new/",
  221. url_encode_title
  222. + "_"
  223. + table["image_path"].lstrip("tables/"),
  224. )
  225. else:
  226. new_image_path = join_path(
  227. "s3://mllm-raw-media/pdf2md_img/table_new/",
  228. url_encode_title + f"_page{i}_{j}.jpg",
  229. )
  230. logger.info(new_image_path, file=sys.stderr)
  231. bucket_name, bucket_key = parse_bucket_key(new_image_path)
  232. s3_client.put_object(
  233. Bucket=bucket_name, Key=bucket_key, Body=image_bytes
  234. )
  235. else:
  236. continue
  237. i += 1
  238. # 把无用的信息清空
  239. jso["doc_layout_result"] = ""
  240. jso["pdf_intermediate_dict"] = ""
  241. jso["pdf_meta"] = ""
  242. except Exception as e:
  243. jso = exception_handler(jso, e)
  244. return jso
  245. def drop_needdrop_pdf(jso: dict) -> dict:
  246. if jso.get("need_drop", False):
  247. logger.info(
  248. f"book_name is:{get_data_source(jso)}/{jso['file_id']} need drop",
  249. file=sys.stderr,
  250. )
  251. jso["dropped"] = True
  252. return jso
  253. def pdf_intermediate_dict_to_markdown(jso: dict, debug_mode=False) -> dict:
  254. if debug_mode:
  255. pass
  256. else: # 如果debug没开,则检测是否有needdrop字段
  257. if jso.get("need_drop", False):
  258. book_name = join_path(get_data_source(jso), jso["file_id"])
  259. logger.info(f"book_name is:{book_name} need drop", file=sys.stderr)
  260. jso["dropped"] = True
  261. return jso
  262. try:
  263. pdf_intermediate_dict = jso["pdf_intermediate_dict"]
  264. # 将 pdf_intermediate_dict 解压
  265. pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
  266. # markdown_content = mk_nlp_markdown(pdf_intermediate_dict)
  267. jso["content_list"] = mk_universal_format(pdf_intermediate_dict)
  268. # jso["content"] = markdown_content
  269. logger.info(f"book_name is:{get_data_source(jso)}/{jso['file_id']}")
  270. # 把无用的信息清空
  271. jso["doc_layout_result"] = ""
  272. jso["pdf_intermediate_dict"] = ""
  273. jso["pdf_meta"] = ""
  274. except Exception as e:
  275. jso = exception_handler(jso, e)
  276. return jso
  277. def parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
  278. # 检测debug开关
  279. if debug_mode:
  280. pass
  281. else: # 如果debug没开,则检测是否有needdrop字段
  282. if jso.get("need_drop", False):
  283. return jso
  284. # 开始正式逻辑
  285. s3_pdf_path = jso.get("file_location")
  286. s3_config = get_s3_config(s3_pdf_path)
  287. model_output_json_list = jso.get("doc_layout_result")
  288. data_source = get_data_source(jso)
  289. file_id = jso.get("file_id")
  290. book_name = f"{data_source}/{file_id}"
  291. # 1.23.22已修复
  292. # if debug_mode:
  293. # pass
  294. # else:
  295. # if book_name == "zlib/zlib_21929367":
  296. # jso['need_drop'] = True
  297. # jso['drop_reason'] = DropReason.SPECIAL_PDF
  298. # return jso
  299. junk_img_bojids = jso["pdf_meta"]["junk_img_bojids"]
  300. # total_page = jso['pdf_meta']['total_page']
  301. # 增加检测 max_svgs 数量的检测逻辑,如果 max_svgs 超过3000则drop
  302. svgs_per_page_list = jso["pdf_meta"]["svgs_per_page"]
  303. max_svgs = max(svgs_per_page_list)
  304. if max_svgs > 3000:
  305. jso["need_drop"] = True
  306. jso["drop_reason"] = DropReason.HIGH_COMPUTATIONAL_lOAD_BY_SVGS
  307. # elif total_page > 1000:
  308. # jso['need_drop'] = True
  309. # jso['drop_reason'] = DropReason.HIGH_COMPUTATIONAL_lOAD_BY_TOTAL_PAGES
  310. else:
  311. try:
  312. save_path = s3_image_save_path
  313. image_s3_config = get_s3_config(save_path)
  314. start_time = time.time() # 记录开始时间
  315. # 先打印一下book_name和解析开始的时间
  316. logger.info(
  317. f"book_name is:{book_name},start_time is:{formatted_time(start_time)}",
  318. file=sys.stderr,
  319. )
  320. pdf_info_dict = parse_pdf_by_model(
  321. s3_pdf_path,
  322. s3_config,
  323. model_output_json_list,
  324. save_path,
  325. book_name,
  326. pdf_model_profile=None,
  327. image_s3_config=image_s3_config,
  328. start_page_id=start_page_id,
  329. junk_img_bojids=junk_img_bojids,
  330. debug_mode=debug_mode,
  331. )
  332. if pdf_info_dict.get(
  333. "need_drop", False
  334. ): # 如果返回的字典里有need_drop,则提取drop_reason并跳过本次解析
  335. jso["need_drop"] = True
  336. jso["drop_reason"] = pdf_info_dict["drop_reason"]
  337. else: # 正常返回,将 pdf_info_dict 压缩并存储
  338. pdf_info_dict = JsonCompressor.compress_json(pdf_info_dict)
  339. jso["pdf_intermediate_dict"] = pdf_info_dict
  340. end_time = time.time() # 记录完成时间
  341. parse_time = int(end_time - start_time) # 计算执行时间
  342. # 解析完成后打印一下book_name和耗时
  343. logger.info(
  344. f"book_name is:{book_name},end_time is:{formatted_time(end_time)},cost_time is:{parse_time}",
  345. file=sys.stderr,
  346. )
  347. jso["parse_time"] = parse_time
  348. except Exception as e:
  349. jso = exception_handler(jso, e)
  350. return jso
  351. """
  352. 统一处理逻辑
  353. 1.先调用parse_pdf对文本类pdf进行处理
  354. 2.再调用ocr_dropped_parse_pdf,对之前drop的pdf进行处理
  355. """
  356. def uni_parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
  357. jso = parse_pdf(jso, start_page_id=start_page_id, debug_mode=debug_mode)
  358. jso = ocr_dropped_parse_pdf(jso, start_page_id=start_page_id, debug_mode=debug_mode)
  359. return jso
  360. def parse_pdf_for_model_train(jso: dict, start_page_id=0, debug_mode=False) -> dict:
  361. # 检测debug开关
  362. if debug_mode:
  363. pass
  364. else: # 如果debug没开,则检测是否有needdrop字段
  365. if jso.get("need_drop", False):
  366. return jso
  367. # 开始正式逻辑
  368. s3_pdf_path = jso.get("file_location")
  369. s3_config = get_s3_config(s3_pdf_path)
  370. model_output_json_list = jso.get("doc_layout_result")
  371. data_source = get_data_source(jso)
  372. file_id = jso.get("file_id")
  373. book_name = f"{data_source}/{file_id}"
  374. # 1.23.22已修复
  375. # if debug_mode:
  376. # pass
  377. # else:
  378. # if book_name == "zlib/zlib_21929367":
  379. # jso['need_drop'] = True
  380. # jso['drop_reason'] = DropReason.SPECIAL_PDF
  381. # return jso
  382. junk_img_bojids = jso["pdf_meta"]["junk_img_bojids"]
  383. # total_page = jso['pdf_meta']['total_page']
  384. # 增加检测 max_svgs 数量的检测逻辑,如果 max_svgs 超过3000则drop
  385. svgs_per_page_list = jso["pdf_meta"]["svgs_per_page"]
  386. max_svgs = max(svgs_per_page_list)
  387. if max_svgs > 3000:
  388. jso["need_drop"] = True
  389. jso["drop_reason"] = DropReason.HIGH_COMPUTATIONAL_lOAD_BY_SVGS
  390. # elif total_page > 1000:
  391. # jso['need_drop'] = True
  392. # jso['drop_reason'] = DropReason.HIGH_COMPUTATIONAL_lOAD_BY_TOTAL_PAGES
  393. else:
  394. try:
  395. save_path = s3_image_save_path
  396. image_s3_config = get_s3_config(save_path)
  397. start_time = time.time() # 记录开始时间
  398. # 先打印一下book_name和解析开始的时间
  399. logger.info(
  400. f"book_name is:{book_name},start_time is:{formatted_time(start_time)}",
  401. file=sys.stderr,
  402. )
  403. pdf_info_dict = parse_pdf_for_train(
  404. s3_pdf_path,
  405. s3_config,
  406. model_output_json_list,
  407. save_path,
  408. book_name,
  409. pdf_model_profile=None,
  410. image_s3_config=image_s3_config,
  411. start_page_id=start_page_id,
  412. junk_img_bojids=junk_img_bojids,
  413. debug_mode=debug_mode,
  414. )
  415. if pdf_info_dict.get(
  416. "need_drop", False
  417. ): # 如果返回的字典里有need_drop,则提取drop_reason并跳过本次解析
  418. jso["need_drop"] = True
  419. jso["drop_reason"] = pdf_info_dict["drop_reason"]
  420. else: # 正常返回,将 pdf_info_dict 压缩并存储
  421. jso["parsed_results"] = convert_to_train_format(pdf_info_dict)
  422. pdf_info_dict = JsonCompressor.compress_json(pdf_info_dict)
  423. jso["pdf_intermediate_dict"] = pdf_info_dict
  424. end_time = time.time() # 记录完成时间
  425. parse_time = int(end_time - start_time) # 计算执行时间
  426. # 解析完成后打印一下book_name和耗时
  427. logger.info(
  428. f"book_name is:{book_name},end_time is:{formatted_time(end_time)},cost_time is:{parse_time}",
  429. file=sys.stderr,
  430. )
  431. jso["parse_time"] = parse_time
  432. except Exception as e:
  433. jso = exception_handler(jso, e)
  434. return jso
  435. # 专门用来跑被drop的pdf,跑完之后需要把need_drop字段置为false
  436. def ocr_dropped_parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
  437. if not jso.get("need_drop", False):
  438. return jso
  439. else:
  440. jso = ocr_parse_pdf_core(
  441. jso, start_page_id=start_page_id, debug_mode=debug_mode
  442. )
  443. jso["need_drop"] = False
  444. return jso
  445. def ocr_parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
  446. # 检测debug开关
  447. if debug_mode:
  448. pass
  449. else: # 如果debug没开,则检测是否有needdrop字段
  450. if jso.get("need_drop", False):
  451. return jso
  452. jso = ocr_parse_pdf_core(jso, start_page_id=start_page_id, debug_mode=debug_mode)
  453. return jso
  454. def ocr_parse_pdf_core(jso: dict, start_page_id=0, debug_mode=False) -> dict:
  455. s3_pdf_path = jso.get("file_location")
  456. s3_config = get_s3_config(s3_pdf_path)
  457. model_output_json_list = jso.get("doc_layout_result")
  458. data_source = get_data_source(jso)
  459. file_id = jso.get("file_id")
  460. book_name = f"{data_source}/{file_id}"
  461. try:
  462. save_path = s3_image_save_path
  463. image_s3_config = get_s3_config(save_path)
  464. start_time = time.time() # 记录开始时间
  465. # 先打印一下book_name和解析开始的时间
  466. logger.info(
  467. f"book_name is:{book_name},start_time is:{formatted_time(start_time)}",
  468. file=sys.stderr,
  469. )
  470. pdf_info_dict = parse_pdf_by_ocr(
  471. s3_pdf_path,
  472. s3_config,
  473. model_output_json_list,
  474. save_path,
  475. book_name,
  476. pdf_model_profile=None,
  477. image_s3_config=image_s3_config,
  478. start_page_id=start_page_id,
  479. debug_mode=debug_mode,
  480. )
  481. pdf_info_dict = JsonCompressor.compress_json(pdf_info_dict)
  482. jso["pdf_intermediate_dict"] = pdf_info_dict
  483. end_time = time.time() # 记录完成时间
  484. parse_time = int(end_time - start_time) # 计算执行时间
  485. # 解析完成后打印一下book_name和耗时
  486. logger.info(
  487. f"book_name is:{book_name},end_time is:{formatted_time(end_time)},cost_time is:{parse_time}",
  488. file=sys.stderr,
  489. )
  490. jso["parse_time"] = parse_time
  491. except Exception as e:
  492. jso = exception_handler(jso, e)
  493. return jso
  494. def ocr_pdf_intermediate_dict_to_markdown(jso: dict, debug_mode=False) -> dict:
  495. if debug_mode:
  496. pass
  497. else: # 如果debug没开,则检测是否有needdrop字段
  498. if jso.get("need_drop", False):
  499. book_name = join_path(get_data_source(jso), jso["file_id"])
  500. logger.info(f"book_name is:{book_name} need drop", file=sys.stderr)
  501. jso["dropped"] = True
  502. return jso
  503. try:
  504. pdf_intermediate_dict = jso["pdf_intermediate_dict"]
  505. # 将 pdf_intermediate_dict 解压
  506. pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
  507. markdown_content = ocr_mk_mm_markdown(pdf_intermediate_dict)
  508. jso["content"] = markdown_content
  509. logger.info(
  510. f"book_name is:{get_data_source(jso)}/{jso['file_id']},markdown content length is {len(markdown_content)}",
  511. file=sys.stderr,
  512. )
  513. # 把无用的信息清空
  514. jso["doc_layout_result"] = ""
  515. jso["pdf_intermediate_dict"] = ""
  516. jso["pdf_meta"] = ""
  517. except Exception as e:
  518. jso = exception_handler(jso, e)
  519. return jso
  520. def ocr_pdf_intermediate_dict_to_markdown_with_para(jso: dict, debug_mode=False) -> dict:
  521. if debug_mode:
  522. pass
  523. else: # 如果debug没开,则检测是否有needdrop字段
  524. if jso.get("need_drop", False):
  525. book_name = join_path(get_data_source(jso), jso["file_id"])
  526. logger.info(f"book_name is:{book_name} need drop", file=sys.stderr)
  527. jso["dropped"] = True
  528. return jso
  529. try:
  530. pdf_intermediate_dict = jso["pdf_intermediate_dict"]
  531. # 将 pdf_intermediate_dict 解压
  532. pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
  533. # markdown_content = ocr_mk_mm_markdown_with_para(pdf_intermediate_dict)
  534. markdown_content = ocr_mk_nlp_markdown_with_para(pdf_intermediate_dict)
  535. jso["content"] = markdown_content
  536. logger.info(
  537. f"book_name is:{get_data_source(jso)}/{jso['file_id']},markdown content length is {len(markdown_content)}",
  538. file=sys.stderr,
  539. )
  540. # 把无用的信息清空
  541. jso["doc_layout_result"] = ""
  542. jso["pdf_intermediate_dict"] = ""
  543. jso["pdf_meta"] = ""
  544. except Exception as e:
  545. jso = exception_handler(jso, e)
  546. return jso
  547. def ocr_pdf_intermediate_dict_to_markdown_with_para_and_pagination(jso: dict, debug_mode=False) -> dict:
  548. if debug_mode:
  549. pass
  550. else: # 如果debug没开,则检测是否有needdrop字段
  551. if jso.get("need_drop", False):
  552. book_name = join_path(get_data_source(jso), jso["file_id"])
  553. logger.info(f"book_name is:{book_name} need drop", file=sys.stderr)
  554. jso["dropped"] = True
  555. return jso
  556. try:
  557. pdf_intermediate_dict = jso["pdf_intermediate_dict"]
  558. # 将 pdf_intermediate_dict 解压
  559. pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
  560. markdown_content = ocr_mk_mm_markdown_with_para_and_pagination(pdf_intermediate_dict)
  561. jso["content"] = markdown_content
  562. logger.info(
  563. f"book_name is:{get_data_source(jso)}/{jso['file_id']},markdown content length is {len(markdown_content)}",
  564. file=sys.stderr,
  565. )
  566. # 把无用的信息清空
  567. # jso["doc_layout_result"] = ""
  568. jso["pdf_intermediate_dict"] = ""
  569. # jso["pdf_meta"] = ""
  570. except Exception as e:
  571. jso = exception_handler(jso, e)
  572. return jso
  573. def ocr_pdf_intermediate_dict_to_markdown_with_para_for_qa(
  574. jso: dict, debug_mode=False
  575. ) -> dict:
  576. if debug_mode:
  577. pass
  578. else: # 如果debug没开,则检测是否有needdrop字段
  579. if jso.get("need_drop", False):
  580. book_name = join_path(get_data_source(jso), jso["file_id"])
  581. logger.info(f"book_name is:{book_name} need drop", file=sys.stderr)
  582. jso["dropped"] = True
  583. return jso
  584. try:
  585. pdf_intermediate_dict = jso["pdf_intermediate_dict"]
  586. # 将 pdf_intermediate_dict 解压
  587. pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
  588. markdown_content = ocr_mk_mm_markdown_with_para(pdf_intermediate_dict)
  589. jso["content_ocr"] = markdown_content
  590. logger.info(
  591. f"book_name is:{get_data_source(jso)}/{jso['file_id']},markdown content length is {len(markdown_content)}",
  592. file=sys.stderr,
  593. )
  594. # 把无用的信息清空
  595. jso["doc_layout_result"] = ""
  596. jso["pdf_intermediate_dict"] = ""
  597. jso["mid_json_ocr"] = pdf_intermediate_dict
  598. jso["pdf_meta"] = ""
  599. except Exception as e:
  600. jso = exception_handler(jso, e)
  601. return jso
  602. def ocr_pdf_intermediate_dict_to_standard_format(jso: dict, debug_mode=False) -> dict:
  603. if debug_mode:
  604. pass
  605. else: # 如果debug没开,则检测是否有needdrop字段
  606. if jso.get("need_drop", False):
  607. book_name = join_path(get_data_source(jso), jso["file_id"])
  608. logger.info(f"book_name is:{book_name} need drop", file=sys.stderr)
  609. jso["dropped"] = True
  610. return jso
  611. try:
  612. pdf_intermediate_dict = jso["pdf_intermediate_dict"]
  613. # 将 pdf_intermediate_dict 解压
  614. pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
  615. standard_format = ocr_mk_mm_standard_format(pdf_intermediate_dict)
  616. jso["content_list"] = standard_format
  617. logger.info(
  618. f"book_name is:{get_data_source(jso)}/{jso['file_id']},content_list length is {len(standard_format)}",
  619. file=sys.stderr,
  620. )
  621. # 把无用的信息清空
  622. jso["doc_layout_result"] = ""
  623. jso["pdf_intermediate_dict"] = ""
  624. jso["pdf_meta"] = ""
  625. except Exception as e:
  626. jso = exception_handler(jso, e)
  627. return jso
  628. def ocr_pdf_intermediate_dict_to_standard_format_with_para(jso: dict, debug_mode=False) -> dict:
  629. if debug_mode:
  630. pass
  631. else: # 如果debug没开,则检测是否有needdrop字段
  632. if jso.get("need_drop", False):
  633. book_name = join_path(get_data_source(jso), jso["file_id"])
  634. logger.info(f"book_name is:{book_name} need drop", file=sys.stderr)
  635. jso["dropped"] = True
  636. return jso
  637. try:
  638. pdf_intermediate_dict = jso["pdf_intermediate_dict"]
  639. # 将 pdf_intermediate_dict 解压
  640. pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
  641. standard_format = make_standard_format_with_para(pdf_intermediate_dict)
  642. jso["content_list"] = standard_format
  643. logger.info(
  644. f"book_name is:{get_data_source(jso)}/{jso['file_id']},content_list length is {len(standard_format)}",
  645. file=sys.stderr,
  646. )
  647. # 把无用的信息清空
  648. jso["doc_layout_result"] = ""
  649. jso["pdf_intermediate_dict"] = ""
  650. jso["pdf_meta"] = ""
  651. except Exception as e:
  652. jso = exception_handler(jso, e)
  653. return jso
  654. if __name__ == "__main__":
  655. pass