|
|
@@ -0,0 +1,210 @@
|
|
|
+# zhch/custom_modules/enhanced_doc_orientation.py
|
|
|
+
|
|
|
+"""增强版文档方向分类器 - 结合 OCR 和 CNN"""
|
|
|
+
|
|
|
+import cv2
|
|
|
+import numpy as np
|
|
|
+import onnxruntime
|
|
|
+from paddlex import create_model
|
|
|
+from typing import Union, List
|
|
|
+from PIL import Image
|
|
|
+
|
|
|
+
|
|
|
+class EnhancedDocOrientationClassify:
|
|
|
+ """
|
|
|
+ 增强版文档方向分类器
|
|
|
+
|
|
|
+ 参考 MinerU 的多阶段判断逻辑:
|
|
|
+ 1. 快速过滤: 宽高比检查
|
|
|
+ 2. OCR 分析: 文本框方向判断
|
|
|
+ 3. CNN 分类: 精确角度预测
|
|
|
+ """
|
|
|
+
|
|
|
+ def __init__(
|
|
|
+ self,
|
|
|
+ cnn_model_name: str = "PP-LCNet_x1_0_doc_ori",
|
|
|
+ ocr_model_name: str = "PP-OCRv5_server_det",
|
|
|
+ aspect_ratio_threshold: float = 1.2,
|
|
|
+ vertical_box_ratio_threshold: float = 0.28,
|
|
|
+ min_vertical_boxes: int = 3,
|
|
|
+ text_box_aspect_ratio: float = 0.8,
|
|
|
+ ):
|
|
|
+ """
|
|
|
+ Args:
|
|
|
+ cnn_model_name: CNN 方向分类模型名称
|
|
|
+ ocr_model_name: OCR 文本检测模型名称
|
|
|
+ aspect_ratio_threshold: 图像宽高比阈值（> 此值才进行 OCR 检测）
|
|
|
+ vertical_box_ratio_threshold: 垂直文本框占比阈值
|
|
|
+ min_vertical_boxes: 最少垂直文本框数量
|
|
|
+ text_box_aspect_ratio: 文本框宽高比阈值（< 此值为垂直文本）
|
|
|
+ """
|
|
|
+ # 1. 加载 CNN 方向分类模型
|
|
|
+ self.cnn_model = create_model(cnn_model_name)
|
|
|
+
|
|
|
+ # 2. 加载 OCR 文本检测模型
|
|
|
+ self.ocr_detector = create_model(ocr_model_name)
|
|
|
+
|
|
|
+ # 3. 参数设置
|
|
|
+ self.aspect_ratio_threshold = aspect_ratio_threshold
|
|
|
+ self.vertical_box_ratio_threshold = vertical_box_ratio_threshold
|
|
|
+ self.min_vertical_boxes = min_vertical_boxes
|
|
|
+ self.text_box_aspect_ratio = text_box_aspect_ratio
|
|
|
+
|
|
|
+ # 4. 标签映射
|
|
|
+ self.labels = ["0", "90", "180", "270"]
|
|
|
+
|
|
|
+ def predict(
|
|
|
+ self,
|
|
|
+ img: Union[str, np.ndarray, Image.Image],
|
|
|
+ use_ocr_filter: bool = True,
|
|
|
+ batch_size: int = 1,
|
|
|
+ ) -> dict:
|
|
|
+ """
|
|
|
+ 预测图像方向
|
|
|
+
|
|
|
+ Args:
|
|
|
+ img: 输入图像
|
|
|
+ use_ocr_filter: 是否使用 OCR 过滤（False 则直接使用 CNN）
|
|
|
+ batch_size: 批处理大小
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ {
|
|
|
+ "orientation": "0", # "0", "90", "180", "270"
|
|
|
+ "confidence": 0.95,
|
|
|
+ "method": "ocr_filter" or "cnn",
|
|
|
+ "details": {...}
|
|
|
+ }
|
|
|
+ """
|
|
|
+ # 统一输入格式为 numpy array
|
|
|
+ if isinstance(img, str):
|
|
|
+ img = cv2.imread(img)
|
|
|
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
|
+ elif isinstance(img, Image.Image):
|
|
|
+ img = np.array(img)
|
|
|
+
|
|
|
+ # ============================================
|
|
|
+ # 阶段1: 快速过滤（宽高比检查）
|
|
|
+ # ============================================
|
|
|
+ img_height, img_width = img.shape[:2]
|
|
|
+ img_aspect_ratio = img_height / img_width if img_width > 0 else 1.0
|
|
|
+
|
|
|
+ if not use_ocr_filter or img_aspect_ratio <= self.aspect_ratio_threshold:
|
|
|
+ # 横向图像，直接返回 0 度
|
|
|
+ return {
|
|
|
+ "orientation": "0",
|
|
|
+ "confidence": 1.0,
|
|
|
+ "method": "aspect_ratio_filter",
|
|
|
+ "details": {
|
|
|
+ "img_aspect_ratio": img_aspect_ratio,
|
|
|
+ "threshold": self.aspect_ratio_threshold,
|
|
|
+ "reason": "Image is landscape, no rotation needed"
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ # ============================================
|
|
|
+ # 阶段2: OCR 文本框分析
|
|
|
+ # ============================================
|
|
|
+ # 转换为 BGR（PaddleOCR 需要）
|
|
|
+ bgr_img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
|
|
+
|
|
|
+ # OCR 检测文本框
|
|
|
+ det_result = self.ocr_detector.predict(bgr_img)
|
|
|
+ det_boxes = det_result.get("boxes", [])
|
|
|
+
|
|
|
+ if not det_boxes:
|
|
|
+ # 没有检测到文本框，使用 CNN
|
|
|
+ return self._cnn_predict(img, reason="No text boxes detected")
|
|
|
+
|
|
|
+ # 分析文本框的方向
|
|
|
+ vertical_count = 0
|
|
|
+ for box in det_boxes:
|
|
|
+ # 提取文本框坐标
|
|
|
+ coords = box.get("coordinate", [])
|
|
|
+ if len(coords) < 4:
|
|
|
+ continue
|
|
|
+
|
|
|
+ # 计算文本框的宽度和高度
|
|
|
+ # PaddleX 返回的是 [x1, y1, x2, y2] 格式
|
|
|
+ x1, y1, x2, y2 = coords[:4]
|
|
|
+ width = abs(x2 - x1)
|
|
|
+ height = abs(y2 - y1)
|
|
|
+
|
|
|
+ aspect_ratio = width / height if height > 0 else 1.0
|
|
|
+
|
|
|
+ # 判断是否为垂直文本框
|
|
|
+ if aspect_ratio < self.text_box_aspect_ratio:
|
|
|
+ vertical_count += 1
|
|
|
+
|
|
|
+ # 计算垂直文本框占比
|
|
|
+ vertical_ratio = vertical_count / len(det_boxes) if det_boxes else 0
|
|
|
+
|
|
|
+ # 判断是否需要旋转
|
|
|
+ is_rotated = (
|
|
|
+ vertical_ratio >= self.vertical_box_ratio_threshold
|
|
|
+ and vertical_count >= self.min_vertical_boxes
|
|
|
+ )
|
|
|
+
|
|
|
+ if not is_rotated:
|
|
|
+ # 文本框正常，不需要旋转
|
|
|
+ return {
|
|
|
+ "orientation": "0",
|
|
|
+ "confidence": 1.0,
|
|
|
+ "method": "ocr_filter",
|
|
|
+ "details": {
|
|
|
+ "total_boxes": len(det_boxes),
|
|
|
+ "vertical_boxes": vertical_count,
|
|
|
+ "vertical_ratio": vertical_ratio,
|
|
|
+ "threshold": self.vertical_box_ratio_threshold,
|
|
|
+ "reason": "Text boxes are mostly horizontal"
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ # ============================================
|
|
|
+ # 阶段3: CNN 精确分类
|
|
|
+ # ============================================
|
|
|
+ return self._cnn_predict(
|
|
|
+ img,
|
|
|
+ reason=f"OCR detected rotation (vertical_ratio={vertical_ratio:.2f})",
|
|
|
+ ocr_details={
|
|
|
+ "total_boxes": len(det_boxes),
|
|
|
+ "vertical_boxes": vertical_count,
|
|
|
+ "vertical_ratio": vertical_ratio,
|
|
|
+ }
|
|
|
+ )
|
|
|
+
|
|
|
+ def _cnn_predict(self, img: np.ndarray, reason: str = "", ocr_details: dict = None) -> dict:
|
|
|
+ """使用 CNN 模型预测方向"""
|
|
|
+ # 转换为 PIL Image（PaddleX 需要）
|
|
|
+ if isinstance(img, np.ndarray):
|
|
|
+ img = Image.fromarray(img)
|
|
|
+
|
|
|
+ # CNN 推理
|
|
|
+ result = self.cnn_model.predict(img)
|
|
|
+
|
|
|
+ # 提取结果
|
|
|
+ orientation = result.get("label", "0")
|
|
|
+ confidence = result.get("score", 0.0)
|
|
|
+
|
|
|
+ return {
|
|
|
+ "orientation": orientation,
|
|
|
+ "confidence": confidence,
|
|
|
+ "method": "cnn",
|
|
|
+ "details": {
|
|
|
+ "reason": reason,
|
|
|
+ "ocr_analysis": ocr_details or {},
|
|
|
+ "cnn_scores": result.get("label_names", [])
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ def batch_predict(
|
|
|
+ self,
|
|
|
+ imgs: List[Union[str, np.ndarray, Image.Image]],
|
|
|
+ use_ocr_filter: bool = True,
|
|
|
+ batch_size: int = 8,
|
|
|
+ ) -> List[dict]:
|
|
|
+ """批量预测"""
|
|
|
+ results = []
|
|
|
+ for img in imgs:
|
|
|
+ result = self.predict(img, use_ocr_filter=use_ocr_filter)
|
|
|
+ results.append(result)
|
|
|
+ return results
|
