from abc import ABC, abstractmethod from typing import Dict, Any, List, Union, Optional, Tuple import numpy as np from PIL import Image class BaseAdapter(ABC): """基础适配器接口""" def __init__(self, config: Dict[str, Any]): self.config = config @abstractmethod def initialize(self): """初始化模型""" pass @abstractmethod def cleanup(self): """清理资源""" pass class BasePreprocessor(BaseAdapter): """预处理器基类""" @abstractmethod def process(self, image: Union[np.ndarray, Image.Image]) -> tuple[np.ndarray, int]: """ 处理图像 返回处理后的图像和旋转角度 """ pass def _apply_rotation(self, image: np.ndarray, rotation_angle: int) -> np.ndarray: """应用旋转""" import cv2 if rotation_angle == 90: # 90度 return cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE) elif rotation_angle == 180: # 180度 return cv2.rotate(image, cv2.ROTATE_180) elif rotation_angle == 270: # 270度 return cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE) return image class BaseLayoutDetector(BaseAdapter): """版式检测器基类""" def detect( self, image: Union[np.ndarray, Image.Image], ocr_spans: Optional[List[Dict[str, Any]]] = None ) -> List[Dict[str, Any]]: """ 检测版式(模板方法,自动执行后处理) 此方法会: 1. 调用子类实现的 _detect_raw() 进行原始检测 2. 自动执行后处理(去除重叠框、文本转表格等) Args: image: 输入图像 ocr_spans: OCR结果(可选,某些detector可能需要) Returns: 后处理后的布局检测结果 """ # 调用子类实现的原始检测方法 layout_results = self._detect_raw(image, ocr_spans) # 自动执行后处理 if layout_results: layout_config = self.config.get('post_process', {}) if hasattr(self, 'config') else {} layout_results = self.post_process(layout_results, image, layout_config) return layout_results @abstractmethod def _detect_raw( self, image: Union[np.ndarray, Image.Image], ocr_spans: Optional[List[Dict[str, Any]]] = None ) -> List[Dict[str, Any]]: """ 原始检测方法(子类必须实现) Args: image: 输入图像 ocr_spans: OCR结果(可选) Returns: 原始检测结果(未后处理) """ pass def post_process( self, layout_results: List[Dict[str, Any]], image: Union[np.ndarray, Image.Image], config: Optional[Dict[str, Any]] = None ) -> List[Dict[str, Any]]: """ 后处理布局检测结果 默认实现包括: 1. 去除重叠框 2. 将大面积文本块转换为表格(如果配置启用) 子类可以重写此方法以自定义后处理逻辑 Args: layout_results: 原始检测结果 image: 输入图像 config: 后处理配置(可选),如果为None则使用self.config中的post_process配置 Returns: 后处理后的布局结果 """ if not layout_results: return layout_results # 获取配置 if config is None: config = self.config.get('post_process', {}) if hasattr(self, 'config') else {} # 导入 CoordinateUtils(适配器可以访问) try: from ocr_utils.coordinate_utils import CoordinateUtils except ImportError: try: from ocr_utils import CoordinateUtils except ImportError: # 如果无法导入,返回原始结果 return layout_results # 1. 去除重叠框 layout_results = self._remove_overlapping_boxes(layout_results, CoordinateUtils) # 2. 将大面积文本块转换为表格(如果配置启用) layout_config = config if config is not None else {} if layout_config.get('convert_large_text_to_table', False): # 获取图像尺寸 if isinstance(image, Image.Image): h, w = image.size[1], image.size[0] else: h, w = image.shape[:2] if len(image.shape) >= 2 else (0, 0) layout_results = self._convert_large_text_to_table( layout_results, (h, w), min_area_ratio=layout_config.get('min_text_area_ratio', 0.25), min_width_ratio=layout_config.get('min_text_width_ratio', 0.4), min_height_ratio=layout_config.get('min_text_height_ratio', 0.3) ) return layout_results def _remove_overlapping_boxes( self, layout_results: List[Dict[str, Any]], coordinate_utils: Any, iou_threshold: float = 0.8, overlap_ratio_threshold: float = 0.8 ) -> List[Dict[str, Any]]: """ 处理重叠的布局框(参考 MinerU 的去重策略) 策略: 1. 高 IoU 重叠:保留置信度高的框 2. 包含关系:小框被大框高度包含时,保留大框并扩展边界 """ if not layout_results or len(layout_results) <= 1: return layout_results # 复制列表避免修改原数据 results = [item.copy() for item in layout_results] need_remove = set() for i in range(len(results)): if i in need_remove: continue for j in range(i + 1, len(results)): if j in need_remove: continue bbox1 = results[i].get('bbox', [0, 0, 0, 0]) bbox2 = results[j].get('bbox', [0, 0, 0, 0]) if len(bbox1) < 4 or len(bbox2) < 4: continue # 计算 IoU iou = coordinate_utils.calculate_iou(bbox1, bbox2) if iou > iou_threshold: # 高度重叠,保留置信度高的 score1 = results[i].get('confidence', results[i].get('score', 0)) score2 = results[j].get('confidence', results[j].get('score', 0)) if score1 >= score2: need_remove.add(j) else: need_remove.add(i) break # i 被移除,跳出内层循环 else: # 检查包含关系 overlap_ratio = coordinate_utils.calculate_overlap_ratio(bbox1, bbox2) if overlap_ratio > overlap_ratio_threshold: # 小框被大框高度包含 area1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1]) area2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1]) if area1 <= area2: small_idx, large_idx = i, j else: small_idx, large_idx = j, i # 扩展大框的边界 small_bbox = results[small_idx]['bbox'] large_bbox = results[large_idx]['bbox'] results[large_idx]['bbox'] = [ min(small_bbox[0], large_bbox[0]), min(small_bbox[1], large_bbox[1]), max(small_bbox[2], large_bbox[2]), max(small_bbox[3], large_bbox[3]) ] need_remove.add(small_idx) if small_idx == i: break # i 被移除,跳出内层循环 # 返回去重后的结果 return [results[i] for i in range(len(results)) if i not in need_remove] def _convert_large_text_to_table( self, layout_results: List[Dict[str, Any]], image_shape: Tuple[int, int], min_area_ratio: float = 0.25, min_width_ratio: float = 0.4, min_height_ratio: float = 0.3 ) -> List[Dict[str, Any]]: """ 将大面积的文本块转换为表格 判断规则: 1. 面积占比:占页面面积超过 min_area_ratio(默认25%) 2. 尺寸比例:宽度和高度都超过一定比例(避免细长条) 3. 不与其他表格重叠:如果已有表格,不转换 """ if not layout_results: return layout_results img_height, img_width = image_shape img_area = img_height * img_width if img_area == 0: return layout_results # 检查是否已有表格 has_table = any( item.get('category', '').lower() in ['table', 'table_body'] for item in layout_results ) # 如果已有表格,不进行转换(避免误判) if has_table: return layout_results # 复制列表避免修改原数据 results = [item.copy() for item in layout_results] converted_count = 0 for item in results: category = item.get('category', '').lower() # 只处理文本类型的元素 if category not in ['text', 'ocr_text']: continue bbox = item.get('bbox', [0, 0, 0, 0]) if len(bbox) < 4: continue x1, y1, x2, y2 = bbox[:4] width = x2 - x1 height = y2 - y1 area = width * height # 计算占比 area_ratio = area / img_area if img_area > 0 else 0 width_ratio = width / img_width if img_width > 0 else 0 height_ratio = height / img_height if img_height > 0 else 0 # 判断是否满足转换条件 if (area_ratio >= min_area_ratio and width_ratio >= min_width_ratio and height_ratio >= min_height_ratio): # 转换为表格 item['category'] = 'table' item['original_category'] = category # 保留原始类别 converted_count += 1 return results def _map_category_id(self, category_id: int) -> str: """映射类别ID到字符串""" category_map = { 0: 'title', 1: 'text', 2: 'abandon', 3: 'image_body', 4: 'image_caption', 5: 'table_body', 6: 'table_caption', 7: 'table_footnote', 8: 'interline_equation', 9: 'interline_equation_number', 13: 'inline_equation', 14: 'interline_equation_yolo', 15: 'ocr_text', 16: 'low_score_text', 101: 'image_footnote' } return category_map.get(category_id, f'unknown_{category_id}') class BaseVLRecognizer(BaseAdapter): """VL识别器基类""" @abstractmethod def recognize_table(self, image: Union[np.ndarray, Image.Image], **kwargs) -> Dict[str, Any]: """识别表格""" pass @abstractmethod def recognize_formula(self, image: Union[np.ndarray, Image.Image], **kwargs) -> Dict[str, Any]: """识别公式""" pass class BaseOCRRecognizer(BaseAdapter): """OCR识别器基类""" @abstractmethod def recognize_text(self, image: Union[np.ndarray, Image.Image]) -> List[Dict[str, Any]]: """识别文本""" pass @abstractmethod def detect_text_boxes(self, image: Union[np.ndarray, Image.Image]) -> List[Dict[str, Any]]: """ 只检测文本框(不识别文字内容) 子类必须实现此方法。建议使用只运行检测模型的方式(不运行识别模型)以优化性能。 如果无法优化,至少实现一个调用 recognize_text() 的版本作为兜底。 Returns: 文本框列表,每项包含 'bbox', 'poly',可能包含 'confidence' """ pass