PaddleX/zhch/unified_pytorch_models/orientation_classifier_v2.py

"""
增强版文档方向分类模块 - 独立版本
无需依赖 PaddleX 内部结构
"""
import cv2
import numpy as np
import onnxruntime as ort
from typing import Dict, Tuple, Optional
from pathlib import Path
from dataclasses import dataclass


@dataclass
class OrientationResult:
    """方向分类结果"""
    rotation_angle: str = "0"  # "0", "90", "180", "270"
    confidence: float = 1.0
    needs_rotation: bool = False
    vertical_text_count: int = 0
    aspect_ratio: float = 1.0
    
    def __str__(self):
        return (
            f"OrientationResult(\n"
            f"  angle={self.rotation_angle}°, "
            f"  confidence={self.confidence:.3f}, "
            f"  needs_rotation={self.needs_rotation}, "
            f"  vertical_texts={self.vertical_text_count}, "
            f"  aspect_ratio={self.aspect_ratio:.2f}\n"
            f")"
        )


class OrientationClassifierV2:
    """
    增强版方向分类器
    参考 MinerU 的两阶段检测策略
    """
    
    def __init__(
        self, 
        model_path: str,
        text_detector=None,  # 可选的 OCR 检测器
        aspect_ratio_threshold: float = 1.2,
        vertical_text_ratio: float = 0.28,
        vertical_text_min_count: int = 3,
        use_gpu: bool = False
    ):
        """
        Args:
            model_path: ONNX 模型路径
            text_detector: 文本检测器(可选,用于辅助判断)
            aspect_ratio_threshold: 长宽比阈值
            vertical_text_ratio: 垂直文本框占比阈值
            vertical_text_min_count: 最小垂直文本框数量
            use_gpu: 是否使用GPU
        """
        # 初始化 ONNX Runtime
        providers = ['CoreMLExecutionProvider', 'CPUExecutionProvider'] if use_gpu else ['CPUExecutionProvider']
        self.session = ort.InferenceSession(model_path, providers=providers)
        
        self.text_detector = text_detector
        self.aspect_ratio_threshold = aspect_ratio_threshold
        self.vertical_text_ratio = vertical_text_ratio
        self.vertical_text_min_count = vertical_text_min_count
        
        # 预计算标准化系数 (ImageNet 标准)
        self.mean = np.array([0.485, 0.456, 0.406])
        self.std = np.array([0.229, 0.224, 0.225])
        self.scale = 1.0 / 255.0
        
        self.target_size = 256  # 缩放后的最短边
        self.crop_size = (224, 224)  # 裁剪尺寸
        self.labels = ["0", "90", "180", "270"]
        
        print(f"✅ Orientation classifier initialized")
        print(f"   Model: {Path(model_path).name}")
        print(f"   Aspect ratio threshold: {aspect_ratio_threshold}")
        print(f"   Vertical text ratio: {vertical_text_ratio}")
    
    def _needs_rotation_check(self, img: np.ndarray) -> Tuple[bool, float]:
        """检查图像是否需要进行旋转检测"""
        h, w = img.shape[:2]
        aspect_ratio = h / w if w > 0 else 1.0
        needs_check = aspect_ratio > self.aspect_ratio_threshold
        return needs_check, aspect_ratio
    
    def _detect_vertical_text(self, img: np.ndarray) -> Tuple[bool, int]:
        """
        使用文本检测判断是否存在大量垂直文本
        
        Returns:
            (is_rotated, vertical_count): 是否旋转, 垂直文本框数量
        """
        if self.text_detector is None:
            return False, 0
        
        try:
            # ✅ 修改：适配 MinerUOCRAdapter 的返回格式
            # 返回格式: [[[box], (text, conf)], ...] 或 [[boxes], ...]
            det_results = self.text_detector.ocr(img, det=True, rec=False)
            
            if not det_results or not det_results[0]:
                return False, 0
            
            boxes = det_results[0]
            
            # ✅ 处理两种格式
            vertical_count = 0
            for item in boxes:
                # 格式1: [box] (仅检测)
                # 格式2: [[box], (text, conf)] (检测+识别)
                if isinstance(item, list) and len(item) > 0:
                    if isinstance(item[0], list):
                        # 格式2: [[box], ...]
                        box = np.array(item[0])
                    else:
                        # 格式1: [box]
                        box = np.array(item)
                else:
                    continue
                
                # 计算文本框的宽高
                if len(box) >= 4:
                    points = box
                    width = np.linalg.norm(points[1] - points[0])
                    height = np.linalg.norm(points[2] - points[1])
                    
                    aspect_ratio = width / height if height > 0 else 1.0
                    
                    # 统计垂直文本框 (高 > 宽)
                    if aspect_ratio < 0.8:
                        vertical_count += 1
            
            # 判断是否需要旋转
            total_boxes = len(boxes)
            is_rotated = (
                vertical_count >= total_boxes * self.vertical_text_ratio 
                and vertical_count >= self.vertical_text_min_count
            )
            
            return is_rotated, vertical_count
            
        except Exception as e:
            print(f"⚠️  Text detection failed: {e}")
            import traceback
            traceback.print_exc()
            return False, 0
    
    def _preprocess(self, img: np.ndarray) -> np.ndarray:
        """
        图像预处理
        1. 缩放最短边到 256
        2. 中心裁剪到 224×224
        3. 标准化
        """
        h, w = img.shape[:2]
        
        # 1. 缩放
        scale = self.target_size / min(h, w)
        new_h = round(h * scale)
        new_w = round(w * scale)
        img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
        
        # 2. 中心裁剪
        h, w = img.shape[:2]
        cw, ch = self.crop_size
        x1 = max(0, (w - cw) // 2)
        y1 = max(0, (h - ch) // 2)
        x2 = min(w, x1 + cw)
        y2 = min(h, y1 + ch)
        
        if w < cw or h < ch:
            # Padding instead of error
            padded = np.ones((ch, cw, 3), dtype=np.uint8) * 114
            paste_h = min(h, ch)
            paste_w = min(w, cw)
            padded[:paste_h, :paste_w] = img[:paste_h, :paste_w]
            img = padded
        else:
            img = img[y1:y2, x1:x2]
        
        # 3. 标准化 (转 RGB + ImageNet 标准化)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = img.astype(np.float32) * self.scale  # [0, 1]
        
        # 分通道标准化
        for c in range(3):
            img[:, :, c] = (img[:, :, c] - self.mean[c]) / self.std[c]
        
        # 4. 转换为 NCHW 格式
        img = img.transpose((2, 0, 1))
        img = np.expand_dims(img, axis=0)
        
        return img.astype(np.float32)
    
    def predict(self, img: np.ndarray, return_debug: bool = False) -> OrientationResult:
        """
        预测图像方向
        
        Args:
            img: BGR 格式的输入图像
            return_debug: 是否返回调试信息
            
        Returns:
            OrientationResult 对象
        """
        result = OrientationResult()
        
        # 1. 检查长宽比
        needs_check, aspect_ratio = self._needs_rotation_check(img)
        result.aspect_ratio = aspect_ratio
        
        if not needs_check:
            if return_debug:
                print(f"   ⏭️  Skipped (aspect_ratio={aspect_ratio:.2f} <= {self.aspect_ratio_threshold})")
            return result
        
        # 2. 使用文本检测判断是否旋转
        if self.text_detector:
            is_rotated, vertical_count = self._detect_vertical_text(img)
            result.vertical_text_count = vertical_count
            
            if not is_rotated:
                if return_debug:
                    print(f"   ⏭️  No rotation needed (vertical_texts={vertical_count})")
                return result
            
        # 3. 使用分类模型预测旋转角度
        input_tensor = self._preprocess(img)
        
        # ONNX 推理
        input_name = self.session.get_inputs()[0].name
        output_name = self.session.get_outputs()[0].name
        outputs = self.session.run([output_name], {input_name: input_tensor})
        
        probabilities = outputs[0][0]  # [4,]
        
        predicted_idx = np.argmax(probabilities)
        result.rotation_angle = self.labels[predicted_idx]
        result.confidence = float(probabilities[predicted_idx])
        result.needs_rotation = result.rotation_angle != '0'
        
        if return_debug:
            print(f"   🎯 Predicted angle: {result.rotation_angle}° (conf={result.confidence:.3f})")
        
        return result
    
    def rotate_image(self, img: np.ndarray, angle: str) -> np.ndarray:
        """根据预测角度旋转图像"""
        if angle == "90":
            return cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
        elif angle == "180":
            return cv2.rotate(img, cv2.ROTATE_180)
        elif angle == "270":
            return cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
        else:
            return img

if __name__ == "__main__":
    # 测试代码
    model_path = "/Users/zhch158/workspace/repository.git/PaddleX/zhch/unified_pytorch_models/Layout/PP-LCNet_x1_0_doc_ori.onnx"  # 替换为实际模型路径
    classifier = OrientationClassifierV2(model_path=model_path, use_gpu=False)
    
    test_image_path = "/Users/zhch158/workspace/data/至远彩色印刷工业有限公司/2023年度报告母公司.img/2023年度报告母公司_page_003.png"  # 替换为实际测试图像路径
    output_image_path = Path(f"/Users/zhch158/workspace/repository.git/PaddleX/zhch/sample_data/PP-LCNet_x1_0_doc_ori.onnx/{Path(test_image_path).name}.jpg")
    img = cv2.imread(test_image_path)
    
    result = classifier.predict(img, return_debug=True)
    print(result)
    
    if result.needs_rotation:
        output_image_path.parent.mkdir(exist_ok=True)
        rotated_img = classifier.rotate_image(img, result.rotation_angle)
        cv2.imwrite(output_image_path.as_posix(), rotated_img)