base.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360
  1. from abc import ABC, abstractmethod
  2. from typing import Dict, Any, List, Union, Optional, Tuple
  3. import numpy as np
  4. from PIL import Image
  5. class BaseAdapter(ABC):
  6. """基础适配器接口"""
  7. def __init__(self, config: Dict[str, Any]):
  8. self.config = config
  9. @abstractmethod
  10. def initialize(self):
  11. """初始化模型"""
  12. pass
  13. @abstractmethod
  14. def cleanup(self):
  15. """清理资源"""
  16. pass
  17. class BasePreprocessor(BaseAdapter):
  18. """预处理器基类"""
  19. @abstractmethod
  20. def process(self, image: Union[np.ndarray, Image.Image]) -> tuple[np.ndarray, int]:
  21. """
  22. 处理图像
  23. 返回处理后的图像和旋转角度
  24. """
  25. pass
  26. def _apply_rotation(self, image: np.ndarray, rotation_angle: int) -> np.ndarray:
  27. """应用旋转"""
  28. import cv2
  29. if rotation_angle == 90: # 90度
  30. return cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)
  31. elif rotation_angle == 180: # 180度
  32. return cv2.rotate(image, cv2.ROTATE_180)
  33. elif rotation_angle == 270: # 270度
  34. return cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
  35. return image
  36. class BaseLayoutDetector(BaseAdapter):
  37. """版式检测器基类"""
  38. def detect(
  39. self,
  40. image: Union[np.ndarray, Image.Image],
  41. ocr_spans: Optional[List[Dict[str, Any]]] = None
  42. ) -> List[Dict[str, Any]]:
  43. """
  44. 检测版式(模板方法,自动执行后处理)
  45. 此方法会:
  46. 1. 调用子类实现的 _detect_raw() 进行原始检测
  47. 2. 自动执行后处理(去除重叠框、文本转表格等)
  48. Args:
  49. image: 输入图像
  50. ocr_spans: OCR结果(可选,某些detector可能需要)
  51. Returns:
  52. 后处理后的布局检测结果
  53. """
  54. # 调用子类实现的原始检测方法
  55. layout_results = self._detect_raw(image, ocr_spans)
  56. # 自动执行后处理
  57. if layout_results:
  58. layout_config = self.config.get('post_process', {}) if hasattr(self, 'config') else {}
  59. layout_results = self.post_process(layout_results, image, layout_config)
  60. return layout_results
  61. @abstractmethod
  62. def _detect_raw(
  63. self,
  64. image: Union[np.ndarray, Image.Image],
  65. ocr_spans: Optional[List[Dict[str, Any]]] = None
  66. ) -> List[Dict[str, Any]]:
  67. """
  68. 原始检测方法(子类必须实现)
  69. Args:
  70. image: 输入图像
  71. ocr_spans: OCR结果(可选)
  72. Returns:
  73. 原始检测结果(未后处理)
  74. """
  75. pass
  76. def post_process(
  77. self,
  78. layout_results: List[Dict[str, Any]],
  79. image: Union[np.ndarray, Image.Image],
  80. config: Optional[Dict[str, Any]] = None
  81. ) -> List[Dict[str, Any]]:
  82. """
  83. 后处理布局检测结果
  84. 默认实现包括:
  85. 1. 去除重叠框
  86. 2. 将大面积文本块转换为表格(如果配置启用)
  87. 子类可以重写此方法以自定义后处理逻辑
  88. Args:
  89. layout_results: 原始检测结果
  90. image: 输入图像
  91. config: 后处理配置(可选),如果为None则使用self.config中的post_process配置
  92. Returns:
  93. 后处理后的布局结果
  94. """
  95. if not layout_results:
  96. return layout_results
  97. # 获取配置
  98. if config is None:
  99. config = self.config.get('post_process', {}) if hasattr(self, 'config') else {}
  100. # 导入 CoordinateUtils(适配器可以访问)
  101. try:
  102. from ocr_utils.coordinate_utils import CoordinateUtils
  103. except ImportError:
  104. try:
  105. from ocr_utils import CoordinateUtils
  106. except ImportError:
  107. # 如果无法导入,返回原始结果
  108. return layout_results
  109. # 1. 去除重叠框
  110. layout_results = self._remove_overlapping_boxes(layout_results, CoordinateUtils)
  111. # 2. 将大面积文本块转换为表格(如果配置启用)
  112. layout_config = config if config is not None else {}
  113. if layout_config.get('convert_large_text_to_table', False):
  114. # 获取图像尺寸
  115. if isinstance(image, Image.Image):
  116. h, w = image.size[1], image.size[0]
  117. else:
  118. h, w = image.shape[:2] if len(image.shape) >= 2 else (0, 0)
  119. layout_results = self._convert_large_text_to_table(
  120. layout_results,
  121. (h, w),
  122. min_area_ratio=layout_config.get('min_text_area_ratio', 0.25),
  123. min_width_ratio=layout_config.get('min_text_width_ratio', 0.4),
  124. min_height_ratio=layout_config.get('min_text_height_ratio', 0.3)
  125. )
  126. return layout_results
  127. def _remove_overlapping_boxes(
  128. self,
  129. layout_results: List[Dict[str, Any]],
  130. coordinate_utils: Any,
  131. iou_threshold: float = 0.8,
  132. overlap_ratio_threshold: float = 0.8
  133. ) -> List[Dict[str, Any]]:
  134. """
  135. 处理重叠的布局框(参考 MinerU 的去重策略)
  136. 策略:
  137. 1. 高 IoU 重叠:保留置信度高的框
  138. 2. 包含关系:小框被大框高度包含时,保留大框并扩展边界
  139. """
  140. if not layout_results or len(layout_results) <= 1:
  141. return layout_results
  142. # 复制列表避免修改原数据
  143. results = [item.copy() for item in layout_results]
  144. need_remove = set()
  145. for i in range(len(results)):
  146. if i in need_remove:
  147. continue
  148. for j in range(i + 1, len(results)):
  149. if j in need_remove:
  150. continue
  151. bbox1 = results[i].get('bbox', [0, 0, 0, 0])
  152. bbox2 = results[j].get('bbox', [0, 0, 0, 0])
  153. if len(bbox1) < 4 or len(bbox2) < 4:
  154. continue
  155. # 计算 IoU
  156. iou = coordinate_utils.calculate_iou(bbox1, bbox2)
  157. if iou > iou_threshold:
  158. # 高度重叠,保留置信度高的
  159. score1 = results[i].get('confidence', results[i].get('score', 0))
  160. score2 = results[j].get('confidence', results[j].get('score', 0))
  161. if score1 >= score2:
  162. need_remove.add(j)
  163. else:
  164. need_remove.add(i)
  165. break # i 被移除,跳出内层循环
  166. else:
  167. # 检查包含关系
  168. overlap_ratio = coordinate_utils.calculate_overlap_ratio(bbox1, bbox2)
  169. if overlap_ratio > overlap_ratio_threshold:
  170. # 小框被大框高度包含
  171. area1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
  172. area2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])
  173. if area1 <= area2:
  174. small_idx, large_idx = i, j
  175. else:
  176. small_idx, large_idx = j, i
  177. # 扩展大框的边界
  178. small_bbox = results[small_idx]['bbox']
  179. large_bbox = results[large_idx]['bbox']
  180. results[large_idx]['bbox'] = [
  181. min(small_bbox[0], large_bbox[0]),
  182. min(small_bbox[1], large_bbox[1]),
  183. max(small_bbox[2], large_bbox[2]),
  184. max(small_bbox[3], large_bbox[3])
  185. ]
  186. need_remove.add(small_idx)
  187. if small_idx == i:
  188. break # i 被移除,跳出内层循环
  189. # 返回去重后的结果
  190. return [results[i] for i in range(len(results)) if i not in need_remove]
  191. def _convert_large_text_to_table(
  192. self,
  193. layout_results: List[Dict[str, Any]],
  194. image_shape: Tuple[int, int],
  195. min_area_ratio: float = 0.25,
  196. min_width_ratio: float = 0.4,
  197. min_height_ratio: float = 0.3
  198. ) -> List[Dict[str, Any]]:
  199. """
  200. 将大面积的文本块转换为表格
  201. 判断规则:
  202. 1. 面积占比:占页面面积超过 min_area_ratio(默认25%)
  203. 2. 尺寸比例:宽度和高度都超过一定比例(避免细长条)
  204. 3. 不与其他表格重叠:如果已有表格,不转换
  205. """
  206. if not layout_results:
  207. return layout_results
  208. img_height, img_width = image_shape
  209. img_area = img_height * img_width
  210. if img_area == 0:
  211. return layout_results
  212. # 检查是否已有表格
  213. has_table = any(
  214. item.get('category', '').lower() in ['table', 'table_body']
  215. for item in layout_results
  216. )
  217. # 如果已有表格,不进行转换(避免误判)
  218. if has_table:
  219. return layout_results
  220. # 复制列表避免修改原数据
  221. results = [item.copy() for item in layout_results]
  222. converted_count = 0
  223. for item in results:
  224. category = item.get('category', '').lower()
  225. # 只处理文本类型的元素
  226. if category not in ['text', 'ocr_text']:
  227. continue
  228. bbox = item.get('bbox', [0, 0, 0, 0])
  229. if len(bbox) < 4:
  230. continue
  231. x1, y1, x2, y2 = bbox[:4]
  232. width = x2 - x1
  233. height = y2 - y1
  234. area = width * height
  235. # 计算占比
  236. area_ratio = area / img_area if img_area > 0 else 0
  237. width_ratio = width / img_width if img_width > 0 else 0
  238. height_ratio = height / img_height if img_height > 0 else 0
  239. # 判断是否满足转换条件
  240. if (area_ratio >= min_area_ratio and
  241. width_ratio >= min_width_ratio and
  242. height_ratio >= min_height_ratio):
  243. # 转换为表格
  244. item['category'] = 'table'
  245. item['original_category'] = category # 保留原始类别
  246. converted_count += 1
  247. return results
  248. def _map_category_id(self, category_id: int) -> str:
  249. """映射类别ID到字符串"""
  250. category_map = {
  251. 0: 'title',
  252. 1: 'text',
  253. 2: 'abandon',
  254. 3: 'image_body',
  255. 4: 'image_caption',
  256. 5: 'table_body',
  257. 6: 'table_caption',
  258. 7: 'table_footnote',
  259. 8: 'interline_equation',
  260. 9: 'interline_equation_number',
  261. 13: 'inline_equation',
  262. 14: 'interline_equation_yolo',
  263. 15: 'ocr_text',
  264. 16: 'low_score_text',
  265. 101: 'image_footnote'
  266. }
  267. return category_map.get(category_id, f'unknown_{category_id}')
  268. class BaseVLRecognizer(BaseAdapter):
  269. """VL识别器基类"""
  270. @abstractmethod
  271. def recognize_table(self, image: Union[np.ndarray, Image.Image], **kwargs) -> Dict[str, Any]:
  272. """识别表格"""
  273. pass
  274. @abstractmethod
  275. def recognize_formula(self, image: Union[np.ndarray, Image.Image], **kwargs) -> Dict[str, Any]:
  276. """识别公式"""
  277. pass
  278. class BaseOCRRecognizer(BaseAdapter):
  279. """OCR识别器基类"""
  280. @abstractmethod
  281. def recognize_text(self, image: Union[np.ndarray, Image.Image]) -> List[Dict[str, Any]]:
  282. """识别文本"""
  283. pass
  284. @abstractmethod
  285. def detect_text_boxes(self, image: Union[np.ndarray, Image.Image]) -> List[Dict[str, Any]]:
  286. """
  287. 只检测文本框(不识别文字内容)
  288. 子类必须实现此方法。建议使用只运行检测模型的方式(不运行识别模型)以优化性能。
  289. 如果无法优化,至少实现一个调用 recognize_text() 的版本作为兜底。
  290. Returns:
  291. 文本框列表,每项包含 'bbox', 'poly',可能包含 'confidence'
  292. """
  293. pass