block_continuation_processor.py 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563
  1. import os
  2. import sys
  3. import unicodedata
  4. from para.commons import *
  5. if sys.version_info[0] >= 3:
  6. sys.stdout.reconfigure(encoding="utf-8") # type: ignore
  7. class BlockContinuationProcessor:
  8. """
  9. This class is used to process the blocks to detect block continuations.
  10. """
  11. def __init__(self) -> None:
  12. pass
  13. def __is_similar_font_type(self, font_type1, font_type2, prefix_length_ratio=0.3):
  14. """
  15. This function checks if the two font types are similar.
  16. Definition of similar font types: the two font types have a common prefix,
  17. and the length of the common prefix is at least a certain ratio of the length of the shorter font type.
  18. Parameters
  19. ----------
  20. font_type1 : str
  21. font type 1
  22. font_type2 : str
  23. font type 2
  24. prefix_length_ratio : float
  25. minimum ratio of the common prefix length to the length of the shorter font type
  26. Returns
  27. -------
  28. bool
  29. True if the two font types are similar, False otherwise.
  30. """
  31. if isinstance(font_type1, list):
  32. font_type1 = font_type1[0] if font_type1 else ""
  33. if isinstance(font_type2, list):
  34. font_type2 = font_type2[0] if font_type2 else ""
  35. if font_type1 == font_type2:
  36. return True
  37. # Find the length of the common prefix
  38. common_prefix_length = len(os.path.commonprefix([font_type1, font_type2]))
  39. # Calculate the minimum prefix length based on the ratio
  40. min_prefix_length = int(min(len(font_type1), len(font_type2)) * prefix_length_ratio)
  41. return common_prefix_length >= min_prefix_length
  42. def __is_same_block_font(self, block1, block2):
  43. """
  44. This function compares the font of block1 and block2
  45. Parameters
  46. ----------
  47. block1 : dict
  48. block1
  49. block2 : dict
  50. block2
  51. Returns
  52. -------
  53. is_same : bool
  54. True if block1 and block2 have the same font, else False
  55. """
  56. block_1_font_type = safe_get(block1, "block_font_type", "")
  57. block_1_font_size = safe_get(block1, "block_font_size", 0)
  58. block_1_avg_char_width = safe_get(block1, "avg_char_width", 0)
  59. block_2_font_type = safe_get(block2, "block_font_type", "")
  60. block_2_font_size = safe_get(block2, "block_font_size", 0)
  61. block_2_avg_char_width = safe_get(block2, "avg_char_width", 0)
  62. if isinstance(block_1_font_size, list):
  63. block_1_font_size = block_1_font_size[0] if block_1_font_size else 0
  64. if isinstance(block_2_font_size, list):
  65. block_2_font_size = block_2_font_size[0] if block_2_font_size else 0
  66. block_1_text = safe_get(block1, "text", "")
  67. block_2_text = safe_get(block2, "text", "")
  68. if block_1_avg_char_width == 0 or block_2_avg_char_width == 0:
  69. return False
  70. if not block_1_text or not block_2_text:
  71. return False
  72. else:
  73. text_len_ratio = len(block_2_text) / len(block_1_text)
  74. if text_len_ratio < 0.2:
  75. avg_char_width_condition = (
  76. abs(block_1_avg_char_width - block_2_avg_char_width) / min(block_1_avg_char_width, block_2_avg_char_width)
  77. < 0.5
  78. )
  79. else:
  80. avg_char_width_condition = (
  81. abs(block_1_avg_char_width - block_2_avg_char_width) / min(block_1_avg_char_width, block_2_avg_char_width)
  82. < 0.2
  83. )
  84. block_font_size_condtion = abs(block_1_font_size - block_2_font_size) < 1
  85. return (
  86. self.__is_similar_font_type(block_1_font_type, block_2_font_type)
  87. and avg_char_width_condition
  88. and block_font_size_condtion
  89. )
  90. def _is_alphabet_char(self, char):
  91. if (char >= "\u0041" and char <= "\u005a") or (char >= "\u0061" and char <= "\u007a"):
  92. return True
  93. else:
  94. return False
  95. def _is_chinese_char(self, char):
  96. if char >= "\u4e00" and char <= "\u9fa5":
  97. return True
  98. else:
  99. return False
  100. def _is_other_letter_char(self, char):
  101. try:
  102. cat = unicodedata.category(char)
  103. if cat == "Lu" or cat == "Ll":
  104. return not self._is_alphabet_char(char) and not self._is_chinese_char(char)
  105. except TypeError:
  106. print("The input to the function must be a single character.")
  107. return False
  108. def _is_year(self, s: str):
  109. try:
  110. number = int(s)
  111. return 1900 <= number <= 2099
  112. except ValueError:
  113. return False
  114. def __is_para_font_consistent(self, para_1, para_2):
  115. """
  116. This function compares the font of para1 and para2
  117. Parameters
  118. ----------
  119. para1 : dict
  120. para1
  121. para2 : dict
  122. para2
  123. Returns
  124. -------
  125. is_same : bool
  126. True if para1 and para2 have the same font, else False
  127. """
  128. if para_1 is None or para_2 is None:
  129. return False
  130. para_1_font_type = safe_get(para_1, "para_font_type", "")
  131. para_1_font_size = safe_get(para_1, "para_font_size", 0)
  132. para_1_font_color = safe_get(para_1, "para_font_color", "")
  133. para_2_font_type = safe_get(para_2, "para_font_type", "")
  134. para_2_font_size = safe_get(para_2, "para_font_size", 0)
  135. para_2_font_color = safe_get(para_2, "para_font_color", "")
  136. if isinstance(para_1_font_type, list): # get the most common font type
  137. para_1_font_type = max(set(para_1_font_type), key=para_1_font_type.count)
  138. if isinstance(para_2_font_type, list):
  139. para_2_font_type = max(set(para_2_font_type), key=para_2_font_type.count)
  140. if isinstance(para_1_font_size, list): # compute average font type
  141. para_1_font_size = sum(para_1_font_size) / len(para_1_font_size)
  142. if isinstance(para_2_font_size, list): # compute average font type
  143. para_2_font_size = sum(para_2_font_size) / len(para_2_font_size)
  144. return (
  145. self.__is_similar_font_type(para_1_font_type, para_2_font_type)
  146. and abs(para_1_font_size - para_2_font_size) < 1.5
  147. # and para_font_color1 == para_font_color2
  148. )
  149. def _is_para_puncs_consistent(self, para_1, para_2):
  150. """
  151. This function determines whether para1 and para2 are originally from the same paragraph by checking the puncs of para1(former) and para2(latter)
  152. Parameters
  153. ----------
  154. para1 : dict
  155. para1
  156. para2 : dict
  157. para2
  158. Returns
  159. -------
  160. is_same : bool
  161. True if para1 and para2 are from the same paragraph by using the puncs, else False
  162. """
  163. para_1_text = safe_get(para_1, "para_text", "").strip()
  164. para_2_text = safe_get(para_2, "para_text", "").strip()
  165. para_1_bboxes = safe_get(para_1, "para_bbox", [])
  166. para_1_font_sizes = safe_get(para_1, "para_font_size", 0)
  167. para_2_bboxes = safe_get(para_2, "para_bbox", [])
  168. para_2_font_sizes = safe_get(para_2, "para_font_size", 0)
  169. # print_yellow(" Features of determine puncs_consistent:")
  170. # print(f" para_1_text: {para_1_text}")
  171. # print(f" para_2_text: {para_2_text}")
  172. # print(f" para_1_bboxes: {para_1_bboxes}")
  173. # print(f" para_2_bboxes: {para_2_bboxes}")
  174. # print(f" para_1_font_sizes: {para_1_font_sizes}")
  175. # print(f" para_2_font_sizes: {para_2_font_sizes}")
  176. if is_nested_list(para_1_bboxes):
  177. x0_1, y0_1, x1_1, y1_1 = para_1_bboxes[-1]
  178. else:
  179. x0_1, y0_1, x1_1, y1_1 = para_1_bboxes
  180. if is_nested_list(para_2_bboxes):
  181. x0_2, y0_2, x1_2, y1_2 = para_2_bboxes[0]
  182. para_2_font_sizes = para_2_font_sizes[0] # type: ignore
  183. else:
  184. x0_2, y0_2, x1_2, y1_2 = para_2_bboxes
  185. right_align_threshold = 0.5 * (para_1_font_sizes + para_2_font_sizes) * 0.8
  186. are_two_paras_right_aligned = abs(x1_1 - x1_2) < right_align_threshold
  187. left_indent_threshold = 0.5 * (para_1_font_sizes + para_2_font_sizes) * 0.8
  188. is_para1_left_indent_than_papa2 = x0_1 - x0_2 > left_indent_threshold
  189. is_para2_left_indent_than_papa1 = x0_2 - x0_1 > left_indent_threshold
  190. # Check if either para_text1 or para_text2 is empty
  191. if not para_1_text or not para_2_text:
  192. return False
  193. # Define the end puncs for a sentence to end and hyphen
  194. end_puncs = [".", "?", "!", "。", "?", "!", "…"]
  195. hyphen = ["-", "—"]
  196. # Check if para_text1 ends with either hyphen or non-end punctuation or spaces
  197. para_1_end_with_hyphen = para_1_text and para_1_text[-1] in hyphen
  198. para_1_end_with_end_punc = para_1_text and para_1_text[-1] in end_puncs
  199. para_1_end_with_space = para_1_text and para_1_text[-1] == " "
  200. para_1_not_end_with_end_punc = para_1_text and para_1_text[-1] not in end_puncs
  201. # print_yellow(f" para_1_end_with_hyphen: {para_1_end_with_hyphen}")
  202. # print_yellow(f" para_1_end_with_end_punc: {para_1_end_with_end_punc}")
  203. # print_yellow(f" para_1_not_end_with_end_punc: {para_1_not_end_with_end_punc}")
  204. # print_yellow(f" para_1_end_with_space: {para_1_end_with_space}")
  205. if para_1_end_with_hyphen: # If para_text1 ends with hyphen
  206. # print_red(f"para_1 is end with hyphen.")
  207. para_2_is_consistent = para_2_text and (
  208. para_2_text[0] in hyphen
  209. or (self._is_alphabet_char(para_2_text[0]) and para_2_text[0].islower())
  210. or (self._is_chinese_char(para_2_text[0]))
  211. or (self._is_other_letter_char(para_2_text[0]))
  212. )
  213. if para_2_is_consistent:
  214. # print(f"para_2 is consistent.\n")
  215. return True
  216. else:
  217. # print(f"para_2 is not consistent.\n")
  218. pass
  219. elif para_1_end_with_end_punc: # If para_text1 ends with ending punctuations
  220. # print_red(f"para_1 is end with end_punc.")
  221. para_2_is_consistent = (
  222. para_2_text
  223. and (
  224. para_2_text[0] == " "
  225. or (self._is_alphabet_char(para_2_text[0]) and para_2_text[0].isupper())
  226. or (self._is_chinese_char(para_2_text[0]))
  227. or (self._is_other_letter_char(para_2_text[0]))
  228. )
  229. and not is_para2_left_indent_than_papa1
  230. )
  231. if para_2_is_consistent:
  232. # print(f"para_2 is consistent.\n")
  233. return True
  234. else:
  235. # print(f"para_2 is not consistent.\n")
  236. pass
  237. elif para_1_not_end_with_end_punc: # If para_text1 is not end with ending punctuations
  238. # print_red(f"para_1 is NOT end with end_punc.")
  239. para_2_is_consistent = para_2_text and (
  240. para_2_text[0] == " "
  241. or (self._is_alphabet_char(para_2_text[0]) and para_2_text[0].islower())
  242. or (self._is_alphabet_char(para_2_text[0]))
  243. or (self._is_year(para_2_text[0:4]))
  244. or (are_two_paras_right_aligned or is_para1_left_indent_than_papa2)
  245. or (self._is_chinese_char(para_2_text[0]))
  246. or (self._is_other_letter_char(para_2_text[0]))
  247. )
  248. if para_2_is_consistent:
  249. # print(f"para_2 is consistent.\n")
  250. return True
  251. else:
  252. # print(f"para_2 is not consistent.\n")
  253. pass
  254. elif para_1_end_with_space: # If para_text1 ends with space
  255. # print_red(f"para_1 is end with space.")
  256. para_2_is_consistent = para_2_text and (
  257. para_2_text[0] == " "
  258. or (self._is_alphabet_char(para_2_text[0]) and para_2_text[0].islower())
  259. or (self._is_chinese_char(para_2_text[0]))
  260. or (self._is_other_letter_char(para_2_text[0]))
  261. )
  262. if para_2_is_consistent:
  263. # print(f"para_2 is consistent.\n")
  264. return True
  265. else:
  266. pass
  267. # print(f"para_2 is not consistent.\n")
  268. return False
  269. def _is_block_consistent(self, block1, block2):
  270. """
  271. This function determines whether block1 and block2 are originally from the same block
  272. Parameters
  273. ----------
  274. block1 : dict
  275. block1s
  276. block2 : dict
  277. block2
  278. Returns
  279. -------
  280. is_same : bool
  281. True if block1 and block2 are from the same block, else False
  282. """
  283. return self.__is_same_block_font(block1, block2)
  284. def _is_para_continued(self, para1, para2):
  285. """
  286. This function determines whether para1 and para2 are originally from the same paragraph
  287. Parameters
  288. ----------
  289. para1 : dict
  290. para1
  291. para2 : dict
  292. para2
  293. Returns
  294. -------
  295. is_same : bool
  296. True if para1 and para2 are from the same paragraph, else False
  297. """
  298. is_para_font_consistent = self.__is_para_font_consistent(para1, para2)
  299. is_para_puncs_consistent = self._is_para_puncs_consistent(para1, para2)
  300. return is_para_font_consistent and is_para_puncs_consistent
  301. def _are_boundaries_of_block_consistent(self, block1, block2):
  302. """
  303. This function checks if the boundaries of block1 and block2 are consistent
  304. Parameters
  305. ----------
  306. block1 : dict
  307. block1
  308. block2 : dict
  309. block2
  310. Returns
  311. -------
  312. is_consistent : bool
  313. True if the boundaries of block1 and block2 are consistent, else False
  314. """
  315. last_line_of_block1 = block1["lines"][-1]
  316. first_line_of_block2 = block2["lines"][0]
  317. spans_of_last_line_of_block1 = last_line_of_block1["spans"]
  318. spans_of_first_line_of_block2 = first_line_of_block2["spans"]
  319. font_type_of_last_line_of_block1 = spans_of_last_line_of_block1[0]["font"].lower()
  320. font_size_of_last_line_of_block1 = spans_of_last_line_of_block1[0]["size"]
  321. font_color_of_last_line_of_block1 = spans_of_last_line_of_block1[0]["color"]
  322. font_flags_of_last_line_of_block1 = spans_of_last_line_of_block1[0]["flags"]
  323. font_type_of_first_line_of_block2 = spans_of_first_line_of_block2[0]["font"].lower()
  324. font_size_of_first_line_of_block2 = spans_of_first_line_of_block2[0]["size"]
  325. font_color_of_first_line_of_block2 = spans_of_first_line_of_block2[0]["color"]
  326. font_flags_of_first_line_of_block2 = spans_of_first_line_of_block2[0]["flags"]
  327. return (
  328. self.__is_similar_font_type(font_type_of_last_line_of_block1, font_type_of_first_line_of_block2)
  329. and abs(font_size_of_last_line_of_block1 - font_size_of_first_line_of_block2) < 1
  330. # and font_color_of_last_line_of_block1 == font_color_of_first_line_of_block2
  331. and font_flags_of_last_line_of_block1 == font_flags_of_first_line_of_block2
  332. )
  333. def _get_last_paragraph(self, block):
  334. """
  335. Retrieves the last paragraph from a block.
  336. Parameters
  337. ----------
  338. block : dict
  339. The block from which to retrieve the paragraph.
  340. Returns
  341. -------
  342. dict
  343. The last paragraph of the block.
  344. """
  345. if block["paras"]:
  346. last_para_key = list(block["paras"].keys())[-1]
  347. return block["paras"][last_para_key]
  348. else:
  349. return None
  350. def _get_first_paragraph(self, block):
  351. """
  352. Retrieves the first paragraph from a block.
  353. Parameters
  354. ----------
  355. block : dict
  356. The block from which to retrieve the paragraph.
  357. Returns
  358. -------
  359. dict
  360. The first paragraph of the block.
  361. """
  362. if block["paras"]:
  363. first_para_key = list(block["paras"].keys())[0]
  364. return block["paras"][first_para_key]
  365. else:
  366. return None
  367. def should_merge_next_para(self, curr_para, next_para):
  368. if self._is_para_continued(curr_para, next_para):
  369. return True
  370. else:
  371. return False
  372. def batch_tag_paras(self, pdf_dict):
  373. the_last_page_id = len(pdf_dict) - 1
  374. for curr_page_idx, (curr_page_id, curr_page_content) in enumerate(pdf_dict.items()):
  375. if curr_page_id.startswith("page_") and curr_page_content.get("para_blocks", []):
  376. para_blocks_of_curr_page = curr_page_content["para_blocks"]
  377. next_page_idx = curr_page_idx + 1
  378. next_page_id = f"page_{next_page_idx}"
  379. next_page_content = pdf_dict.get(next_page_id, {})
  380. for i, current_block in enumerate(para_blocks_of_curr_page):
  381. for para_id, curr_para in current_block["paras"].items():
  382. curr_para["curr_para_location"] = [
  383. curr_page_idx,
  384. current_block["block_id"],
  385. int(para_id.split("_")[-1]),
  386. ]
  387. curr_para["next_para_location"] = None # 默认设置为None
  388. curr_para["merge_next_para"] = False # 默认设置为False
  389. next_block = para_blocks_of_curr_page[i + 1] if i < len(para_blocks_of_curr_page) - 1 else None
  390. if next_block:
  391. curr_block_last_para_key = list(current_block["paras"].keys())[-1]
  392. curr_blk_last_para = current_block["paras"][curr_block_last_para_key]
  393. next_block_first_para_key = list(next_block["paras"].keys())[0]
  394. next_blk_first_para = next_block["paras"][next_block_first_para_key]
  395. if self.should_merge_next_para(curr_blk_last_para, next_blk_first_para):
  396. curr_blk_last_para["next_para_location"] = [
  397. curr_page_idx,
  398. next_block["block_id"],
  399. int(next_block_first_para_key.split("_")[-1]),
  400. ]
  401. curr_blk_last_para["merge_next_para"] = True
  402. else:
  403. # Handle the case where the next block is in a different page
  404. curr_block_last_para_key = list(current_block["paras"].keys())[-1]
  405. curr_blk_last_para = current_block["paras"][curr_block_last_para_key]
  406. while not next_page_content.get("para_blocks", []) and next_page_idx <= the_last_page_id:
  407. next_page_idx += 1
  408. next_page_id = f"page_{next_page_idx}"
  409. next_page_content = pdf_dict.get(next_page_id, {})
  410. if next_page_content.get("para_blocks", []):
  411. next_blk_first_para_key = list(next_page_content["para_blocks"][0]["paras"].keys())[0]
  412. next_blk_first_para = next_page_content["para_blocks"][0]["paras"][next_blk_first_para_key]
  413. if self.should_merge_next_para(curr_blk_last_para, next_blk_first_para):
  414. curr_blk_last_para["next_para_location"] = [
  415. next_page_idx,
  416. next_page_content["para_blocks"][0]["block_id"],
  417. int(next_blk_first_para_key.split("_")[-1]),
  418. ]
  419. curr_blk_last_para["merge_next_para"] = True
  420. return pdf_dict
  421. def find_block_by_id(self, para_blocks, block_id):
  422. for block in para_blocks:
  423. if block.get("block_id") == block_id:
  424. return block
  425. return None
  426. def batch_merge_paras(self, pdf_dict):
  427. for page_id, page_content in pdf_dict.items():
  428. if page_id.startswith("page_") and page_content.get("para_blocks", []):
  429. para_blocks_of_page = page_content["para_blocks"]
  430. for i in range(len(para_blocks_of_page)):
  431. current_block = para_blocks_of_page[i]
  432. paras = current_block["paras"]
  433. for para_id, curr_para in list(paras.items()):
  434. # 跳过标题段落
  435. if curr_para.get("is_para_title"):
  436. continue
  437. while curr_para.get("merge_next_para"):
  438. next_para_location = curr_para.get("next_para_location")
  439. if not next_para_location:
  440. break
  441. next_page_idx, next_block_id, next_para_id = next_para_location
  442. next_page_id = f"page_{next_page_idx}"
  443. next_page_content = pdf_dict.get(next_page_id)
  444. if not next_page_content:
  445. break
  446. next_block = self.find_block_by_id(next_page_content.get("para_blocks", []), next_block_id)
  447. if not next_block:
  448. break
  449. next_para = next_block["paras"].get(f"para_{next_para_id}")
  450. if not next_para or next_para.get("is_para_title"):
  451. break
  452. # 合并段落文本
  453. curr_para_text = curr_para.get("para_text", "")
  454. next_para_text = next_para.get("para_text", "")
  455. curr_para["para_text"] = curr_para_text + " " + next_para_text
  456. # 更新 next_para_location
  457. curr_para["next_para_location"] = next_para.get("next_para_location")
  458. # 将下一个段落文本置为空,表示已被合并
  459. next_para["para_text"] = ""
  460. # 更新 merge_next_para 标记
  461. curr_para["merge_next_para"] = next_para.get("merge_next_para", False)
  462. return pdf_dict