|
|
@@ -0,0 +1,856 @@
|
|
|
+#!/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
|
|
|
+
|
|
|
+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
|
|
|
+
|
|
|
+# 从本地文件导入 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_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) > 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"
|
|
|
+
|
|
|
+ highlighted_content = highlighted_content.replace(
|
|
|
+ selected_text,
|
|
|
+ f'<span class="{highlight_class}" title="{selected_text}">{selected_text}</span>'
|
|
|
+ )
|
|
|
+
|
|
|
+ # 2. 如果有 matched_text 且不同,也高亮
|
|
|
+ if matched_text and matched_text != selected_text and matched_text in highlighted_content:
|
|
|
+ highlighted_content = highlighted_content.replace(
|
|
|
+ matched_text,
|
|
|
+ f'<span class="highlight-text ocr-match" title="OCR: {matched_text}">{matched_text}</span>'
|
|
|
+ )
|
|
|
+
|
|
|
+ # 🎯 调用渲染方法(样式已内置)
|
|
|
+ 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:
|
|
|
+ if st.button("🧹 清除选择", key=f"{layout_type}_clear_selection"):
|
|
|
+ # 清除选中的文本
|
|
|
+ st.session_state.selected_text = None
|
|
|
+ # 清除搜索框内容
|
|
|
+ st.session_state.compact_search_query = None
|
|
|
+ st.rerun()
|
|
|
+
|
|
|
+ # 使用增强的图像加载方法
|
|
|
+ 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
|