block_termination_processor.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486
  1. import sys
  2. from libs.commons import fitz
  3. from termcolor import cprint
  4. from para.commons import *
  5. if sys.version_info[0] >= 3:
  6. sys.stdout.reconfigure(encoding="utf-8") # type: ignore
  7. class BlockTerminationProcessor:
  8. def __init__(self) -> None:
  9. pass
  10. def _is_consistent_lines(
  11. self,
  12. curr_line,
  13. prev_line,
  14. next_line,
  15. consistent_direction, # 0 for prev, 1 for next, 2 for both
  16. ):
  17. """
  18. This function checks if the line is consistent with its neighbors
  19. Parameters
  20. ----------
  21. curr_line : dict
  22. current line
  23. prev_line : dict
  24. previous line
  25. next_line : dict
  26. next line
  27. consistent_direction : int
  28. 0 for prev, 1 for next, 2 for both
  29. Returns
  30. -------
  31. bool
  32. True if the line is consistent with its neighbors, False otherwise.
  33. """
  34. curr_line_font_size = curr_line["spans"][0]["size"]
  35. curr_line_font_type = curr_line["spans"][0]["font"].lower()
  36. if consistent_direction == 0:
  37. if prev_line:
  38. prev_line_font_size = prev_line["spans"][0]["size"]
  39. prev_line_font_type = prev_line["spans"][0]["font"].lower()
  40. return curr_line_font_size == prev_line_font_size and curr_line_font_type == prev_line_font_type
  41. else:
  42. return False
  43. elif consistent_direction == 1:
  44. if next_line:
  45. next_line_font_size = next_line["spans"][0]["size"]
  46. next_line_font_type = next_line["spans"][0]["font"].lower()
  47. return curr_line_font_size == next_line_font_size and curr_line_font_type == next_line_font_type
  48. else:
  49. return False
  50. elif consistent_direction == 2:
  51. if prev_line and next_line:
  52. prev_line_font_size = prev_line["spans"][0]["size"]
  53. prev_line_font_type = prev_line["spans"][0]["font"].lower()
  54. next_line_font_size = next_line["spans"][0]["size"]
  55. next_line_font_type = next_line["spans"][0]["font"].lower()
  56. return (curr_line_font_size == prev_line_font_size and curr_line_font_type == prev_line_font_type) and (
  57. curr_line_font_size == next_line_font_size and curr_line_font_type == next_line_font_type
  58. )
  59. else:
  60. return False
  61. else:
  62. return False
  63. def _is_regular_line(self, curr_line_bbox, prev_line_bbox, next_line_bbox, avg_char_width, X0, X1, avg_line_height):
  64. """
  65. This function checks if the line is a regular line
  66. Parameters
  67. ----------
  68. curr_line_bbox : list
  69. bbox of the current line
  70. prev_line_bbox : list
  71. bbox of the previous line
  72. next_line_bbox : list
  73. bbox of the next line
  74. avg_char_width : float
  75. average of char widths
  76. X0 : float
  77. median of x0 values, which represents the left average boundary of the page
  78. X1 : float
  79. median of x1 values, which represents the right average boundary of the page
  80. avg_line_height : float
  81. average of line heights
  82. Returns
  83. -------
  84. bool
  85. True if the line is a regular line, False otherwise.
  86. """
  87. horizontal_ratio = 0.5
  88. vertical_ratio = 0.5
  89. horizontal_thres = horizontal_ratio * avg_char_width
  90. vertical_thres = vertical_ratio * avg_line_height
  91. x0, y0, x1, y1 = curr_line_bbox
  92. x0_near_X0 = abs(x0 - X0) < horizontal_thres
  93. x1_near_X1 = abs(x1 - X1) < horizontal_thres
  94. prev_line_is_end_of_para = prev_line_bbox and (abs(prev_line_bbox[2] - X1) > avg_char_width)
  95. sufficient_spacing_above = False
  96. if prev_line_bbox:
  97. vertical_spacing_above = y1 - prev_line_bbox[3]
  98. sufficient_spacing_above = vertical_spacing_above > vertical_thres
  99. sufficient_spacing_below = False
  100. if next_line_bbox:
  101. vertical_spacing_below = next_line_bbox[1] - y0
  102. sufficient_spacing_below = vertical_spacing_below > vertical_thres
  103. return (
  104. (sufficient_spacing_above or sufficient_spacing_below)
  105. or (not x0_near_X0 and not x1_near_X1)
  106. or prev_line_is_end_of_para
  107. )
  108. def _is_possible_start_of_para(self, curr_line, prev_line, next_line, X0, X1, avg_char_width, avg_font_size):
  109. """
  110. This function checks if the line is a possible start of a paragraph
  111. Parameters
  112. ----------
  113. curr_line : dict
  114. current line
  115. prev_line : dict
  116. previous line
  117. next_line : dict
  118. next line
  119. X0 : float
  120. median of x0 values, which represents the left average boundary of the page
  121. X1 : float
  122. median of x1 values, which represents the right average boundary of the page
  123. avg_char_width : float
  124. average of char widths
  125. avg_line_height : float
  126. average of line heights
  127. Returns
  128. -------
  129. bool
  130. True if the line is a possible start of a paragraph, False otherwise.
  131. """
  132. start_confidence = 0.5 # Initial confidence of the line being a start of a paragraph
  133. decision_path = [] # Record the decision path
  134. curr_line_bbox = curr_line["bbox"]
  135. prev_line_bbox = prev_line["bbox"] if prev_line else None
  136. next_line_bbox = next_line["bbox"] if next_line else None
  137. indent_ratio = 1
  138. vertical_ratio = 1.5
  139. vertical_thres = vertical_ratio * avg_font_size
  140. left_horizontal_ratio = 0.5
  141. left_horizontal_thres = left_horizontal_ratio * avg_char_width
  142. right_horizontal_ratio = 2.5
  143. right_horizontal_thres = right_horizontal_ratio * avg_char_width
  144. x0, y0, x1, y1 = curr_line_bbox
  145. indent_condition = x0 > X0 + indent_ratio * avg_char_width
  146. if indent_condition:
  147. start_confidence += 0.2
  148. decision_path.append("indent_condition_met")
  149. x0_near_X0 = abs(x0 - X0) < left_horizontal_thres
  150. if x0_near_X0:
  151. start_confidence += 0.1
  152. decision_path.append("x0_near_X0")
  153. x1_near_X1 = abs(x1 - X1) < right_horizontal_thres
  154. if x1_near_X1:
  155. start_confidence += 0.1
  156. decision_path.append("x1_near_X1")
  157. if prev_line is None:
  158. prev_line_is_end_of_para = True
  159. start_confidence += 0.2
  160. decision_path.append("no_prev_line")
  161. else:
  162. prev_line_is_end_of_para, _, _ = self._is_possible_end_of_para(prev_line, next_line, X0, X1, avg_char_width)
  163. if prev_line_is_end_of_para:
  164. start_confidence += 0.1
  165. decision_path.append("prev_line_is_end_of_para")
  166. sufficient_spacing_above = False
  167. if prev_line_bbox:
  168. vertical_spacing_above = y1 - prev_line_bbox[3]
  169. sufficient_spacing_above = vertical_spacing_above > vertical_thres
  170. if sufficient_spacing_above:
  171. start_confidence += 0.2
  172. decision_path.append("sufficient_spacing_above")
  173. sufficient_spacing_below = False
  174. if next_line_bbox:
  175. vertical_spacing_below = next_line_bbox[1] - y0
  176. sufficient_spacing_below = vertical_spacing_below > vertical_thres
  177. if sufficient_spacing_below:
  178. start_confidence += 0.2
  179. decision_path.append("sufficient_spacing_below")
  180. is_regular_line = self._is_regular_line(
  181. curr_line_bbox, prev_line_bbox, next_line_bbox, avg_char_width, X0, X1, avg_font_size
  182. )
  183. if is_regular_line:
  184. start_confidence += 0.1
  185. decision_path.append("is_regular_line")
  186. is_start_of_para = (
  187. (sufficient_spacing_above or sufficient_spacing_below)
  188. or (indent_condition)
  189. or (not indent_condition and x0_near_X0 and x1_near_X1 and not is_regular_line)
  190. or prev_line_is_end_of_para
  191. )
  192. return (is_start_of_para, start_confidence, decision_path)
  193. def _is_possible_end_of_para(self, curr_line, next_line, X0, X1, avg_char_width):
  194. """
  195. This function checks if the line is a possible end of a paragraph
  196. Parameters
  197. ----------
  198. curr_line : dict
  199. current line
  200. next_line : dict
  201. next line
  202. X0 : float
  203. median of x0 values, which represents the left average boundary of the page
  204. X1 : float
  205. median of x1 values, which represents the right average boundary of the page
  206. avg_char_width : float
  207. average of char widths
  208. Returns
  209. -------
  210. bool
  211. True if the line is a possible end of a paragraph, False otherwise.
  212. """
  213. end_confidence = 0.5 # Initial confidence of the line being a end of a paragraph
  214. decision_path = [] # Record the decision path
  215. curr_line_bbox = curr_line["bbox"]
  216. next_line_bbox = next_line["bbox"] if next_line else None
  217. left_horizontal_ratio = 0.5
  218. right_horizontal_ratio = 0.5
  219. x0, _, x1, y1 = curr_line_bbox
  220. next_x0, next_y0, _, _ = next_line_bbox if next_line_bbox else (0, 0, 0, 0)
  221. x0_near_X0 = abs(x0 - X0) < left_horizontal_ratio * avg_char_width
  222. if x0_near_X0:
  223. end_confidence += 0.1
  224. decision_path.append("x0_near_X0")
  225. x1_smaller_than_X1 = x1 < X1 - right_horizontal_ratio * avg_char_width
  226. if x1_smaller_than_X1:
  227. end_confidence += 0.1
  228. decision_path.append("x1_smaller_than_X1")
  229. next_line_is_start_of_para = (
  230. next_line_bbox
  231. and (next_x0 > X0 + left_horizontal_ratio * avg_char_width)
  232. and (not is_line_left_aligned_from_neighbors(curr_line_bbox, None, next_line_bbox, avg_char_width, direction=1))
  233. )
  234. if next_line_is_start_of_para:
  235. end_confidence += 0.2
  236. decision_path.append("next_line_is_start_of_para")
  237. is_line_left_aligned_from_neighbors_bool = is_line_left_aligned_from_neighbors(
  238. curr_line_bbox, None, next_line_bbox, avg_char_width
  239. )
  240. if is_line_left_aligned_from_neighbors_bool:
  241. end_confidence += 0.1
  242. decision_path.append("line_is_left_aligned_from_neighbors")
  243. is_line_right_aligned_from_neighbors_bool = is_line_right_aligned_from_neighbors(
  244. curr_line_bbox, None, next_line_bbox, avg_char_width
  245. )
  246. if not is_line_right_aligned_from_neighbors_bool:
  247. end_confidence += 0.1
  248. decision_path.append("line_is_not_right_aligned_from_neighbors")
  249. is_end_of_para = end_with_punctuation(curr_line["text"]) and (
  250. (x0_near_X0 and x1_smaller_than_X1)
  251. or (is_line_left_aligned_from_neighbors_bool and not is_line_right_aligned_from_neighbors_bool)
  252. )
  253. return (is_end_of_para, end_confidence, decision_path)
  254. def _cut_paras_per_block(
  255. self,
  256. block,
  257. ):
  258. """
  259. Processes a raw block from PyMuPDF and returns the processed block.
  260. Parameters
  261. ----------
  262. raw_block : dict
  263. A raw block from pymupdf.
  264. Returns
  265. -------
  266. processed_block : dict
  267. """
  268. def _construct_para(lines, is_block_title, para_title_level):
  269. """
  270. Construct a paragraph from given lines.
  271. """
  272. font_sizes = [span["size"] for line in lines for span in line["spans"]]
  273. avg_font_size = sum(font_sizes) / len(font_sizes) if font_sizes else 0
  274. font_colors = [span["color"] for line in lines for span in line["spans"]]
  275. most_common_font_color = max(set(font_colors), key=font_colors.count) if font_colors else None
  276. # font_types = [span["font"] for line in lines for span in line["spans"]]
  277. # most_common_font_type = max(set(font_types), key=font_types.count) if font_types else None
  278. font_type_lengths = {}
  279. for line in lines:
  280. for span in line["spans"]:
  281. font_type = span["font"]
  282. bbox_width = span["bbox"][2] - span["bbox"][0]
  283. if font_type in font_type_lengths:
  284. font_type_lengths[font_type] += bbox_width
  285. else:
  286. font_type_lengths[font_type] = bbox_width
  287. # get the font type with the longest bbox width
  288. most_common_font_type = max(font_type_lengths, key=font_type_lengths.get) if font_type_lengths else None # type: ignore
  289. para_bbox = calculate_para_bbox(lines)
  290. para_text = " ".join(line["text"] for line in lines)
  291. return {
  292. "para_bbox": para_bbox,
  293. "para_text": para_text,
  294. "para_font_type": most_common_font_type,
  295. "para_font_size": avg_font_size,
  296. "para_font_color": most_common_font_color,
  297. "is_para_title": is_block_title,
  298. "para_title_level": para_title_level,
  299. }
  300. block_bbox = block["bbox"]
  301. block_text = block["text"]
  302. block_lines = block["lines"]
  303. X0 = safe_get(block, "X0", 0)
  304. X1 = safe_get(block, "X1", 0)
  305. avg_char_width = safe_get(block, "avg_char_width", 0)
  306. avg_char_height = safe_get(block, "avg_char_height", 0)
  307. avg_font_size = safe_get(block, "avg_font_size", 0)
  308. is_block_title = safe_get(block, "is_block_title", False)
  309. para_title_level = safe_get(block, "block_title_level", 0)
  310. # Segment into paragraphs
  311. para_ranges = []
  312. in_paragraph = False
  313. start_idx_of_para = None
  314. # Create the processed paragraphs
  315. processed_paras = {}
  316. para_bboxes = []
  317. end_idx_of_para = 0
  318. for line_index, line in enumerate(block_lines):
  319. curr_line = line
  320. prev_line = block_lines[line_index - 1] if line_index > 0 else None
  321. next_line = block_lines[line_index + 1] if line_index < len(block_lines) - 1 else None
  322. """
  323. Start processing paragraphs.
  324. """
  325. # Check if the line is the start of a paragraph
  326. is_start_of_para, start_confidence, decision_path = self._is_possible_start_of_para(
  327. curr_line, prev_line, next_line, X0, X1, avg_char_width, avg_font_size
  328. )
  329. if not in_paragraph and is_start_of_para:
  330. in_paragraph = True
  331. start_idx_of_para = line_index
  332. # print_green(">>> Start of a paragraph")
  333. # print(" curr_line_text: ", curr_line["text"])
  334. # print(" start_confidence: ", start_confidence)
  335. # print(" decision_path: ", decision_path)
  336. # Check if the line is the end of a paragraph
  337. is_end_of_para, end_confidence, decision_path = self._is_possible_end_of_para(
  338. curr_line, next_line, X0, X1, avg_char_width
  339. )
  340. if in_paragraph and (is_end_of_para or not next_line):
  341. para_ranges.append((start_idx_of_para, line_index))
  342. start_idx_of_para = None
  343. in_paragraph = False
  344. # print_red(">>> End of a paragraph")
  345. # print(" curr_line_text: ", curr_line["text"])
  346. # print(" end_confidence: ", end_confidence)
  347. # print(" decision_path: ", decision_path)
  348. # Add the last paragraph if it is not added
  349. if in_paragraph and start_idx_of_para is not None:
  350. para_ranges.append((start_idx_of_para, len(block_lines) - 1))
  351. # Process the matched paragraphs
  352. for para_index, (start_idx, end_idx) in enumerate(para_ranges):
  353. matched_lines = block_lines[start_idx : end_idx + 1]
  354. para_properties = _construct_para(matched_lines, is_block_title, para_title_level)
  355. para_key = f"para_{len(processed_paras)}"
  356. processed_paras[para_key] = para_properties
  357. para_bboxes.append(para_properties["para_bbox"])
  358. end_idx_of_para = end_idx + 1
  359. # Deal with the remaining lines
  360. if end_idx_of_para < len(block_lines):
  361. unmatched_lines = block_lines[end_idx_of_para:]
  362. unmatched_properties = _construct_para(unmatched_lines, is_block_title, para_title_level)
  363. unmatched_key = f"para_{len(processed_paras)}"
  364. processed_paras[unmatched_key] = unmatched_properties
  365. para_bboxes.append(unmatched_properties["para_bbox"])
  366. block["paras"] = processed_paras
  367. return block
  368. def batch_process_blocks(self, pdf_dict):
  369. """
  370. Parses the blocks of all pages.
  371. Parameters
  372. ----------
  373. pdf_dict : dict
  374. PDF dictionary.
  375. filter_blocks : list
  376. List of bounding boxes to filter.
  377. Returns
  378. -------
  379. result_dict : dict
  380. Result dictionary.
  381. """
  382. num_paras = 0
  383. for page_id, page in pdf_dict.items():
  384. if page_id.startswith("page_"):
  385. para_blocks = []
  386. if "para_blocks" in page.keys():
  387. input_blocks = page["para_blocks"]
  388. for input_block in input_blocks:
  389. new_block = self._cut_paras_per_block(input_block)
  390. para_blocks.append(new_block)
  391. num_paras += len(new_block["paras"])
  392. page["para_blocks"] = para_blocks
  393. pdf_dict["statistics"]["num_paras"] = num_paras
  394. return pdf_dict