ocr_verify/ocr_validator_layout.py

#!/usr/bin/env python3
"""
OCR验证工具的布局管理模块
包含标准布局、滚动布局、紧凑布局的实现
"""

import streamlit as st
from pathlib import Path
from PIL import Image
from typing import Dict, List, Optional
import plotly.graph_objects as go

from ocr_validator_utils import (
    convert_html_table_to_markdown, 
    parse_html_tables,
    draw_bbox_on_image,
    rotate_image_and_coordinates
)


class OCRLayoutManager:
    """OCR布局管理器"""
    
    def __init__(self, validator):
        self.validator = validator
        self.config = validator.config
        self._rotated_image_cache = {}  # 缓存旋转后的图像
    
    def get_rotation_angle(self) -> float:
        """从OCR数据中获取旋转角度"""
        if self.validator.ocr_data:
            for item in self.validator.ocr_data:
                if isinstance(item, dict) and 'rotation_angle' in item:
                    return item['rotation_angle']
        return 0.0
    
    def load_and_rotate_image(self, image_path: str) -> Optional[Image.Image]:
        """加载并根据需要旋转图像"""
        if not image_path or not Path(image_path).exists():
            return None
            
        # 检查缓存
        rotation_angle = self.get_rotation_angle()
        cache_key = f"{image_path}_{rotation_angle}"
        
        if cache_key in self._rotated_image_cache:
            return self._rotated_image_cache[cache_key]
        
        try:
            image = Image.open(image_path)
            
            # 如果需要旋转
            if rotation_angle != 0:
                st.info(f"🔄 检测到文档旋转角度: {rotation_angle}°，正在自动旋转图像...")
                
                # 收集所有bbox坐标
                all_bboxes = []
                text_to_bbox_map = {}  # 记录文本到bbox索引的映射
                
                bbox_index = 0
                for text, info_list in self.validator.text_bbox_mapping.items():
                    text_to_bbox_map[text] = []
                    for info in info_list:
                        all_bboxes.append(info['bbox'])
                        text_to_bbox_map[text].append(bbox_index)
                        bbox_index += 1
                
                # 旋转图像和坐标
                rotated_image, rotated_bboxes = rotate_image_and_coordinates(
                    image, rotation_angle, all_bboxes
                )
                
                # 更新bbox映射 - 使用映射关系确保正确对应
                for text, bbox_indices in text_to_bbox_map.items():
                    for i, bbox_idx in enumerate(bbox_indices):
                        if bbox_idx < len(rotated_bboxes) and i < len(self.validator.text_bbox_mapping[text]):
                            self.validator.text_bbox_mapping[text][i]['bbox'] = rotated_bboxes[bbox_idx]
                
                # 缓存结果
                self._rotated_image_cache[cache_key] = rotated_image
                return rotated_image
            else:
                # 无需旋转，直接缓存原图
                self._rotated_image_cache[cache_key] = image
                return image
                
        except Exception as e:
            st.error(f"❌ 图像加载失败: {e}")
            return None
    
    def render_content_section(self, layout_type: str = "standard"):
        """渲染内容区域 - 统一方法"""
        st.header("📄 OCR识别内容")
        
        # 显示旋转信息
        # rotation_angle = self.get_rotation_angle()
        # if rotation_angle != 0:
        #     st.info(f"📐 文档旋转角度: {rotation_angle}°")
        
        # 文本选择器
        if self.validator.text_bbox_mapping:
            text_options = ["请选择文本..."] + list(self.validator.text_bbox_mapping.keys())
            selected_index = st.selectbox(
                "选择要校验的文本",
                range(len(text_options)),
                format_func=lambda x: text_options[x][:50] + "..." if len(text_options[x]) > 50 else text_options[x],
                key=f"{layout_type}_text_selector"
            )
            
            if selected_index > 0:
                st.session_state.selected_text = text_options[selected_index]
        else:
            st.warning("没有找到可点击的文本")
    
    def render_md_content(self, layout_type: str):
        """渲染Markdown内容 - 统一方法"""
        if not self.validator.md_content:
            return None, None
            
        # 搜索功能
        search_term = st.text_input(
            "🔍 搜索文本内容", 
            placeholder="输入关键词搜索...", 
            key=f"{layout_type}_search"
        )
        
        display_content = self.validator.md_content
        if search_term:
            lines = display_content.split('\n')
            filtered_lines = [line for line in lines if search_term.lower() in line.lower()]
            display_content = '\n'.join(filtered_lines)
            if filtered_lines:
                st.success(f"找到 {len(filtered_lines)} 行包含 '{search_term}'")
            else:
                st.warning(f"未找到包含 '{search_term}' 的内容")
        
        # 渲染方式选择
        render_mode = st.radio(
            "选择渲染方式",
            ["HTML渲染", "Markdown渲染", "DataFrame表格", "原始文本"],
            horizontal=True,
            key=f"{layout_type}_render_mode"
        )
        
        return display_content, render_mode
    
    def render_content_by_mode(self, content: str, render_mode: str, font_size: int, layout_type: str):
        """根据渲染模式显示内容 - 统一方法"""
        if content is None or render_mode is None:
            return
            
        if render_mode == "HTML渲染":
            content_style = f"""
            <style>
            .{layout_type}-content-display {{
                font-size: {font_size}px !important;
                line-height: 1.4;
                color: #333333 !important;
                background-color: #fafafa !important;
                padding: 10px;
                border-radius: 5px;
                border: 1px solid #ddd;
            }}
            </style>
            """
            st.markdown(content_style, unsafe_allow_html=True)
            st.markdown(f'<div class="{layout_type}-content-display">{content}</div>', unsafe_allow_html=True)
            
        elif render_mode == "Markdown渲染":
            converted_content = convert_html_table_to_markdown(content)
            content_style = f"""
            <style>
            .{layout_type}-content-display {{
                font-size: {font_size}px !important;
                line-height: 1.4;
                color: #333333 !important;
                background-color: #fafafa !important;
                padding: 10px;
                border-radius: 5px;
                border: 1px solid #ddd;
            }}
            </style>
            """
            st.markdown(content_style, unsafe_allow_html=True)
            st.markdown(f'<div class="{layout_type}-content-display">{converted_content}</div>', unsafe_allow_html=True)
            
        elif render_mode == "DataFrame表格":
            if '<table' in content.lower():
                self.validator.display_html_table_as_dataframe(content)
            else:
                st.info("当前内容中没有检测到HTML表格")
                st.markdown(content, unsafe_allow_html=True)
        else:  # 原始文本
            st.text_area(
                "MD内容预览",
                content,
                height=300,
                key=f"{layout_type}_text_area"
            )
    
    # 布局实现
    def create_standard_layout(self, font_size: int = 10, zoom_level: float = 1.0):
        """创建标准布局"""
        if zoom_level is None:
            zoom_level = self.config['styles']['layout']['default_zoom']
            
        # 主要内容区域
        layout = self.config['styles']['layout']
        left_col, right_col = st.columns([layout['content_width'], layout['sidebar_width']])
        
        with left_col:
            self.render_content_section("standard")
            
            # 显示内容
            if self.validator.md_content:
                display_content, render_mode = self.render_md_content("standard")
                self.render_content_by_mode(display_content, render_mode, font_size, "standard")
        
        with right_col:
            self.create_aligned_image_display(zoom_level, "compact")
    
    def create_compact_layout(self, font_size: int = 10, zoom_level: float = 1.0):
        """创建紧凑的对比布局"""
        # 主要内容区域
        layout = self.config['styles']['layout']
        left_col, right_col = st.columns([layout['content_width'], layout['sidebar_width']])                

        with left_col:
            self.render_content_section("compact")

            # 只保留一个内容区域高度选择
            container_height = st.selectbox(
                "选择内容区域高度", 
                [400, 600, 800, 1000, 1200], 
                index=2,
                key="compact_content_height"
            )
            
            # 快速定位文本选择器（使用不同的key）
            if self.validator.text_bbox_mapping:
                text_options = ["请选择文本..."] + list(self.validator.text_bbox_mapping.keys())
                selected_index = st.selectbox(
                    "快速定位文本",
                    range(len(text_options)),
                    format_func=lambda x: text_options[x][:30] + "..." if len(text_options[x]) > 30 else text_options[x],
                    key="compact_quick_text_selector"  # 使用不同的key
                )
                
                if selected_index > 0:
                    st.session_state.selected_text = text_options[selected_index]
            
            # 自定义CSS样式
            st.markdown(f"""
            <style>
            .compact-content {{
                height: {container_height}px;
                overflow-y: auto;
                font-size: {font_size}px !important;
                line-height: 1.4;
                border: 1px solid #ddd;
                padding: 10px;
                background-color: #fafafa !important;
                font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', monospace;
                color: #333333 !important;
            }}
            
            .highlight-text {{
                background-color: #ffeb3b !important;
                padding: 2px 4px;
                border-radius: 3px;
                cursor: pointer;
                color: #333333 !important;
            }}
            
            .selected-highlight {{
                background-color: #4caf50 !important;
                color: white !important;
            }}
            </style>
            """, unsafe_allow_html=True)
            
            # 处理并显示OCR内容
            if self.validator.md_content:
                # 高亮可点击文本
                highlighted_content = self.validator.md_content
                for text in self.validator.text_bbox_mapping.keys():
                    if len(text) > 2:  # 避免高亮过短的文本
                        css_class = "highlight-text selected-highlight" if text == st.session_state.selected_text else "highlight-text"
                        highlighted_content = highlighted_content.replace(
                            text, 
                            f'<span class="{css_class}" title="{text[:50]}...">{text}</span>'
                        )
                st.markdown(
                    f'<div class="compact-content">{highlighted_content}</div>', 
                    unsafe_allow_html=True
                )
        
        with right_col:
            # 修复的对齐图片显示
            self.create_aligned_image_display(zoom_level, "compact")
    
    def create_aligned_image_display(self, zoom_level: float = 1.0, layout_type: str = "aligned"):
        """创建与左侧对齐的图片显示"""
        # 精确对齐CSS
        st.markdown(f"""
        <style>
        .aligned-image-container-{layout_type} {{
            margin-top: -70px;
            padding-top: 0px;
        }}
        .aligned-image-container-{layout_type} h1 {{
            margin-top: 0px !important;
            padding-top: 0px !important;
        }}
        </style>
        """, unsafe_allow_html=True)
        
        st.markdown(f'<div class="aligned-image-container-{layout_type}">', unsafe_allow_html=True)
        st.header("🖼️ 原图标注")
        
        # 图片缩放控制
        col1, col2 = st.columns(2)
        with col1:
            current_zoom = st.slider("图片缩放", 0.3, 2.0, zoom_level, 0.1, key=f"{layout_type}_zoom_level")
        with col2:
            show_all_boxes = st.checkbox("显示所有框", value=False, key=f"{layout_type}_show_all_boxes")
        
        # 使用新的图像加载方法
        image = self.load_and_rotate_image(self.validator.image_path)
        
        if image:
            try:
                # 根据缩放级别调整图片大小
                new_width = int(image.width * current_zoom)
                new_height = int(image.height * current_zoom)
                resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
                
                # 计算选中的bbox - 注意bbox已经是旋转后的坐标
                selected_bbox = None
                if st.session_state.selected_text and st.session_state.selected_text in self.validator.text_bbox_mapping:
                    info = self.validator.text_bbox_mapping[st.session_state.selected_text][0]
                    # bbox已经是旋转后的坐标，只需要应用缩放
                    bbox = info['bbox']
                    selected_bbox = [int(coord * current_zoom) for coord in bbox]
                
                # 创建交互式图片
                fig = self.create_resized_interactive_plot(resized_image, selected_bbox, current_zoom, show_all_boxes)
                st.plotly_chart(fig, use_container_width=True, key=f"{layout_type}_plot")
                
                # 显示选中文本的详细信息
                if st.session_state.selected_text and st.session_state.selected_text in self.validator.text_bbox_mapping:
                    st.subheader("📍 选中文本详情")
                    
                    info = self.validator.text_bbox_mapping[st.session_state.selected_text][0]
                    bbox = info['bbox']
                    
                    info_col1, info_col2 = st.columns(2)
                    with info_col1:
                        st.write(f"**文本内容:** {st.session_state.selected_text[:30]}...")
                        st.write(f"**类别:** {info['category']}")
                        # 显示旋转信息
                        rotation_angle = self.get_rotation_angle()
                        if rotation_angle != 0:
                            st.write(f"**旋转角度:** {rotation_angle}°")
                
                    with info_col2:
                        st.write(f"**位置:** [{', '.join(map(str, bbox))}]")
                        if len(bbox) >= 4:
                            st.write(f"**大小:** {bbox[2] - bbox[0]} x {bbox[3] - bbox[1]} px")
                
                # 错误标记功能
                col1, col2 = st.columns(2)
                with col1:
                    if st.button("❌ 标记为错误", key=f"{layout_type}_mark_error"):
                        st.session_state.marked_errors.add(st.session_state.selected_text)
                        st.rerun()
                
                with col2:
                    if st.button("✅ 取消错误标记", key=f"{layout_type}_unmark_error"):
                        st.session_state.marked_errors.discard(st.session_state.selected_text)
                        st.rerun()
                        
            except Exception as e:
                st.error(f"❌ 图片处理失败: {e}")
                st.error(f"详细错误: {str(e)}")
        else:
            st.error("未找到对应的图片文件")
            if self.validator.image_path:
                st.write(f"期望路径: {self.validator.image_path}")
    
        st.markdown('</div>', unsafe_allow_html=True)
    
    def create_resized_interactive_plot(self, image: Image.Image, selected_bbox: Optional[List[int]], zoom_level: float, show_all_boxes: bool) -> go.Figure:
        """创建可调整大小的交互式图片 - 修复图像显示和bbox对齐问题"""
        fig = go.Figure()
        
        # 添加图片 - 修正图像定位，确保与工具栏距离一致
        fig.add_layout_image(
            dict(
                source=image,
                xref="x", yref="y",
                x=0, y=image.height * zoom_level,  # 修正：图片左上角位置
                sizex=image.width * zoom_level, 
                sizey=image.height * zoom_level,
                sizing="stretch", 
                opacity=1.0, 
                layer="below"
            )
        )
        
        # 显示所有bbox
        if show_all_boxes:
            for text, info_list in self.validator.text_bbox_mapping.items():
                for info in info_list:
                    bbox = info['bbox']
                    if len(bbox) >= 4:
                        # bbox已经是旋转后的坐标，需要应用缩放并转换坐标系
                        x1, y1, x2, y2 = bbox[:4]
                        
                        # 应用缩放
                        scaled_x1 = x1 * zoom_level
                        scaled_y1 = y1 * zoom_level
                        scaled_x2 = x2 * zoom_level
                        scaled_y2 = y2 * zoom_level
                        
                        # 转换为plotly坐标系（原点在左下角）
                        plot_x1 = scaled_x1
                        plot_y1 = (image.height * zoom_level) - scaled_y2  # 翻转Y坐标
                        plot_x2 = scaled_x2
                        plot_y2 = (image.height * zoom_level) - scaled_y1  # 翻转Y坐标
                        
                        color = "rgba(0, 100, 200, 0.2)"
                        if text in self.validator.marked_errors:
                            color = "rgba(255, 0, 0, 0.3)"
                        
                        fig.add_shape(
                            type="rect",
                            x0=plot_x1, y0=plot_y1,
                            x1=plot_x2, y1=plot_y2,
                            line=dict(color=color.replace('0.2', '0.8').replace('0.3', '1.0'), width=1),
                            fillcolor=color,
                        )

        # 高亮显示选中的bbox
        if selected_bbox and len(selected_bbox) >= 4:
            x1, y1, x2, y2 = selected_bbox[:4]
            
            # 转换为plotly坐标系（selected_bbox已经是缩放后的坐标）
            plot_x1 = x1
            plot_y1 = (image.height * zoom_level) - y2  # 翻转Y坐标
            plot_x2 = x2
            plot_y2 = (image.height * zoom_level) - y1  # 翻转Y坐标
            
            fig.add_shape(
                type="rect",
                x0=plot_x1, y0=plot_y1,
                x1=plot_x2, y1=plot_y2,
                line=dict(color="red", width=3),
                fillcolor="rgba(255, 0, 0, 0.3)",
            )
        
        # 计算合适的显示尺寸
        aspect_ratio = image.width / image.height
        display_height = min(800, max(400, image.height // 2))
        display_width = int(display_height * aspect_ratio)
        
        # 设置布局 - 确保图像完全可见，使用缩放后的尺寸
        fig.update_layout(
            width=display_width,
            height=display_height,
            margin=dict(l=0, r=0, t=0, b=0),
            showlegend=False,
            plot_bgcolor='white',
            dragmode="pan",
            
            # X轴设置 - 使用缩放后的图像尺寸
            xaxis=dict(
                visible=False,
                range=[0, image.width * zoom_level],
                constrain="domain",
                fixedrange=False,
                autorange=False
            ),
            
            # Y轴设置 - plotly坐标系（原点在左下角）
            yaxis=dict(
                visible=False,
                range=[0, image.height * zoom_level],
                constrain="domain",
                scaleanchor="x",
                scaleratio=1,
                fixedrange=False,
                autorange=False
            )
        )
        
        return fig