| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845 |
- #!/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,
- )
- from ocr_validator_file_utils import (
- convert_html_table_to_markdown,
- parse_html_tables,
- draw_bbox_on_image,
- 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']
-
- # 如果没有预设角度,尝试自动检测
- if hasattr(self, 'rotated_angle'):
- return self.rotated_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:
- # 获取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] = rotated_image
- self._manage_cache_size()
- return rotated_image
- image = rotated_image # 继续使用旋转后的图像进行后续处理
-
- # Dots OCR: 需要同时旋转图像和坐标
- # 收集所有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] = rotated_image
- self._manage_cache_size()
- return rotated_image
-
- else:
- # 无需旋转,直接缓存原图
- self._rotated_image_cache[cache_key] = image
- 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,
- use_container_width=True, # 🎯 使用容器宽度
- 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 = 1,
- 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 = 1,
- 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=1,
- 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=1,
- 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
|