from abc import ABC, abstractmethod from typing import Dict, Any, List, Union, Optional, Tuple import numpy as np from PIL import Image from loguru import logger from pathlib import Path import cv2 import json 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 __init__(self, config: Dict[str, Any]): """初始化版式检测器 Args: config: 配置字典 """ super().__init__(config) # 初始化 debug 相关属性(支持从配置或运行时设置) self.debug_mode = None # 将在 detect() 方法中从配置读取 self.output_dir = None # 将在 detect() 方法中从配置读取 self.page_name = None # 将在 detect() 方法中从配置读取 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) # Debug 模式:打印和可视化后处理前的检测结果 # 优先从实例属性读取(如果存在),否则从配置读取 # 支持两种配置方式:debug_mode 或 debug_options.enabled debug_mode = getattr(self, 'debug_mode', None) if debug_mode is None: if hasattr(self, 'config'): # 优先从 debug_mode 读取 debug_mode = self.config.get('debug_mode', False) # 如果没有 debug_mode,尝试从 debug_options.enabled 读取 if not debug_mode: debug_options = self.config.get('debug_options', {}) if isinstance(debug_options, dict): debug_mode = debug_options.get('enabled', False) else: debug_mode = False if debug_mode: logger.debug(f"🔍 Layout detection raw results (before post-processing): {len(layout_results)} elements") # logger.debug(f"Raw layout_results: {layout_results}") # 可视化 layout 结果 output_dir = getattr(self, 'output_dir', None) if output_dir is None: if hasattr(self, 'config'): # 优先从 output_dir 读取 output_dir = self.config.get('output_dir', None) # 如果没有 output_dir,尝试从 debug_options.output_dir 读取 if output_dir is None: debug_options = self.config.get('debug_options', {}) if isinstance(debug_options, dict): output_dir = debug_options.get('output_dir', None) else: output_dir = None page_name = getattr(self, 'page_name', None) if page_name is None: if hasattr(self, 'config'): # 优先从 page_name 读取 page_name = self.config.get('page_name', None) # 如果没有 page_name,尝试从 debug_options.prefix 读取 if page_name is None: debug_options = self.config.get('debug_options', {}) if isinstance(debug_options, dict): prefix = debug_options.get('prefix', '') page_name = prefix if prefix else 'layout_detection' if page_name is None: page_name = 'layout_detection' else: page_name = 'layout_detection' if output_dir: self._visualize_layout_results(image, layout_results, output_dir, page_name, suffix='raw') # 自动执行后处理 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_removed_overlapping = 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_converted_large_text = self._convert_large_text_to_table( layout_results_removed_overlapping, (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_converted_large_text 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}') def _visualize_layout_results( self, image: Union[np.ndarray, Image.Image], layout_results: List[Dict[str, Any]], output_dir: str, page_name: str, suffix: str = 'raw' ) -> None: """ 可视化 layout 检测结果 Args: image: 输入图像 layout_results: 布局检测结果 output_dir: 输出目录 page_name: 页面名称 suffix: 文件名后缀(如 'raw', 'postprocessed') """ if not layout_results: return try: # 转换为 numpy 数组 if isinstance(image, Image.Image): vis_image = np.array(image) if len(vis_image.shape) == 3 and vis_image.shape[2] == 3: # PIL RGB -> OpenCV BGR vis_image = cv2.cvtColor(vis_image, cv2.COLOR_RGB2BGR) else: vis_image = image.copy() if len(vis_image.shape) == 3 and vis_image.shape[2] == 3: # 如果是 RGB,转换为 BGR vis_image = cv2.cvtColor(vis_image, cv2.COLOR_RGB2BGR) # 定义类别颜色映射 (BGR格式) category_colors = { 'table_body': (0, 0, 255), # 红色 'table_caption': (0, 0, 200), # 暗红色 'table_footnote': (0, 0, 150), # 更暗的红色 'text': (255, 0, 0), # 蓝色 'title': (0, 255, 255), # 黄色 'header': (255, 0, 255), # 紫色 'footer': (0, 165, 255), # 橙色 'image_body': (0, 255, 0), # 绿色 'image_caption': (0, 200, 0), # 暗绿色 'image_footnote': (0, 150, 0), # 更暗的绿色 'abandon': (128, 128, 128), # 灰色 } # 绘制检测框 for result in layout_results: bbox = result.get('bbox', []) if not bbox or len(bbox) < 4: continue category = result.get('category', 'unknown') color = category_colors.get(category, (128, 128, 128)) # 默认灰色 thickness = 2 x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]) cv2.rectangle(vis_image, (x1, y1), (x2, y2), color, thickness) # 添加类别标签 label = f"{category}" confidence = result.get('confidence', result.get('score', 0)) if confidence: label += f":{confidence:.2f}" # 计算文本大小 font = cv2.FONT_HERSHEY_SIMPLEX font_scale = 0.4 text_thickness = 1 (text_width, text_height), baseline = cv2.getTextSize(label, font, font_scale, text_thickness) # 在框的上方绘制文本背景 text_y = max(y1 - baseline - 1, text_height + baseline) cv2.rectangle(vis_image, (x1, text_y - text_height - baseline - 2), (x1 + text_width, text_y), color, -1) # 绘制文本 cv2.putText(vis_image, label, (x1, text_y - baseline - 1), font, font_scale, (255, 255, 255), text_thickness) # 保存图像 debug_dir = Path(output_dir) / "debug_comparison" / "layout_detection" debug_dir.mkdir(parents=True, exist_ok=True) output_path = debug_dir / f"{page_name}_layout_{suffix}.jpg" cv2.imwrite(str(output_path), vis_image) logger.info(f"📊 Saved layout detection image ({suffix}): {output_path}") # 保存 JSON 数据 json_data = { 'page_name': page_name, 'suffix': suffix, 'count': len(layout_results), 'results': [ { 'category': r.get('category'), 'bbox': r.get('bbox'), 'confidence': r.get('confidence', r.get('score', 0.0)) } for r in layout_results ] } json_path = debug_dir / f"{page_name}_layout_{suffix}.json" with open(json_path, 'w', encoding='utf-8') as f: json.dump(json_data, f, ensure_ascii=False, indent=2) logger.info(f"📊 Saved layout detection JSON ({suffix}): {json_path}") except Exception as e: logger.warning(f"⚠️ Failed to visualize layout results: {e}") 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]]: """ 改进版重叠框处理算法(基于优先级和决策规则的清晰算法) 策略: 1. 定义类别优先级(abandon < text/image < table_body) 2. 使用统一的决策规则 3. 按综合评分排序处理,优先保留大的聚合框 Args: layout_results: 布局检测结果 coordinate_utils: 坐标工具类 iou_threshold: IoU阈值(默认0.8) overlap_ratio_threshold: 重叠比例阈值(默认0.8) Returns: 去重后的布局结果 """ if not layout_results or len(layout_results) <= 1: return layout_results # 常量定义 CATEGORY_PRIORITY = { 'abandon': 0, 'text': 1, 'image_body': 1, 'title': 2, 'footer': 2, 'header': 2, 'table_body': 3, } AGGREGATE_LABELS = {'key-value region', 'form'} MAX_AREA = 4000000.0 # 用于面积归一化 AREA_WEIGHT = 0.5 CONFIDENCE_WEIGHT = 0.5 AGGREGATE_BONUS = 0.1 AREA_RATIO_THRESHOLD = 3.0 # 大框面积需大于小框的倍数 def get_bbox_area(bbox: List[float]) -> float: """计算bbox面积""" if len(bbox) < 4: return 0.0 return (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]) def is_aggregate_type(box: Dict[str, Any]) -> bool: """检查是否是聚合类型""" original_label = box.get('raw', {}).get('original_label', '').lower() return original_label in AGGREGATE_LABELS def is_bbox_inside(inner: List[float], outer: List[float]) -> bool: """检查inner是否完全包含在outer内""" if len(inner) < 4 or len(outer) < 4: return False return (inner[0] >= outer[0] and inner[1] >= outer[1] and inner[2] <= outer[2] and inner[3] <= outer[3]) def calculate_composite_score(box: Dict[str, Any], area: float) -> float: """计算text类型的综合评分(面积+置信度)""" if box.get('category') != 'text': return box.get('confidence', box.get('score', 0)) normalized_area = min(area / MAX_AREA, 1.0) area_score = (normalized_area ** 0.5) * AREA_WEIGHT confidence_score = box.get('confidence', box.get('score', 0)) * CONFIDENCE_WEIGHT bonus = AGGREGATE_BONUS if is_aggregate_type(box) else 0.0 return area_score + confidence_score + bonus def should_keep_box1(box1: Dict[str, Any], box2: Dict[str, Any], iou: float, overlap_ratio: float, contained_1_in_2: bool, contained_2_in_1: bool) -> bool: """判断是否应该保留box1""" # 提取基本信息 cat1, cat2 = box1.get('category', ''), box2.get('category', '') score1 = box1.get('confidence', box1.get('score', 0)) score2 = box2.get('confidence', box2.get('score', 0)) bbox1, bbox2 = box1.get('bbox', [0, 0, 0, 0]), box2.get('bbox', [0, 0, 0, 0]) area1, area2 = get_bbox_area(bbox1), get_bbox_area(bbox2) is_agg1, is_agg2 = is_aggregate_type(box1), is_aggregate_type(box2) # 规则1: 类别优先级 priority1 = CATEGORY_PRIORITY.get(cat1, 1) priority2 = CATEGORY_PRIORITY.get(cat2, 1) if priority1 != priority2: return priority1 > priority2 # 规则2: 包含关系 + 聚合类型优先 if contained_2_in_1 and is_agg1 and not is_agg2: return True if contained_1_in_2 and is_agg2 and not is_agg1: return False # 规则3: 包含关系 + 面积比例 if contained_2_in_1 and area1 > area2 * AREA_RATIO_THRESHOLD: return True if contained_1_in_2 and area2 > area1 * AREA_RATIO_THRESHOLD: return False # 规则4: text类型使用综合评分 if cat1 == 'text' or cat2 == 'text': comp_score1 = calculate_composite_score(box1, area1) comp_score2 = calculate_composite_score(box2, area2) if abs(comp_score1 - comp_score2) > 0.05: return comp_score1 > comp_score2 # 规则5: 置信度比较 if abs(score1 - score2) > 0.1: return score1 > score2 # 规则6: 面积比较 return area1 >= area2 # 主处理逻辑 results = [item.copy() for item in layout_results] need_remove = set() # 按综合评分排序(高分优先) def get_sort_key(i: int) -> float: item = results[i] if item.get('category') == 'text': return -calculate_composite_score(item, get_bbox_area(item.get('bbox', []))) return -item.get('confidence', item.get('score', 0)) sorted_indices = sorted(range(len(results)), key=get_sort_key) # 比较每对框 for idx_i, i in enumerate(sorted_indices): if i in need_remove: continue for idx_j, j in enumerate(sorted_indices): if j == i or j in need_remove or idx_j >= idx_i: continue bbox1, bbox2 = results[i].get('bbox', []), results[j].get('bbox', []) if len(bbox1) < 4 or len(bbox2) < 4: continue # 计算重叠指标 iou = coordinate_utils.calculate_iou(bbox1, bbox2) overlap_ratio = coordinate_utils.calculate_overlap_ratio(bbox1, bbox2) contained_1_in_2 = is_bbox_inside(bbox1, bbox2) contained_2_in_1 = is_bbox_inside(bbox2, bbox1) # 检查是否有显著重叠 if not (iou > iou_threshold or overlap_ratio > overlap_ratio_threshold or contained_1_in_2 or contained_2_in_1): continue # 应用决策规则 if should_keep_box1(results[i], results[j], iou, overlap_ratio, contained_1_in_2, contained_2_in_1): need_remove.add(j) else: need_remove.add(i) break return [results[i] for i in range(len(results)) if i not in need_remove] 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