| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360 |
- 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
|