ocr_platform/ocr_validator/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 typing import Tuple
import re
import html

from ocr_validator_utils import (
    rotate_image_and_coordinates,
    get_ocr_tool_rotation_config,
)

# 添加 ocr_platform 根目录到 Python 路径（用于导入 ocr_utils）
# 使用 resolve() 确保路径是绝对路径，避免相对路径导致的 IndexError
import sys
_file_path = Path(__file__).resolve()
ocr_platform_root = _file_path.parents[1]  # ocr_validator_layout.py -> ocr_validator -> ocr_platform
if str(ocr_platform_root) not in sys.path:
    sys.path.insert(0, str(ocr_platform_root))

# 从 ocr_utils 导入通用工具
from ocr_utils.html_utils import convert_html_table_to_markdown, parse_html_tables
from ocr_utils.visualization_utils import VisualizationUtils

# BeautifulSoup用于精确HTML表格处理
from bs4 import BeautifulSoup

# 从本地文件导入 Streamlit 特定函数
from ocr_validator_file_utils import load_css_styles

# 为了向后兼容，提供函数别名
draw_bbox_on_image = VisualizationUtils.draw_bbox_on_image

# detect_image_orientation_by_opencv 保留在 ocr_validator_file_utils
from ocr_validator_file_utils import detect_image_orientation_by_opencv

class OCRLayoutManager:
    """OCR布局管理器"""
    
    def __init__(self, validator):
        self.validator = validator
        self.config = validator.config
        self._rotated_image_cache = {}
        self._cache_max_size = 10
        self._orientation_cache = {}  # 缓存方向检测结果
        self.rotated_angle = 0.0  # 自动检测的旋转角度缓存
        self.show_all_boxes = False
        self.fit_to_container = False
        self.zoom_level = 1.0

    def _clear_selection_callback(self, layout_type: str):
        """在按钮回调中清理选择/搜索状态（避免 widget 实例化后修改报错）"""
        # 业务态
        st.session_state.selected_text = ""
        # 紧凑布局的搜索态
        st.session_state.compact_search_query = ""

        # widget key 对应的状态（必须在 on_click 回调里改）
        search_key = f"{layout_type}_search_input"
        quick_select_key = f"{layout_type}_quick_text_selector"
        st.session_state[search_key] = ""
        st.session_state[quick_select_key] = 0
    
    def _highlight_text_safely(self, content: str, text_to_highlight: str, 
                               highlight_class: str, title: Optional[str] = None) -> str:
        """
        安全地高亮文本，保护Markdown语法（特别是图片）
        
        策略：
        1. 保护特殊内容（HTML注释、Markdown图片）
        2. 只对HTML表格使用BeautifulSoup精确处理
        3. 其他部分使用简单字符串替换，保持Markdown格式
        
        Args:
            content: 要处理的Markdown/HTML混合内容
            text_to_highlight: 要高亮的文本
            highlight_class: 高亮样式类名
            title: 鼠标悬停提示文本
        
        Returns:
            处理后的内容
        """
        if not text_to_highlight or text_to_highlight not in content:
            return content
        
        if title is None:
            title = text_to_highlight
        
        try:
            import re
            
            # 1. 提取并保护特殊内容
            protected_parts = []
            
            # 保护 HTML 注释
            def protect_comment(match):
                protected_parts.append(match.group(0))
                return f"__PROTECTED_{len(protected_parts) - 1}__"
            
            content = re.sub(r'<!--.*?-->', protect_comment, content, flags=re.DOTALL)
            
            # 保护 Markdown 图片（完整语法）
            def protect_image(match):
                protected_parts.append(match.group(0))
                return f"__PROTECTED_{len(protected_parts) - 1}__"
            
            content = re.sub(r'!\[.*?\]\([^)]+\)', protect_image, content)
            
            # 2. 提取表格并单独处理
            tables = []
            def extract_table(match):
                tables.append(match.group(0))
                return f"__TABLE_{len(tables) - 1}__"
            
            content = re.sub(r'<table[^>]*>.*?</table>', extract_table, content, flags=re.DOTALL)
            
            # 3. 对表格使用 BeautifulSoup 精确处理（只高亮文本，不高亮整个单元格）
            highlighted_tables = []
            
            for table_html in tables:
                soup = BeautifulSoup(table_html, 'html.parser')
                
                # 在表格单元格中查找完全匹配
                for td in soup.find_all(['td', 'th']):
                    cell_text = td.get_text(strip=True)
                    if cell_text == text_to_highlight:
                        # 🎯 只高亮文本，不高亮整个单元格
                        # 清空单元格内容
                        td.clear()
                        # 创建高亮 span 包裹文本
                        span = soup.new_tag('span')
                        span['class'] = highlight_class.split()
                        if title:
                            span['title'] = title
                        span.string = text_to_highlight
                        # 将 span 添加到单元格
                        td.append(span)
                
                highlighted_tables.append(str(soup))
            
            # 4. 对普通文本进行简单替换（保持Markdown格式，跳过占位符）
            if text_to_highlight in content:
                highlight_span = f'<span class="{highlight_class}"'
                if title:
                    highlight_span += f' title="{title}"'
                highlight_span += f'>{text_to_highlight}</span>'
                
                # 🎯 安全替换：使用正则表达式，排除占位符内的匹配
                # 负向前瞻：确保前面不是占位符的一部分
                pattern = f'(?<!__PROTECTED_)(?<!__TABLE_){re.escape(text_to_highlight)}(?!__)'
                content = re.sub(pattern, highlight_span, content)
            
            # 5. 恢复表格
            for i, table in enumerate(highlighted_tables):
                content = content.replace(f"__TABLE_{i}__", table)
            
            # 6. 恢复受保护的内容（图片和注释）
            for i, protected in enumerate(protected_parts):
                content = content.replace(f"__PROTECTED_{i}__", protected)
            
            return content
            
        except Exception as e:
            st.warning(f"文本高亮时出错: {str(e)}")
            return content
            
    
    def clear_image_cache(self):
        """清理所有图像缓存"""
        self._rotated_image_cache.clear()
        
    def clear_cache_for_image(self, image_path: str):
        """清理指定图像的所有缓存"""
        keys_to_remove = [key for key in self._rotated_image_cache.keys() if key.startswith(image_path)]
        for key in keys_to_remove:
            del self._rotated_image_cache[key]
    
    def get_cache_info(self) -> dict:
        """获取缓存信息"""
        return {
            'cache_size': len(self._rotated_image_cache),
            'cached_images': list(self._rotated_image_cache.keys()),
            'max_size': self._cache_max_size
        }
    
    def _manage_cache_size(self):
        """管理缓存大小，超出限制时清理最旧的缓存"""
        if len(self._rotated_image_cache) > self._cache_max_size:
            # 删除最旧的缓存项（FIFO策略）
            oldest_key = next(iter(self._rotated_image_cache))
            del self._rotated_image_cache[oldest_key]
    
    def detect_and_suggest_rotation(self, image_path: str) -> Dict:
        """检测并建议图片旋转角度"""
        if image_path in self._orientation_cache:
            return self._orientation_cache[image_path]
        
        # 使用自动检测功能
        detection_result = detect_image_orientation_by_opencv(image_path)
        
        # 缓存结果
        self._orientation_cache[image_path] = detection_result
        return detection_result
    
    def get_rotation_angle(self) -> float:
        """获取旋转角度 - 增强版本支持自动检测"""
        # 如果没有预设角度，优先人工设置
        if hasattr(self, 'rotated_angle') and self.rotated_angle != 0:
            return self.rotated_angle

        # 尝试从OCR数据中获取（PPStructV3等）
        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:
            self.validator.text_bbox_mapping = self._rotated_image_cache[cache_key]['text_bbox_mapping']
            return self._rotated_image_cache[cache_key]['image']
        
        try:
            image = Image.open(image_path)
            
            # 如果需要旋转
            if rotation_angle != 0:
                # 获取OCR工具的旋转配置
                rotation_config = get_ocr_tool_rotation_config(self.validator.ocr_data, self.config)
                
                # st.info(f"🔄 检测到文档旋转角度: {rotation_angle}°，正在处理图像和坐标...")
                # st.info(f"📋 OCR工具配置: 坐标{'已预旋转' if rotation_config['coordinates_are_pre_rotated'] else '需要旋转'}")
                
                # 判断是否需要旋转坐标
                if rotation_config['coordinates_are_pre_rotated']:
                    # 图片的角度与坐标的角度不一致，比如PPStructV3，图片0度，坐标已旋转270度
                    # 这种情况下，只需要旋转图片，坐标不变
                    # PPStructV3: 坐标已经是旋转后的，只旋转图像
                    img_rotation_angle = (rotation_angle + self.rotated_angle) % 360
                    if img_rotation_angle == 270:
                        rotated_image = image.rotate(-90, expand=True)  # 顺时针90度
                    elif img_rotation_angle == 90:
                        rotated_image = image.rotate(90, expand=True)   # 逆时针90度
                    elif img_rotation_angle == 180:
                        rotated_image = image.rotate(180, expand=True)  # 180度
                    else:
                        rotated_image = image.rotate(-img_rotation_angle, expand=True)
                    
                    if self.rotated_angle == 0:
                        # 坐标不需要变换，因为JSON中已经是正确的坐标
                        self._rotated_image_cache[cache_key] = {'image': rotated_image, 'text_bbox_mapping': self.validator.text_bbox_mapping}
                        self._manage_cache_size()
                        return rotated_image

                    image = rotated_image  # 继续使用旋转后的图像进行后续处理
                
                # VLM: 需要同时旋转图像和坐标
                # 收集所有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, 
                    rotate_coordinates=not rotation_config['coordinates_are_pre_rotated']
                )
                
                # 更新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] = {'image': rotated_image, 'text_bbox_mapping': self.validator.text_bbox_mapping}
                self._manage_cache_size()
                return rotated_image
                    
            else:
                # 无需旋转，直接缓存原图
                self._rotated_image_cache[cache_key] = {'image': image, 'text_bbox_mapping': self.validator.text_bbox_mapping}
                self._manage_cache_size()  # 检查并管理缓存大小
                return image
                
        except Exception as e:
            st.error(f"❌ 图像加载失败: {e}")
            return None

    def render_content_by_mode(self, content: str, render_mode: str, font_size: int, 
                          container_height: int, layout_type: str, 
                          highlight_config: Optional[Dict] = None):
        """
        根据渲染模式显示内容 - 增强版本
        
        Args:
            content: 要渲染的内容
            render_mode: 渲染模式
            font_size: 字体大小
            container_height: 容器高度
            layout_type: 布局类型
            highlight_config: 高亮配置 {'has_bbox': bool, 'match_type': str}
        """
        if content is None or render_mode is None:
            return
            
        if render_mode == "HTML渲染":
            # 🎯 构建样式 - 包含基础样式和高亮样式
            content_style = f"""
            <style>
            /* ========== 基础容器样式 ========== */
            .{layout_type}-content-display {{
                height: {container_height}px;
                overflow-x: auto;
                overflow-y: auto;
                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;
                max-width: 100%;
            }}
            
            /* ========== 表格样式 ========== */
            .{layout_type}-content-display table {{
                width: 100%;
                border-collapse: collapse;
                margin: 10px 0;
                white-space: nowrap;
            }}
            
            .{layout_type}-content-display th,
            .{layout_type}-content-display td {{
                border: 1px solid #ddd;
                padding: 8px;
                text-align: left;
                max-width: 300px;
                word-wrap: break-word;
                word-break: break-all;
                vertical-align: top;
            }}
            
            .{layout_type}-content-display th {{
                background-color: #f5f5f5;
                position: sticky;
                top: 0;
                z-index: 1;
                font-weight: bold;
            }}
            
            /* 数字列右对齐 */
            .{layout_type}-content-display td.number {{
                text-align: right;
                white-space: nowrap;
                font-family: 'Monaco', 'Menlo', monospace;
            }}
            
            /* 短文本列不换行 */
            .{layout_type}-content-display td.short-text {{
                white-space: nowrap;
                min-width: 80px;
            }}
            
            /* ========== 图片样式 ========== */
            .{layout_type}-content-display img {{
                max-width: 100%;
                height: auto;
                border-radius: 4px;
                margin: 10px 0;
            }}
            
            /* ========== 响应式设计 ========== */
            @media (max-width: 768px) {{
                .{layout_type}-content-display table {{
                    font-size: {max(font_size-2, 8)}px;
                }}
                .{layout_type}-content-display th,
                .{layout_type}-content-display td {{
                    padding: 4px;
                    max-width: 150px;
                }}
            }}
            
            /* ========== 高亮文本样式 ========== */
            .{layout_type}-content-display .highlight-text {{
                padding: 2px 4px;
                border-radius: 3px;
                cursor: pointer;
                font-weight: 500;
                transition: all 0.2s ease;
            }}
            
            .{layout_type}-content-display .highlight-text:hover {{
                opacity: 0.8;
                transform: scale(1.02);
            }}
            
            /* 🎯 精确匹配且有框 - 绿色 */
            .{layout_type}-content-display .highlight-text.selected-highlight {{
                background-color: #4caf50 !important;
                color: white !important;
                border: 1px solid #2e7d32 !important;
            }}
            
            /* 🎯 OCR匹配 - 蓝色 */
            .{layout_type}-content-display .highlight-text.ocr-match {{
                background-color: #2196f3 !important;
                color: white !important;
                border: 1px solid #1565c0 !important;
            }}
            
            /* 🎯 无边界框 - 橙色虚线 */
            .{layout_type}-content-display .highlight-text.no-bbox {{
                background-color: #ff9800 !important;
                color: white !important;
                border: 1px dashed #f57c00 !important;
            }}
            
            /* 🎯 默认高亮 - 黄色 */
            .{layout_type}-content-display .highlight-text.default {{
                background-color: #ffeb3b !important;
                color: #333333 !important;
                border: 1px solid #fbc02d !important;
            }}
            </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)
            st.markdown(converted_content, 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_compact_layout(self, config: Dict):
        """创建紧凑的对比布局 - 增强搜索功能"""
        layout = config['styles']['layout']
        font_size = config['styles'].get('font_size', 10)
        container_height = layout.get('default_height', 600)
        zoom_level = layout.get('default_zoom', 1.0)
        layout_type = "compact"

        left_col, right_col = st.columns([layout['content_width'], layout['sidebar_width']], 
                                         vertical_alignment='top', border=True)

        with left_col:
            if self.validator.text_bbox_mapping:
                # 搜索输入框
                search_col, select_col = st.columns([1, 2])
                
                if "compact_search_query" not in st.session_state:
                    st.session_state.compact_search_query = ""
                
                with search_col:
                    search_query = st.text_input(
                        "搜索文本",
                        placeholder="输入关键词...",
                        value=st.session_state.compact_search_query,
                        key=f"{layout_type}_search_input",
                        label_visibility="collapsed"
                    )
                    st.session_state.compact_search_query = search_query
                
                # 🎯 增强搜索逻辑：构建选项列表
                text_options = ["请选择文本..."]
                text_display = ["请选择文本..."]
                match_info = [None]  # 记录匹配信息
                
                for text, info_list in self.validator.text_bbox_mapping.items():
                    # 🔑 关键改进：同时搜索 text 和 matched_text
                    if search_query and search_query.strip():
                        query_lower = search_query.lower()
                        
                        # 1. 检查原始文本
                        text_match = query_lower in text.lower()
                        
                        # 2. 检查 matched_text（OCR识别文本）
                        matched_text_match = False
                        matched_text = None
                        if info_list and isinstance(info_list[0], dict):
                            matched_text = info_list[0].get('matched_text', '')
                            matched_text_match = query_lower in matched_text.lower() if matched_text else False
                        
                        # 如果都不匹配，跳过
                        if not text_match and not matched_text_match:
                            continue
                        
                        # 记录匹配类型
                        if text_match:
                            match_type = "exact"
                            match_source = text
                        else:
                            match_type = "ocr"
                            match_source = matched_text
                    else:
                        match_type = None
                        match_source = text
                    
                    text_options.append(text)
                    
                    # 🎯 构建显示文本（带匹配提示）
                    if info_list and isinstance(info_list[0], dict):
                        first_info = info_list[0]
                        
                        # 检查是否有 bbox
                        has_bbox = 'bbox' in first_info and first_info['bbox']
                        
                        # 表格单元格显示
                        if 'row' in first_info and 'col' in first_info:
                            display_text = f"[R{first_info['row']},C{first_info['col']}] {text}"
                        else:
                            display_text = text
                        
                        # 🎯 添加匹配提示
                        if match_type == "ocr":
                            display_text = f"🔍 {display_text} (OCR: {match_source[:20]}...)"
                        elif not has_bbox:
                            display_text = f"⚠️ {display_text} (无框)"
                        
                        # 截断过长文本
                        if len(display_text) > 60:
                            display_text = display_text[:57] + "..."
                    else:
                        display_text = text[:57] + "..." if len(text) > 60 else text
                    
                    text_display.append(display_text)
                    match_info.append({
                        'type': match_type,
                        'source': match_source,
                        'has_bbox': has_bbox if info_list else False
                    })
                
                # 🎯 显示搜索统计
                if search_query and search_query.strip():
                    ocr_matches = sum(1 for m in match_info[1:] if m and m['type'] == 'ocr')
                    no_bbox_count = sum(1 for m in match_info[1:] if m and not m['has_bbox'])
                    
                    stat_parts = [f"找到 {len(text_options)-1} 个匹配项"]
                    if ocr_matches > 0:
                        stat_parts.append(f"🔍 {ocr_matches} 个OCR匹配")
                    if no_bbox_count > 0:
                        stat_parts.append(f"⚠️ {no_bbox_count} 个无框")
                    
                    st.caption(" | ".join(stat_parts))
                
                # 确定默认选中的索引
                default_index = 0
                if st.session_state.selected_text and st.session_state.selected_text in text_options:
                    default_index = text_options.index(st.session_state.selected_text)
                
                with select_col:
                    selected_index = st.selectbox(
                        "快速定位文本",
                        range(len(text_options)),
                        index=default_index,
                        format_func=lambda x: text_display[x] if x < len(text_display) else "",
                        label_visibility="collapsed",
                        key=f"{layout_type}_quick_text_selector"
                    )
                
                # 🎯 显示匹配详情
                if selected_index > 0:
                    st.session_state.selected_text = text_options[selected_index]
                    
                    # 获取匹配信息
                    selected_match_info = match_info[selected_index]
                    if selected_match_info:
                        if selected_match_info['type'] == 'ocr':
                            st.info(f"🔍 **OCR识别文本匹配**: `{selected_match_info['source']}`")
                        elif not selected_match_info['has_bbox']:
                            st.warning(f"⚠️ **未找到边界框**: 文本在MD中存在，但没有对应的坐标信息")
            
            # 🎯 增强高亮显示逻辑
            if self.validator.md_content:
                highlighted_content = self.validator.md_content
                
                if st.session_state.selected_text:
                    selected_text = st.session_state.selected_text
                    
                    # 获取匹配信息
                    info_list = self.validator.text_bbox_mapping.get(selected_text, [])
                    has_bbox = False
                    matched_text = None
                    match_type = None
                    
                    if info_list and isinstance(info_list[0], dict):
                        has_bbox = 'bbox' in info_list[0] and info_list[0]['bbox']
                        matched_text = info_list[0].get('matched_text', '')
                        
                        # 🔑 判断匹配类型
                        if matched_text and matched_text != selected_text:
                            match_type = "ocr"
                        elif has_bbox:
                            match_type = "exact"
                        else:
                            match_type = "no_bbox"
                    
                    # 🎯 应用高亮
                    if len(selected_text) >= self.config.get('ocr', {}).get('min_text_length', 2):
                        # 1. 高亮原始文本
                        if selected_text in highlighted_content:
                            if match_type == "exact":
                                highlight_class = "highlight-text selected-highlight"
                            elif match_type == "no_bbox":
                                highlight_class = "highlight-text no-bbox"
                            else:
                                highlight_class = "highlight-text default"
                            
                            # 使用正则表达式避免替换base64编码中的内容
                            highlighted_content = self._highlight_text_safely(
                                highlighted_content,
                                selected_text,
                                highlight_class
                            )
                        
                        # 2. 如果有 matched_text 且不同，也高亮
                        if matched_text and matched_text != selected_text and matched_text in highlighted_content:
                            # 使用正则表达式避免替换base64编码中的内容
                            highlighted_content = self._highlight_text_safely(
                                highlighted_content,
                                matched_text,
                                "highlight-text ocr-match",
                                f"OCR: {matched_text}"
                            )
                
                # 🎯 调用渲染方法（样式已内置）
                self.render_content_by_mode(
                    highlighted_content, 
                    "HTML渲染", 
                    font_size, 
                    container_height, 
                    layout_type
                )
    
        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"):
        """创建响应式图片显示"""
    
        # st.header("🖼️ 原图标注")
        
        # 图片控制选项
        col1, col2, col3, col4, col5 = st.columns(5, vertical_alignment="center", border= False)

        with col1:
            # 判断{layout_type}_show_all_boxes是否有值，如果有值直接使用，否则默认False
            # if f"{layout_type}_show_all_boxes" not in st.session_state:
            #     st.session_state[f"{layout_type}_show_all_boxes"] = False

            show_all_boxes = st.checkbox(
                "显示所有框",
                # value=st.session_state[f"{layout_type}_show_all_boxes"],
                value = self.show_all_boxes,
                key=f"{layout_type}_show_all_boxes"
            )
            if show_all_boxes != self.show_all_boxes:
                self.show_all_boxes = show_all_boxes

        with col2:
            if st.button("🔄 旋转90度", type="secondary", key=f"{layout_type}_manual_angle"):
                self.rotated_angle = (self.rotated_angle + 90) % 360
                # 需要清除图片缓存，以及text_bbox_mapping中的bbox
                self.clear_image_cache()
                self.validator.process_data()
                st.rerun()
                
        with col3:
            # 显示当前角度状态
            current_angle = self.get_rotation_angle()
            st.metric("当前角度", f"{current_angle}°", label_visibility="collapsed")

        with col4:
            if st.button("↺ 重置角度", key=f"{layout_type}_reset_angle"):
                self.rotated_angle = 0.0
                st.success("已重置旋转角度")
                # 需要清除图片缓存，以及text_bbox_mapping中的bbox
                self.clear_image_cache()
                self.validator.process_data()
                st.rerun()
 
        with col5:
            st.button(
                "🧹 清除选择",
                key=f"{layout_type}_clear_selection",
                on_click=self._clear_selection_callback,
                kwargs={"layout_type": layout_type},
            )

        # 使用增强的图像加载方法
        image = self.load_and_rotate_image(self.validator.image_path)
        
        if image:
            try:
                resized_image, all_boxes, selected_boxes = self.zoom_image(image, self.zoom_level)
                
                # 创建交互式图片
                fig = self.create_resized_interactive_plot(resized_image, selected_boxes, self.zoom_level, all_boxes)

                plot_config = {
                    'displayModeBar': True,
                    'modeBarButtonsToRemove': ['zoom2d', 'select2d', 'lasso2d', 'autoScale2d'],
                    'scrollZoom': True,
                    'doubleClick': 'reset',
                    'responsive': False,  # 关键：禁用响应式，使用固定尺寸
                    'toImageButtonOptions': {
                        'format': 'png',
                        'filename': 'ocr_image',
                        'height': None,  # 使用当前高度
                        'width': None,   # 使用当前宽度
                        'scale': 1
                    }
                }
                
                # 🔧 修复：使用 use_container_width 替代废弃的参数
                st.plotly_chart(
                    fig, 
                    width='stretch',  # 🎯 使用容器宽度
                    config=plot_config,
                    key=f"{layout_type}_plot"
                )
                        
            except Exception as e:
                st.error(f"❌ 图片处理失败: {e}")
                st.exception(e)
        else:
            st.error("未找到对应的图片文件")
            if self.validator.image_path:
                st.write(f"期望路径: {self.validator.image_path}")

    # st.markdown('</div>', unsafe_allow_html=True)

    def zoom_image(self, image: Image.Image, current_zoom: float) -> Tuple[Image.Image, List[List[int]], List[List[int]]]:
        """缩放图像"""
        # 根据缩放级别调整图片大小
        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
        selected_boxes = []
        if st.session_state.selected_text and st.session_state.selected_text in self.validator.text_bbox_mapping:
            info_list = self.validator.text_bbox_mapping[st.session_state.selected_text]
            for info in info_list:
                if 'bbox' in info:
                    bbox = info['bbox']
                    selected_box = [int(coord * current_zoom) for coord in bbox]
                    selected_boxes.append(selected_box)

        # 收集所有框
        all_boxes = []
        if self.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:
                        scaled_bbox = [coord * current_zoom for coord in bbox]
                        all_boxes.append(scaled_bbox)

        return resized_image, all_boxes, selected_boxes

    def _add_bboxes_to_plot_batch(self, fig: go.Figure, bboxes: List[List[int]], 
                                image_height: int, 
                                line_color: str = "blue", 
                                line_width: int = 2, 
                                fill_color: str = "rgba(0, 100, 200, 0.2)"):
        """
        批量添加边界框（性能优化版）
        """
        if not bboxes or len(bboxes) == 0:
            return
        
        # 🎯 关键优化：构建 shapes 列表，一次性添加
        shapes = []
        for bbox in bboxes:
            if len(bbox) < 4:
                continue
            
            x1, y1, x2, y2 = bbox[:4]
            
            # 转换坐标
            plot_x1 = x1
            plot_x2 = x2
            plot_y1 = image_height - y2
            plot_y2 = image_height - y1
            
            shapes.append(dict(
                type="rect",
                x0=plot_x1, y0=plot_y1,
                x1=plot_x2, y1=plot_y2,
                line=dict(color=line_color, width=line_width),
                fillcolor=fill_color,
            ))
        
        # 🎯 一次性更新所有形状
        fig.update_layout(shapes=fig.layout.shapes + tuple(shapes))

    def _add_bboxes_as_scatter(self, fig: go.Figure, bboxes: List[List[int]], 
                          image_height: int,
                          line_color: str = "blue", 
                          line_width: int = 2,
                          name: str = "boxes"):
        """
        使用 Scatter 绘制边界框（极致性能优化）
        """
        if not bboxes or len(bboxes) == 0:
            return
        
        # 🎯 收集所有矩形的边框线坐标
        x_coords = []
        y_coords = []
        
        for bbox in bboxes:
            if len(bbox) < 4:
                continue
            
            x1, y1, x2, y2 = bbox[:4]
            
            # 转换坐标
            plot_y1 = image_height - y2
            plot_y2 = image_height - y1
            
            # 绘制矩形：5个点（闭合）
            x_coords.extend([x1, x2, x2, x1, x1, None])  # None用于断开线段
            y_coords.extend([plot_y1, plot_y1, plot_y2, plot_y2, plot_y1, None])
        
        # 🎯 一次性添加所有边框
        fig.add_trace(go.Scatter(
            x=x_coords,
            y=y_coords,
            mode='lines',
            line=dict(color=line_color, width=line_width),
            name=name,
            showlegend=False,
            hoverinfo='skip'
        ))

    def create_resized_interactive_plot(self, image: Image.Image, selected_boxes: List[List[int]], 
                                       zoom_level: float, all_boxes: List[List[int]]) -> go.Figure:
        """创建可调整大小的交互式图片 - 修复容器溢出问题"""
        fig = go.Figure()
        
        # 添加图片 - Plotly坐标系，原点在左下角
        fig.add_layout_image(
            dict(
                source=image,
                xref="x", yref="y",
                x=0, y=image.height,  # 图片左下角在Plotly坐标系中的位置
                sizex=image.width, 
                sizey=image.height,
                sizing="stretch", 
                opacity=1.0, 
                layer="below",
                yanchor="top"  # 确保图片顶部对齐
            )
        )
        
        # 显示所有bbox（淡蓝色）
        if all_boxes:
            self._add_bboxes_as_scatter(
                fig=fig,
                bboxes=all_boxes,
                image_height=image.height,
                line_color="rgba(0, 100, 200, 0.8)",
                line_width=2,
                name="all_boxes"
            )

        # 高亮显示选中的bbox（红色）
        if selected_boxes:
            self._add_bboxes_to_plot_batch(
                fig=fig,
                bboxes=selected_boxes,
                image_height=image.height,
                line_color="red",
                line_width=2,
                fill_color="rgba(255, 0, 0, 0.3)"
            )
    
        # 修复：优化显示尺寸计算
        max_display_width = 1500
        max_display_height = 1000
        
        # 计算合适的显示尺寸，保持宽高比
        aspect_ratio = image.width / image.height
        
        if self.fit_to_container:
            # 自适应容器模式
            if aspect_ratio > 1:  # 宽图
                display_width = min(max_display_width, image.width)
                display_height = int(display_width / aspect_ratio)
            else:  # 高图
                display_height = min(max_display_height, image.height)
                display_width = int(display_height * aspect_ratio)
            
            # 确保不会太小
            display_width = max(display_width, 800)
            display_height = max(display_height, 600)
        else:
            # 固定尺寸模式，但仍要考虑容器限制
            display_width = min(image.width, max_display_width)
            display_height = min(image.height, max_display_height)
        
        # 设置布局 - 关键修改
        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",
            
            # 关键：让图表自适应容器
            # autosize=True,  # 启用自动调整大小
            
            xaxis=dict(
                visible=False,
                range=[0, image.width],
                constrain="domain",
                fixedrange=False,
                autorange=False,
                showgrid=False,
                zeroline=False,
            ),
            
            # 修复：Y轴设置，确保范围正确
            yaxis=dict(
                visible=False,
                range=[0, image.height],  # 确保Y轴范围从0到图片高度
                constrain="domain",
                scaleanchor="x",
                scaleratio=1,
                fixedrange=False,
                autorange=False,
                showgrid=False,
                zeroline=False
            )
        )
        
        return fig