ocr_platform/legacy_table_line_generator/smart_generator.py

"""
智能表格线生成器
整合列检测 + 自适应行分割，提供统一入口
"""
import json
import cv2
import numpy as np
from pathlib import Path
from typing import List, Dict, Tuple, Optional, Union
from dataclasses import dataclass, asdict

from .adaptive_row_splitter import AdaptiveRowSplitter, RowRegion
from .column_detector import ColumnBoundaryDetector, ColumnRegion


@dataclass
class TableStructure:
    """表格结构"""
    table_region: Tuple[int, int, int, int]  # (x1, y1, x2, y2)
    columns: List[ColumnRegion]
    rows: List[RowRegion]
    page_size: Tuple[int, int]  # (width, height)
    
    def get_vertical_lines(self) -> List[int]:
        """获取所有竖线的 X 坐标"""
        if not self.columns:
            return []
        lines = [self.columns[0].x_left]
        for col in self.columns:
            lines.append(col.x_right)
        return lines
    
    def get_horizontal_lines(self) -> List[int]:
        """获取所有横线的 Y 坐标"""
        if not self.rows:
            return []
        lines = [self.rows[0].y_top]
        for row in self.rows:
            lines.append(row.y_bottom)
        return lines
    
    def get_cell_bboxes(self) -> List[List[Tuple[int, int, int, int]]]:
        """获取所有单元格的 bbox"""
        cells = []
        for row in self.rows:
            row_cells = []
            for col in self.columns:
                cell_bbox = (col.x_left, row.y_top, col.x_right, row.y_bottom)
                row_cells.append(cell_bbox)
            cells.append(row_cells)
        return cells
    
    def to_dict(self) -> Dict:
        return {
            'table_region': self.table_region,
            'columns': [c.to_dict() for c in self.columns],
            'rows': [r.to_dict() for r in self.rows],
            'page_size': self.page_size,
            'vertical_lines': self.get_vertical_lines(),
            'horizontal_lines': self.get_horizontal_lines()
        }


class SmartTableLineGenerator:
    """
    智能表格线生成器
    
    功能：
    1. 自动检测表格区域
    2. 自动检测列边界
    3. 自适应行分割（支持可变行高）
    4. 生成表格线坐标
    5. 可选：在图片上绘制表格线
    """
    
    def __init__(self,
                 # 列检测参数
                 x_tolerance: int = 20,
                 min_boxes_per_column: int = 3,
                 # 行分割参数
                 min_gap_height: int = 6,
                 density_threshold: float = 0.05,
                 min_row_height: int = 15,
                 # 其他参数
                 table_margin: int = 10,
                 header_detection: bool = True):
        
        self.column_detector = ColumnBoundaryDetector(
            x_tolerance=x_tolerance,
            min_boxes_per_column=min_boxes_per_column
        )
        
        self.row_splitter = AdaptiveRowSplitter(
            min_gap_height=min_gap_height,
            density_threshold=density_threshold,
            min_row_height=min_row_height
        )
        
        self.table_margin = table_margin
        self.header_detection = header_detection
    
    def generate(self,
                 ocr_boxes: List[Dict],
                 page_size: Tuple[int, int],
                 table_region: Optional[Tuple[int, int, int, int]] = None,
                 header_row_count: int = 1,
                 debug: bool = False) -> Tuple[TableStructure, Optional[Dict]]:
        """
        生成表格结构
        
        Args:
            ocr_boxes: OCR 结果 [{text, bbox: [x1,y1,x2,y2]}]
            page_size: (width, height)
            table_region: 表格区域 (可选，自动检测)
            header_row_count: 表头行数
            debug: 是否返回调试信息
            
        Returns:
            (TableStructure, debug_info)
        """
        width, height = page_size
        debug_info = {} if debug else None
        
        # 1. 自动检测表格区域
        if table_region is None:
            table_region = self._detect_table_region(ocr_boxes, width, height)
        
        if debug:
            debug_info['table_region'] = table_region
        
        # 2. 检测列边界
        columns = self.column_detector.detect(ocr_boxes, table_region, width)
        
        if debug:
            debug_info['columns'] = [c.to_dict() for c in columns]
        
        # 3. 自适应行分割
        rows, row_debug = self.row_splitter.split_rows(
            ocr_boxes, table_region, debug=debug
        )
        
        if debug and row_debug:
            debug_info['row_splitter'] = row_debug
        
        # 4. 构建表格结构
        structure = TableStructure(
            table_region=table_region,
            columns=columns,
            rows=rows,
            page_size=page_size
        )
        
        return structure, debug_info
    
    def generate_from_image(self,
                            image_path: Union[str, Path],
                            ocr_boxes: List[Dict],
                            table_region: Optional[Tuple[int, int, int, int]] = None,
                            debug: bool = False) -> Tuple[TableStructure, Optional[Dict]]:
        """
        从图片生成表格结构
        
        自动获取图片尺寸
        """
        image = cv2.imread(str(image_path))
        if image is None:
            raise ValueError(f"无法读取图片: {image_path}")
        
        height, width = image.shape[:2]
        return self.generate(ocr_boxes, (width, height), table_region, debug=debug)
    
    def draw_table_lines(self,
                         image: np.ndarray,
                         structure: TableStructure,
                         line_color: Tuple[int, int, int] = (0, 0, 255),
                         line_thickness: int = 1,
                         draw_cells: bool = False) -> np.ndarray:
        """
        在图片上绘制表格线
        
        Args:
            image: 原始图片
            structure: 表格结构
            line_color: 线条颜色 (BGR)
            line_thickness: 线条粗细
            draw_cells: 是否绘制单元格边框
            
        Returns:
            绘制后的图片
        """
        result = image.copy()
        
        v_lines = structure.get_vertical_lines()
        h_lines = structure.get_horizontal_lines()
        
        if not v_lines or not h_lines:
            return result
        
        y_top = h_lines[0]
        y_bottom = h_lines[-1]
        x_left = v_lines[0]
        x_right = v_lines[-1]
        
        # 绘制竖线
        for x in v_lines:
            cv2.line(result, (x, y_top), (x, y_bottom), line_color, line_thickness)
        
        # 绘制横线
        for y in h_lines:
            cv2.line(result, (x_left, y), (x_right, y), line_color, line_thickness)
        
        return result
    
    def build_table_data(self,
                         ocr_boxes: List[Dict],
                         structure: TableStructure) -> List[List[str]]:
        """
        根据表格结构构建结构化数据
        
        Returns:
            二维列表 table[row_idx][col_idx] = cell_text
        """
        n_rows = len(structure.rows)
        n_cols = len(structure.columns)
        
        # 初始化表格
        table = [[[] for _ in range(n_cols)] for _ in range(n_rows)]
        
        # 将每个 box 分配到对应的单元格
        for box in ocr_boxes:
            row_idx = self._find_row_index(box, structure.rows)
            col_idx = self._find_column_index(box, structure.columns)
            
            if row_idx >= 0 and col_idx >= 0:
                table[row_idx][col_idx].append(box['text'])
        
        # 合并每个单元格的文本
        result = []
        for row in table:
            result.append([' '.join(texts) for texts in row])
        
        return result
    
    def _detect_table_region(self,
                             boxes: List[Dict],
                             width: int,
                             height: int) -> Tuple[int, int, int, int]:
        """自动检测表格区域"""
        if not boxes:
            return (0, 0, width, height)
        
        # 使用所有 boxes 的边界框
        x1 = min(b['bbox'][0] for b in boxes)
        y1 = min(b['bbox'][1] for b in boxes)
        x2 = max(b['bbox'][2] for b in boxes)
        y2 = max(b['bbox'][3] for b in boxes)
        
        # 留边距
        return (
            max(0, int(x1 - self.table_margin)),
            max(0, int(y1 - self.table_margin)),
            min(width, int(x2 + self.table_margin)),
            min(height, int(y2 + self.table_margin))
        )
    
    def _find_row_index(self, box: Dict, rows: List[RowRegion]) -> int:
        """找到 box 所属的行索引"""
        cy = (box['bbox'][1] + box['bbox'][3]) / 2
        for i, row in enumerate(rows):
            if row.y_top <= cy < row.y_bottom:
                return i
        return -1
    
    def _find_column_index(self, box: Dict, columns: List[ColumnRegion]) -> int:
        """找到 box 所属的列索引"""
        cx = (box['bbox'][0] + box['bbox'][2]) / 2
        for i, col in enumerate(columns):
            if col.x_left - 10 <= cx <= col.x_right + 10:
                return i
        return -1
    
    def save_structure(self, structure: TableStructure, output_path: Union[str, Path]):
        """保存表格结构到 JSON"""
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(structure.to_dict(), f, ensure_ascii=False, indent=2)
    
    def load_structure(self, input_path: Union[str, Path]) -> TableStructure:
        """从 JSON 加载表格结构"""
        with open(input_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        columns = [
            ColumnRegion(x_left=c['x_left'], x_right=c['x_right'], column_index=c['column_index'])
            for c in data['columns']
        ]
        
        rows = [
            RowRegion(y_top=r['y_top'], y_bottom=r['y_bottom'], row_index=r['row_index'], boxes=[])
            for r in data['rows']
        ]
        
        return TableStructure(
            table_region=tuple(data['table_region']),
            columns=columns,
            rows=rows,
            page_size=tuple(data['page_size'])
        )