denoise.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247
  1. import sys
  2. import math
  3. from collections import defaultdict
  4. from para.commons import *
  5. if sys.version_info[0] >= 3:
  6. sys.stdout.reconfigure(encoding="utf-8") # type: ignore
  7. class HeaderFooterProcessor:
  8. def __init__(self) -> None:
  9. pass
  10. def get_most_common_bboxes(self, bboxes, page_height, position="top", threshold=0.25, num_bboxes=3, min_frequency=2):
  11. """
  12. This function gets the most common bboxes from the bboxes
  13. Parameters
  14. ----------
  15. bboxes : list
  16. bboxes
  17. page_height : float
  18. height of the page
  19. position : str, optional
  20. "top" or "bottom", by default "top"
  21. threshold : float, optional
  22. threshold, by default 0.25
  23. num_bboxes : int, optional
  24. number of bboxes to return, by default 3
  25. min_frequency : int, optional
  26. minimum frequency of the bbox, by default 2
  27. Returns
  28. -------
  29. common_bboxes : list
  30. common bboxes
  31. """
  32. # Filter bbox by position
  33. if position == "top":
  34. filtered_bboxes = [bbox for bbox in bboxes if bbox[1] < page_height * threshold]
  35. else:
  36. filtered_bboxes = [bbox for bbox in bboxes if bbox[3] > page_height * (1 - threshold)]
  37. # Find the most common bbox
  38. bbox_count = defaultdict(int)
  39. for bbox in filtered_bboxes:
  40. bbox_count[tuple(bbox)] += 1
  41. # Get the most frequently occurring bbox, but only consider it when the frequency exceeds min_frequency
  42. common_bboxes = [
  43. bbox for bbox, count in sorted(bbox_count.items(), key=lambda item: item[1], reverse=True) if count >= min_frequency
  44. ][:num_bboxes]
  45. return common_bboxes
  46. def detect_footer_header(self, result_dict, similarity_threshold=0.5):
  47. """
  48. This function detects the header and footer of the document.
  49. Parameters
  50. ----------
  51. result_dict : dict
  52. result dictionary
  53. Returns
  54. -------
  55. result_dict : dict
  56. result dictionary
  57. """
  58. def compare_bbox_with_list(bbox, bbox_list, tolerance=1):
  59. return any(all(abs(a - b) < tolerance for a, b in zip(bbox, common_bbox)) for common_bbox in bbox_list)
  60. def is_single_line_block(block):
  61. # Determine based on the width and height of the block
  62. block_width = block["X1"] - block["X0"]
  63. block_height = block["bbox"][3] - block["bbox"][1]
  64. # If the height of the block is close to the average character height and the width is large, it is considered a single line
  65. return block_height <= block["avg_char_height"] * 3 and block_width > block["avg_char_width"] * 3
  66. # Traverse all blocks in the document
  67. single_preproc_blocks = 0
  68. total_blocks = 0
  69. single_preproc_blocks = 0
  70. for page_id, blocks in result_dict.items():
  71. if page_id.startswith("page_"):
  72. for block_key, block in blocks.items():
  73. if block_key.startswith("block_"):
  74. total_blocks += 1
  75. if is_single_line_block(block):
  76. single_preproc_blocks += 1
  77. # If there are no blocks, skip the header and footer detection
  78. if total_blocks == 0:
  79. print("No blocks found. Skipping header/footer detection.")
  80. return result_dict
  81. # If most of the blocks are single-line, skip the header and footer detection
  82. if single_preproc_blocks / total_blocks > 0.5: # 50% of the blocks are single-line
  83. return result_dict
  84. # Collect the bounding boxes of all blocks
  85. all_bboxes = []
  86. all_texts = []
  87. for page_id, blocks in result_dict.items():
  88. if page_id.startswith("page_"):
  89. for block_key, block in blocks.items():
  90. if block_key.startswith("block_"):
  91. all_bboxes.append(block["bbox"])
  92. # Get the height of the page
  93. page_height = max(bbox[3] for bbox in all_bboxes)
  94. # Get the most common bbox lists for headers and footers
  95. common_header_bboxes = self.get_most_common_bboxes(all_bboxes, page_height, position="top") if all_bboxes else []
  96. common_footer_bboxes = self.get_most_common_bboxes(all_bboxes, page_height, position="bottom") if all_bboxes else []
  97. # Detect and mark headers and footers
  98. for page_id, blocks in result_dict.items():
  99. if page_id.startswith("page_"):
  100. for block_key, block in blocks.items():
  101. if block_key.startswith("block_"):
  102. bbox = block["bbox"]
  103. text = block["text"]
  104. is_header = compare_bbox_with_list(bbox, common_header_bboxes)
  105. is_footer = compare_bbox_with_list(bbox, common_footer_bboxes)
  106. block["is_header"] = int(is_header)
  107. block["is_footer"] = int(is_footer)
  108. return result_dict
  109. class NonHorizontalTextProcessor:
  110. def __init__(self) -> None:
  111. pass
  112. def detect_non_horizontal_texts(self, result_dict):
  113. """
  114. This function detects watermarks and vertical margin notes in the document.
  115. Watermarks are identified by finding blocks with the same coordinates and frequently occurring identical texts across multiple pages.
  116. If these conditions are met, the blocks are highly likely to be watermarks, as opposed to headers or footers, which can change from page to page.
  117. If the direction of these blocks is not horizontal, they are definitely considered to be watermarks.
  118. Vertical margin notes are identified by finding blocks with the same coordinates and frequently occurring identical texts across multiple pages.
  119. If these conditions are met, the blocks are highly likely to be vertical margin notes, which typically appear on the left and right sides of the page.
  120. If the direction of these blocks is vertical, they are definitely considered to be vertical margin notes.
  121. Parameters
  122. ----------
  123. result_dict : dict
  124. The result dictionary.
  125. Returns
  126. -------
  127. result_dict : dict
  128. The updated result dictionary.
  129. """
  130. # Dictionary to store information about potential watermarks
  131. potential_watermarks = {}
  132. potential_margin_notes = {}
  133. for page_id, page_content in result_dict.items():
  134. if page_id.startswith("page_"):
  135. for block_id, block_data in page_content.items():
  136. if block_id.startswith("block_"):
  137. if "dir" in block_data:
  138. coordinates_text = (block_data["bbox"], block_data["text"]) # Tuple of coordinates and text
  139. angle = math.atan2(block_data["dir"][1], block_data["dir"][0])
  140. angle = abs(math.degrees(angle))
  141. if angle > 5 and angle < 85: # Check if direction is watermarks
  142. if coordinates_text in potential_watermarks:
  143. potential_watermarks[coordinates_text] += 1
  144. else:
  145. potential_watermarks[coordinates_text] = 1
  146. if angle > 85 and angle < 105: # Check if direction is vertical
  147. if coordinates_text in potential_margin_notes:
  148. potential_margin_notes[coordinates_text] += 1 # Increment count
  149. else:
  150. potential_margin_notes[coordinates_text] = 1 # Initialize count
  151. # Identify watermarks by finding entries with counts higher than a threshold (e.g., appearing on more than half of the pages)
  152. watermark_threshold = len(result_dict) // 2
  153. watermarks = {k: v for k, v in potential_watermarks.items() if v > watermark_threshold}
  154. # Identify margin notes by finding entries with counts higher than a threshold (e.g., appearing on more than half of the pages)
  155. margin_note_threshold = len(result_dict) // 2
  156. margin_notes = {k: v for k, v in potential_margin_notes.items() if v > margin_note_threshold}
  157. # Add watermark information to the result dictionary
  158. for page_id, blocks in result_dict.items():
  159. if page_id.startswith("page_"):
  160. for block_id, block_data in blocks.items():
  161. coordinates_text = (block_data["bbox"], block_data["text"])
  162. if coordinates_text in watermarks:
  163. block_data["is_watermark"] = 1
  164. else:
  165. block_data["is_watermark"] = 0
  166. if coordinates_text in margin_notes:
  167. block_data["is_vertical_margin_note"] = 1
  168. else:
  169. block_data["is_vertical_margin_note"] = 0
  170. return result_dict
  171. class NoiseRemover:
  172. def __init__(self) -> None:
  173. pass
  174. def skip_data_noises(self, result_dict):
  175. """
  176. This function skips the data noises, including overlap blocks, header, footer, watermark, vertical margin note, title
  177. """
  178. filtered_result_dict = {}
  179. for page_id, blocks in result_dict.items():
  180. if page_id.startswith("page_"):
  181. filtered_blocks = {}
  182. for block_id, block in blocks.items():
  183. if block_id.startswith("block_"):
  184. if any(
  185. block.get(key, 0)
  186. for key in [
  187. "is_overlap",
  188. "is_header",
  189. "is_footer",
  190. "is_watermark",
  191. "is_vertical_margin_note",
  192. "is_block_title",
  193. ]
  194. ):
  195. continue
  196. filtered_blocks[block_id] = block
  197. if filtered_blocks:
  198. filtered_result_dict[page_id] = filtered_blocks
  199. return filtered_result_dict