ocr_verify/merger/bbox_extractor.py

"""
bbox 提取模块
负责从 PaddleOCR 结果中提取文字框信息
"""
from typing import List, Dict
import numpy as np
from pathlib import Path


class BBoxExtractor:
    """bbox 提取器"""
    
    @staticmethod
    def extract_paddle_text_boxes(paddle_data: Dict) -> List[Dict]:
        """
        提取 PaddleOCR 的文字框信息
        
        Args:
            paddle_data: PaddleOCR 输出的数据
        
        Returns:
            文字框列表（坐标已转换为 angle=0 时的坐标）
        """
        text_boxes = []
        
        if 'overall_ocr_res' not in paddle_data:
            return text_boxes
        
        ocr_res = paddle_data['overall_ocr_res']
        rec_texts = ocr_res.get('rec_texts', [])
        rec_polys = ocr_res.get('rec_polys', [])
        rec_scores = ocr_res.get('rec_scores', [])
        
        # 🎯 获取旋转角度
        rotation_angle = BBoxExtractor._get_rotation_angle(paddle_data)
        
        # 🎯 如果有旋转，需要获取原始图像尺寸
        orig_image_size = None
        
        if rotation_angle != 0:
            orig_image_size = BBoxExtractor._get_original_image_size(paddle_data)
            print(f"🔄 检测到旋转角度: {rotation_angle}°")
            print(f"📐 原始图像尺寸: {orig_image_size[0]} x {orig_image_size[1]}")
        
        for i, (text, poly, score) in enumerate(zip(rec_texts, rec_polys, rec_scores)):
            if text and text.strip():
                # 🎯 如果有旋转角度，转换坐标
                if rotation_angle != 0 and orig_image_size:
                    poly = BBoxExtractor._inverse_rotate_coordinates(
                        poly, rotation_angle, orig_image_size
                    )
                
                # 计算 bbox (x_min, y_min, x_max, y_max)
                bbox = BBoxExtractor._poly_to_bbox(poly)
                
                text_boxes.append({
                    'text': text,
                    'bbox': bbox,
                    'poly': poly,
                    'score': score,
                    'paddle_bbox_index': i,
                    'used': False
                })
        
        return text_boxes
    
    @staticmethod
    def _get_rotation_angle(paddle_data: Dict) -> float:
        """获取旋转角度"""
        if 'doc_preprocessor_res' not in paddle_data:
            return 0.0
        
        doc_res = paddle_data['doc_preprocessor_res']
        if isinstance(doc_res, dict) and 'angle' in doc_res:
            return float(doc_res['angle'])
        
        return 0.0
    
    @staticmethod
    def _get_original_image_size(paddle_data: Dict) -> tuple:
        """
        获取原始图像尺寸（从图片文件读取）
        
        Args:
            paddle_data: PaddleOCR 数据
        
        Returns:
            (width, height) 元组
        """
        from PIL import Image
        
        # 🎯 从 input_path 读取图像
        input_path = paddle_data.get('input_path')
        
        if input_path and Path(input_path).exists():
            try:
                with Image.open(input_path) as img:
                    # 返回原始图像尺寸
                    return img.size  # (width, height)
            except Exception as e:
                print(f"⚠️ 无法读取图像文件 {input_path}: {e}")
        
        # 🎯 降级方案：从 layout_det_res 推断
        if 'layout_det_res' in paddle_data:
            layout_res = paddle_data['layout_det_res']
            if 'boxes' in layout_res and layout_res['boxes']:
                max_x = 0
                max_y = 0
                for box in layout_res['boxes']:
                    coord = box.get('coordinate', [])
                    if len(coord) >= 4:
                        max_x = max(max_x, coord[2])
                        max_y = max(max_y, coord[3])
                
                if max_x > 0 and max_y > 0:
                    return (int(max_x) + 50, int(max_y) + 50)
        
        # 🎯 最后降级：从 overall_ocr_res 推断
        if 'overall_ocr_res' in paddle_data:
            ocr_res = paddle_data['overall_ocr_res']
            rec_polys = ocr_res.get('rec_polys', [])
            if rec_polys:
                max_x = 0
                max_y = 0
                for poly in rec_polys:
                    for point in poly:
                        max_x = max(max_x, point[0])
                        max_y = max(max_y, point[1])
                
                if max_x > 0 and max_y > 0:
                    return (int(max_x) + 50, int(max_y) + 50)
        
        # 🎯 默认 A4 尺寸
        print("⚠️ 无法确定原始图像尺寸，使用默认值")
        return (2480, 3508)
    
    @staticmethod
    def _inverse_rotate_coordinates(poly: List[List[float]], 
                                    angle: float,
                                    orig_image_size: tuple) -> List[List[float]]:
        """
        反向旋转坐标
        
        参考 ocr_validator_utils.rotate_image_and_coordinates 的逆操作
        
        PaddleOCR 在旋转后的图像上识别，坐标是旋转后的
        我们需要将坐标转换回原始图像（未旋转）
        
        Args:
            poly: 旋转后图像上的多边形坐标 [[x',y'], ...]
            angle: 旋转角度（度数，PaddleX 使用的角度）
            orig_image_size: 原始图像尺寸 (width, height)
        
        Returns:
            原始图像上的多边形坐标 [[x,y], ...]
        """
        orig_width, orig_height = orig_image_size
        
        # 🎯 根据旋转角度计算旋转后的图像尺寸
        if angle == 90:
            rotated_width, rotated_height = orig_height, orig_width
        elif angle == 270:
            rotated_width, rotated_height = orig_height, orig_width
        else:
            rotated_width, rotated_height = orig_width, orig_height
        
        inverse_poly = []
        
        for point in poly:
            x_rot, y_rot = point[0], point[1]  # 旋转后的坐标
            
            # 🎯 反向旋转（参考 rotate_image_and_coordinates 的逆操作）
            if angle == 90:
                # 正向: rotated = image.rotate(90, expand=True)
                #      x_rot = y_orig
                #      y_rot = rotated_width - x_orig = orig_height - x_orig
                # 反向: x_orig = rotated_width - y_rot = orig_height - y_rot
                #      y_orig = x_rot
                x_orig = rotated_width - y_rot
                y_orig = x_rot
                
            elif angle == 270:
                # 正向: rotated = image.rotate(-90, expand=True)
                #      x_rot = rotated_width - y_orig = orig_height - y_orig
                #      y_rot = x_orig
                # 反向: y_orig = rotated_width - x_rot = orig_height - x_rot
                #      x_orig = y_rot
                x_orig = y_rot
                y_orig = rotated_width - x_rot
                
            elif angle == 180:
                # 正向: rotated = image.rotate(180)
                #      x_rot = orig_width - x_orig
                #      y_rot = orig_height - y_orig
                # 反向: x_orig = orig_width - x_rot
                #      y_orig = orig_height - y_rot
                x_orig = orig_width - x_rot
                y_orig = orig_height - y_rot
                
            else:
                # 其他角度或0度，不转换
                x_orig = x_rot
                y_orig = y_rot
            
            inverse_poly.append([x_orig, y_orig])
        
        return inverse_poly
    
    @staticmethod
    def _poly_to_bbox(poly: List[List[float]]) -> List[float]:
        """将多边形转换为 bbox [x_min, y_min, x_max, y_max]"""
        xs = [p[0] for p in poly]
        ys = [p[1] for p in poly]
        return [min(xs), min(ys), max(xs), max(ys)]
    
    @staticmethod
    def extract_table_cells_with_bbox(merged_data: List[Dict]) -> List[Dict]:
        """
        提取所有表格单元格及其 bbox 信息
        
        Args:
            merged_data: 合并后的数据
        
        Returns:
            单元格列表
        """
        import json
        from bs4 import BeautifulSoup
        
        cells = []
        
        for item in merged_data:
            if item['type'] != 'table':
                continue
            
            html = item.get('table_body_with_bbox', item.get('table_body', ''))
            soup = BeautifulSoup(html, 'html.parser')
            
            for row_idx, row in enumerate(soup.find_all('tr')):
                for col_idx, cell in enumerate(row.find_all(['td', 'th'])):
                    cell_text = cell.get_text(strip=True)
                    bbox_str = cell.get('data-bbox', '')
                    
                    if bbox_str:
                        try:
                            bbox = json.loads(bbox_str)
                            cells.append({
                                'text': cell_text,
                                'bbox': bbox,
                                'row': row_idx,
                                'col': col_idx,
                                'score': float(cell.get('data-score', 0)),
                                'paddle_index': int(cell.get('data-paddle-index', -1))
                            })
                        except (json.JSONDecodeError, ValueError):
                            pass
        
        return cells