|
|
@@ -1,5 +1,5 @@
|
|
|
from abc import ABC, abstractmethod
|
|
|
-from typing import Dict, Any, List, Union
|
|
|
+from typing import Dict, Any, List, Union, Optional, Tuple
|
|
|
import numpy as np
|
|
|
from PIL import Image
|
|
|
|
|
|
@@ -44,11 +44,266 @@ class BasePreprocessor(BaseAdapter):
|
|
|
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(self, image: Union[np.ndarray, Image.Image]) -> List[Dict[str, Any]]:
|
|
|
- """检测版式"""
|
|
|
+ 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 = {
|
|
|
@@ -89,4 +344,17 @@ class BaseOCRRecognizer(BaseAdapter):
|
|
|
@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
|