ocr_validator_layout.py 35 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842
  1. #!/usr/bin/env python3
  2. """
  3. OCR验证工具的布局管理模块
  4. 包含标准布局、滚动布局、紧凑布局的实现
  5. """
  6. import streamlit as st
  7. from pathlib import Path
  8. from PIL import Image
  9. from typing import Dict, List, Optional
  10. import plotly.graph_objects as go
  11. from typing import Tuple
  12. from ocr_validator_utils import (
  13. rotate_image_and_coordinates,
  14. get_ocr_tool_rotation_config,
  15. )
  16. from ocr_validator_file_utils import (
  17. convert_html_table_to_markdown,
  18. parse_html_tables,
  19. draw_bbox_on_image,
  20. detect_image_orientation_by_opencv # 新增导入
  21. )
  22. class OCRLayoutManager:
  23. """OCR布局管理器"""
  24. def __init__(self, validator):
  25. self.validator = validator
  26. self.config = validator.config
  27. self._rotated_image_cache = {}
  28. self._cache_max_size = 10
  29. self._orientation_cache = {} # 缓存方向检测结果
  30. self.rotated_angle = 0.0 # 自动检测的旋转角度缓存
  31. self.show_all_boxes = False
  32. self.fit_to_container = False
  33. self.zoom_level = 1.0
  34. def clear_image_cache(self):
  35. """清理所有图像缓存"""
  36. self._rotated_image_cache.clear()
  37. def clear_cache_for_image(self, image_path: str):
  38. """清理指定图像的所有缓存"""
  39. keys_to_remove = [key for key in self._rotated_image_cache.keys() if key.startswith(image_path)]
  40. for key in keys_to_remove:
  41. del self._rotated_image_cache[key]
  42. def get_cache_info(self) -> dict:
  43. """获取缓存信息"""
  44. return {
  45. 'cache_size': len(self._rotated_image_cache),
  46. 'cached_images': list(self._rotated_image_cache.keys()),
  47. 'max_size': self._cache_max_size
  48. }
  49. def _manage_cache_size(self):
  50. """管理缓存大小,超出限制时清理最旧的缓存"""
  51. if len(self._rotated_image_cache) > self._cache_max_size:
  52. # 删除最旧的缓存项(FIFO策略)
  53. oldest_key = next(iter(self._rotated_image_cache))
  54. del self._rotated_image_cache[oldest_key]
  55. def detect_and_suggest_rotation(self, image_path: str) -> Dict:
  56. """检测并建议图片旋转角度"""
  57. if image_path in self._orientation_cache:
  58. return self._orientation_cache[image_path]
  59. # 使用自动检测功能
  60. detection_result = detect_image_orientation_by_opencv(image_path)
  61. # 缓存结果
  62. self._orientation_cache[image_path] = detection_result
  63. return detection_result
  64. def get_rotation_angle(self) -> float:
  65. """获取旋转角度 - 增强版本支持自动检测"""
  66. # 如果没有预设角度,优先人工设置
  67. if hasattr(self, 'rotated_angle') and self.rotated_angle != 0:
  68. return self.rotated_angle
  69. # 尝试从OCR数据中获取(PPStructV3等)
  70. if self.validator.ocr_data:
  71. for item in self.validator.ocr_data:
  72. if isinstance(item, dict) and 'rotation_angle' in item:
  73. return item['rotation_angle']
  74. return 0.0
  75. def load_and_rotate_image(self, image_path: str) -> Optional[Image.Image]:
  76. """加载并根据需要旋转图像"""
  77. if not image_path or not Path(image_path).exists():
  78. return None
  79. # 检查缓存
  80. rotation_angle = self.get_rotation_angle()
  81. cache_key = f"{image_path}_{rotation_angle}"
  82. if cache_key in self._rotated_image_cache:
  83. self.validator.text_bbox_mapping = self._rotated_image_cache[cache_key]['text_bbox_mapping']
  84. return self._rotated_image_cache[cache_key]['image']
  85. try:
  86. image = Image.open(image_path)
  87. # 如果需要旋转
  88. if rotation_angle != 0:
  89. # 获取OCR工具的旋转配置
  90. rotation_config = get_ocr_tool_rotation_config(self.validator.ocr_data, self.config)
  91. # st.info(f"🔄 检测到文档旋转角度: {rotation_angle}°,正在处理图像和坐标...")
  92. # st.info(f"📋 OCR工具配置: 坐标{'已预旋转' if rotation_config['coordinates_are_pre_rotated'] else '需要旋转'}")
  93. # 判断是否需要旋转坐标
  94. if rotation_config['coordinates_are_pre_rotated']:
  95. # 图片的角度与坐标的角度不一致,比如PPStructV3,图片0度,坐标已旋转270度
  96. # 这种情况下,只需要旋转图片,坐标不变
  97. # PPStructV3: 坐标已经是旋转后的,只旋转图像
  98. img_rotation_angle = (rotation_angle + self.rotated_angle) % 360
  99. if img_rotation_angle == 270:
  100. rotated_image = image.rotate(-90, expand=True) # 顺时针90度
  101. elif img_rotation_angle == 90:
  102. rotated_image = image.rotate(90, expand=True) # 逆时针90度
  103. elif img_rotation_angle == 180:
  104. rotated_image = image.rotate(180, expand=True) # 180度
  105. else:
  106. rotated_image = image.rotate(-img_rotation_angle, expand=True)
  107. if self.rotated_angle == 0:
  108. # 坐标不需要变换,因为JSON中已经是正确的坐标
  109. self._rotated_image_cache[cache_key] = {'image': rotated_image, 'text_bbox_mapping': self.validator.text_bbox_mapping}
  110. self._manage_cache_size()
  111. return rotated_image
  112. image = rotated_image # 继续使用旋转后的图像进行后续处理
  113. # VLM: 需要同时旋转图像和坐标
  114. # 收集所有bbox坐标
  115. all_bboxes = []
  116. text_to_bbox_map = {} # 记录文本到bbox索引的映射
  117. bbox_index = 0
  118. for text, info_list in self.validator.text_bbox_mapping.items():
  119. text_to_bbox_map[text] = []
  120. for info in info_list:
  121. all_bboxes.append(info['bbox'])
  122. text_to_bbox_map[text].append(bbox_index)
  123. bbox_index += 1
  124. # 旋转图像和坐标
  125. rotated_image, rotated_bboxes = rotate_image_and_coordinates(
  126. image, rotation_angle, all_bboxes,
  127. rotate_coordinates=not rotation_config['coordinates_are_pre_rotated']
  128. )
  129. # 更新bbox映射 - 使用映射关系确保正确对应
  130. for text, bbox_indices in text_to_bbox_map.items():
  131. for i, bbox_idx in enumerate(bbox_indices):
  132. if bbox_idx < len(rotated_bboxes) and i < len(self.validator.text_bbox_mapping[text]):
  133. self.validator.text_bbox_mapping[text][i]['bbox'] = rotated_bboxes[bbox_idx]
  134. # 缓存结果
  135. self._rotated_image_cache[cache_key] = {'image': rotated_image, 'text_bbox_mapping': self.validator.text_bbox_mapping}
  136. self._manage_cache_size()
  137. return rotated_image
  138. else:
  139. # 无需旋转,直接缓存原图
  140. self._rotated_image_cache[cache_key] = {'image': image, 'text_bbox_mapping': self.validator.text_bbox_mapping}
  141. self._manage_cache_size() # 检查并管理缓存大小
  142. return image
  143. except Exception as e:
  144. st.error(f"❌ 图像加载失败: {e}")
  145. return None
  146. def render_content_by_mode(self, content: str, render_mode: str, font_size: int,
  147. container_height: int, layout_type: str,
  148. highlight_config: Optional[Dict] = None):
  149. """
  150. 根据渲染模式显示内容 - 增强版本
  151. Args:
  152. content: 要渲染的内容
  153. render_mode: 渲染模式
  154. font_size: 字体大小
  155. container_height: 容器高度
  156. layout_type: 布局类型
  157. highlight_config: 高亮配置 {'has_bbox': bool, 'match_type': str}
  158. """
  159. if content is None or render_mode is None:
  160. return
  161. if render_mode == "HTML渲染":
  162. # 🎯 构建样式 - 包含基础样式和高亮样式
  163. content_style = f"""
  164. <style>
  165. /* ========== 基础容器样式 ========== */
  166. .{layout_type}-content-display {{
  167. height: {container_height}px;
  168. overflow-x: auto;
  169. overflow-y: auto;
  170. font-size: {font_size}px !important;
  171. line-height: 1.4;
  172. color: #333333 !important;
  173. background-color: #fafafa !important;
  174. padding: 10px;
  175. border-radius: 5px;
  176. border: 1px solid #ddd;
  177. max-width: 100%;
  178. }}
  179. /* ========== 表格样式 ========== */
  180. .{layout_type}-content-display table {{
  181. width: 100%;
  182. border-collapse: collapse;
  183. margin: 10px 0;
  184. white-space: nowrap;
  185. }}
  186. .{layout_type}-content-display th,
  187. .{layout_type}-content-display td {{
  188. border: 1px solid #ddd;
  189. padding: 8px;
  190. text-align: left;
  191. max-width: 300px;
  192. word-wrap: break-word;
  193. word-break: break-all;
  194. vertical-align: top;
  195. }}
  196. .{layout_type}-content-display th {{
  197. background-color: #f5f5f5;
  198. position: sticky;
  199. top: 0;
  200. z-index: 1;
  201. font-weight: bold;
  202. }}
  203. /* 数字列右对齐 */
  204. .{layout_type}-content-display td.number {{
  205. text-align: right;
  206. white-space: nowrap;
  207. font-family: 'Monaco', 'Menlo', monospace;
  208. }}
  209. /* 短文本列不换行 */
  210. .{layout_type}-content-display td.short-text {{
  211. white-space: nowrap;
  212. min-width: 80px;
  213. }}
  214. /* ========== 图片样式 ========== */
  215. .{layout_type}-content-display img {{
  216. max-width: 100%;
  217. height: auto;
  218. border-radius: 4px;
  219. margin: 10px 0;
  220. }}
  221. /* ========== 响应式设计 ========== */
  222. @media (max-width: 768px) {{
  223. .{layout_type}-content-display table {{
  224. font-size: {max(font_size-2, 8)}px;
  225. }}
  226. .{layout_type}-content-display th,
  227. .{layout_type}-content-display td {{
  228. padding: 4px;
  229. max-width: 150px;
  230. }}
  231. }}
  232. /* ========== 高亮文本样式 ========== */
  233. .{layout_type}-content-display .highlight-text {{
  234. padding: 2px 4px;
  235. border-radius: 3px;
  236. cursor: pointer;
  237. font-weight: 500;
  238. transition: all 0.2s ease;
  239. }}
  240. .{layout_type}-content-display .highlight-text:hover {{
  241. opacity: 0.8;
  242. transform: scale(1.02);
  243. }}
  244. /* 🎯 精确匹配且有框 - 绿色 */
  245. .{layout_type}-content-display .highlight-text.selected-highlight {{
  246. background-color: #4caf50 !important;
  247. color: white !important;
  248. border: 1px solid #2e7d32 !important;
  249. }}
  250. /* 🎯 OCR匹配 - 蓝色 */
  251. .{layout_type}-content-display .highlight-text.ocr-match {{
  252. background-color: #2196f3 !important;
  253. color: white !important;
  254. border: 1px solid #1565c0 !important;
  255. }}
  256. /* 🎯 无边界框 - 橙色虚线 */
  257. .{layout_type}-content-display .highlight-text.no-bbox {{
  258. background-color: #ff9800 !important;
  259. color: white !important;
  260. border: 1px dashed #f57c00 !important;
  261. }}
  262. /* 🎯 默认高亮 - 黄色 */
  263. .{layout_type}-content-display .highlight-text.default {{
  264. background-color: #ffeb3b !important;
  265. color: #333333 !important;
  266. border: 1px solid #fbc02d !important;
  267. }}
  268. </style>
  269. """
  270. st.markdown(content_style, unsafe_allow_html=True)
  271. st.markdown(f'<div class="{layout_type}-content-display">{content}</div>',
  272. unsafe_allow_html=True)
  273. elif render_mode == "Markdown渲染":
  274. converted_content = convert_html_table_to_markdown(content)
  275. st.markdown(converted_content, unsafe_allow_html=True)
  276. elif render_mode == "DataFrame表格":
  277. if '<table' in content.lower():
  278. self.validator.display_html_table_as_dataframe(content)
  279. else:
  280. st.info("当前内容中没有检测到HTML表格")
  281. st.markdown(content, unsafe_allow_html=True)
  282. else: # 原始文本
  283. st.text_area(
  284. "MD内容预览",
  285. content,
  286. height=300,
  287. key=f"{layout_type}_text_area"
  288. )
  289. def create_compact_layout(self, config: Dict):
  290. """创建紧凑的对比布局 - 增强搜索功能"""
  291. layout = config['styles']['layout']
  292. font_size = config['styles'].get('font_size', 10)
  293. container_height = layout.get('default_height', 600)
  294. zoom_level = layout.get('default_zoom', 1.0)
  295. layout_type = "compact"
  296. left_col, right_col = st.columns([layout['content_width'], layout['sidebar_width']],
  297. vertical_alignment='top', border=True)
  298. with left_col:
  299. if self.validator.text_bbox_mapping:
  300. # 搜索输入框
  301. search_col, select_col = st.columns([1, 2])
  302. if "compact_search_query" not in st.session_state:
  303. st.session_state.compact_search_query = ""
  304. with search_col:
  305. search_query = st.text_input(
  306. "搜索文本",
  307. placeholder="输入关键词...",
  308. value=st.session_state.compact_search_query,
  309. key=f"{layout_type}_search_input",
  310. label_visibility="collapsed"
  311. )
  312. st.session_state.compact_search_query = search_query
  313. # 🎯 增强搜索逻辑:构建选项列表
  314. text_options = ["请选择文本..."]
  315. text_display = ["请选择文本..."]
  316. match_info = [None] # 记录匹配信息
  317. for text, info_list in self.validator.text_bbox_mapping.items():
  318. # 🔑 关键改进:同时搜索 text 和 matched_text
  319. if search_query and search_query.strip():
  320. query_lower = search_query.lower()
  321. # 1. 检查原始文本
  322. text_match = query_lower in text.lower()
  323. # 2. 检查 matched_text(OCR识别文本)
  324. matched_text_match = False
  325. matched_text = None
  326. if info_list and isinstance(info_list[0], dict):
  327. matched_text = info_list[0].get('matched_text', '')
  328. matched_text_match = query_lower in matched_text.lower() if matched_text else False
  329. # 如果都不匹配,跳过
  330. if not text_match and not matched_text_match:
  331. continue
  332. # 记录匹配类型
  333. if text_match:
  334. match_type = "exact"
  335. match_source = text
  336. else:
  337. match_type = "ocr"
  338. match_source = matched_text
  339. else:
  340. match_type = None
  341. match_source = text
  342. text_options.append(text)
  343. # 🎯 构建显示文本(带匹配提示)
  344. if info_list and isinstance(info_list[0], dict):
  345. first_info = info_list[0]
  346. # 检查是否有 bbox
  347. has_bbox = 'bbox' in first_info and first_info['bbox']
  348. # 表格单元格显示
  349. if 'row' in first_info and 'col' in first_info:
  350. display_text = f"[R{first_info['row']},C{first_info['col']}] {text}"
  351. else:
  352. display_text = text
  353. # 🎯 添加匹配提示
  354. if match_type == "ocr":
  355. display_text = f"🔍 {display_text} (OCR: {match_source[:20]}...)"
  356. elif not has_bbox:
  357. display_text = f"⚠️ {display_text} (无框)"
  358. # 截断过长文本
  359. if len(display_text) > 60:
  360. display_text = display_text[:57] + "..."
  361. else:
  362. display_text = text[:57] + "..." if len(text) > 60 else text
  363. text_display.append(display_text)
  364. match_info.append({
  365. 'type': match_type,
  366. 'source': match_source,
  367. 'has_bbox': has_bbox if info_list else False
  368. })
  369. # 🎯 显示搜索统计
  370. if search_query and search_query.strip():
  371. ocr_matches = sum(1 for m in match_info[1:] if m and m['type'] == 'ocr')
  372. no_bbox_count = sum(1 for m in match_info[1:] if m and not m['has_bbox'])
  373. stat_parts = [f"找到 {len(text_options)-1} 个匹配项"]
  374. if ocr_matches > 0:
  375. stat_parts.append(f"🔍 {ocr_matches} 个OCR匹配")
  376. if no_bbox_count > 0:
  377. stat_parts.append(f"⚠️ {no_bbox_count} 个无框")
  378. st.caption(" | ".join(stat_parts))
  379. # 确定默认选中的索引
  380. default_index = 0
  381. if st.session_state.selected_text and st.session_state.selected_text in text_options:
  382. default_index = text_options.index(st.session_state.selected_text)
  383. with select_col:
  384. selected_index = st.selectbox(
  385. "快速定位文本",
  386. range(len(text_options)),
  387. index=default_index,
  388. format_func=lambda x: text_display[x] if x < len(text_display) else "",
  389. label_visibility="collapsed",
  390. key=f"{layout_type}_quick_text_selector"
  391. )
  392. # 🎯 显示匹配详情
  393. if selected_index > 0:
  394. st.session_state.selected_text = text_options[selected_index]
  395. # 获取匹配信息
  396. selected_match_info = match_info[selected_index]
  397. if selected_match_info:
  398. if selected_match_info['type'] == 'ocr':
  399. st.info(f"🔍 **OCR识别文本匹配**: `{selected_match_info['source']}`")
  400. elif not selected_match_info['has_bbox']:
  401. st.warning(f"⚠️ **未找到边界框**: 文本在MD中存在,但没有对应的坐标信息")
  402. # 🎯 增强高亮显示逻辑
  403. if self.validator.md_content:
  404. highlighted_content = self.validator.md_content
  405. if st.session_state.selected_text:
  406. selected_text = st.session_state.selected_text
  407. # 获取匹配信息
  408. info_list = self.validator.text_bbox_mapping.get(selected_text, [])
  409. has_bbox = False
  410. matched_text = None
  411. match_type = None
  412. if info_list and isinstance(info_list[0], dict):
  413. has_bbox = 'bbox' in info_list[0] and info_list[0]['bbox']
  414. matched_text = info_list[0].get('matched_text', '')
  415. # 🔑 判断匹配类型
  416. if matched_text and matched_text != selected_text:
  417. match_type = "ocr"
  418. elif has_bbox:
  419. match_type = "exact"
  420. else:
  421. match_type = "no_bbox"
  422. # 🎯 应用高亮
  423. if len(selected_text) > 2:
  424. # 1. 高亮原始文本
  425. if selected_text in highlighted_content:
  426. if match_type == "exact":
  427. highlight_class = "highlight-text selected-highlight"
  428. elif match_type == "no_bbox":
  429. highlight_class = "highlight-text no-bbox"
  430. else:
  431. highlight_class = "highlight-text default"
  432. highlighted_content = highlighted_content.replace(
  433. selected_text,
  434. f'<span class="{highlight_class}" title="{selected_text}">{selected_text}</span>'
  435. )
  436. # 2. 如果有 matched_text 且不同,也高亮
  437. if matched_text and matched_text != selected_text and matched_text in highlighted_content:
  438. highlighted_content = highlighted_content.replace(
  439. matched_text,
  440. f'<span class="highlight-text ocr-match" title="OCR: {matched_text}">{matched_text}</span>'
  441. )
  442. # 🎯 调用渲染方法(样式已内置)
  443. self.render_content_by_mode(
  444. highlighted_content,
  445. "HTML渲染",
  446. font_size,
  447. container_height,
  448. layout_type
  449. )
  450. with right_col:
  451. self.create_aligned_image_display(zoom_level, "compact")
  452. def create_aligned_image_display(self, zoom_level: float = 1.0, layout_type: str = "aligned"):
  453. """创建响应式图片显示"""
  454. # st.header("🖼️ 原图标注")
  455. # 图片控制选项
  456. col1, col2, col3, col4, col5 = st.columns(5, vertical_alignment="center", border= False)
  457. with col1:
  458. # 判断{layout_type}_show_all_boxes是否有值,如果有值直接使用,否则默认False
  459. # if f"{layout_type}_show_all_boxes" not in st.session_state:
  460. # st.session_state[f"{layout_type}_show_all_boxes"] = False
  461. show_all_boxes = st.checkbox(
  462. "显示所有框",
  463. # value=st.session_state[f"{layout_type}_show_all_boxes"],
  464. value = self.show_all_boxes,
  465. key=f"{layout_type}_show_all_boxes"
  466. )
  467. if show_all_boxes != self.show_all_boxes:
  468. self.show_all_boxes = show_all_boxes
  469. with col2:
  470. if st.button("🔄 旋转90度", type="secondary", key=f"{layout_type}_manual_angle"):
  471. self.rotated_angle = (self.rotated_angle + 90) % 360
  472. # 需要清除图片缓存,以及text_bbox_mapping中的bbox
  473. self.clear_image_cache()
  474. self.validator.process_data()
  475. st.rerun()
  476. with col3:
  477. # 显示当前角度状态
  478. current_angle = self.get_rotation_angle()
  479. st.metric("当前角度", f"{current_angle}°", label_visibility="collapsed")
  480. with col4:
  481. if st.button("↺ 重置角度", key=f"{layout_type}_reset_angle"):
  482. self.rotated_angle = 0.0
  483. st.success("已重置旋转角度")
  484. # 需要清除图片缓存,以及text_bbox_mapping中的bbox
  485. self.clear_image_cache()
  486. self.validator.process_data()
  487. st.rerun()
  488. with col5:
  489. if st.button("🧹 清除选择", key=f"{layout_type}_clear_selection"):
  490. # 清除选中的文本
  491. st.session_state.selected_text = None
  492. # 清除搜索框内容
  493. st.session_state.compact_search_query = None
  494. st.rerun()
  495. # 使用增强的图像加载方法
  496. image = self.load_and_rotate_image(self.validator.image_path)
  497. if image:
  498. try:
  499. resized_image, all_boxes, selected_boxes = self.zoom_image(image, self.zoom_level)
  500. # 创建交互式图片
  501. fig = self.create_resized_interactive_plot(resized_image, selected_boxes, self.zoom_level, all_boxes)
  502. plot_config = {
  503. 'displayModeBar': True,
  504. 'modeBarButtonsToRemove': ['zoom2d', 'select2d', 'lasso2d', 'autoScale2d'],
  505. 'scrollZoom': True,
  506. 'doubleClick': 'reset',
  507. 'responsive': False, # 关键:禁用响应式,使用固定尺寸
  508. 'toImageButtonOptions': {
  509. 'format': 'png',
  510. 'filename': 'ocr_image',
  511. 'height': None, # 使用当前高度
  512. 'width': None, # 使用当前宽度
  513. 'scale': 1
  514. }
  515. }
  516. # 🔧 修复:使用 use_container_width 替代废弃的参数
  517. st.plotly_chart(
  518. fig,
  519. use_container_width=True, # 🎯 使用容器宽度
  520. config=plot_config,
  521. key=f"{layout_type}_plot"
  522. )
  523. except Exception as e:
  524. st.error(f"❌ 图片处理失败: {e}")
  525. st.exception(e)
  526. else:
  527. st.error("未找到对应的图片文件")
  528. if self.validator.image_path:
  529. st.write(f"期望路径: {self.validator.image_path}")
  530. # st.markdown('</div>', unsafe_allow_html=True)
  531. def zoom_image(self, image: Image.Image, current_zoom: float) -> Tuple[Image.Image, List[List[int]], List[List[int]]]:
  532. """缩放图像"""
  533. # 根据缩放级别调整图片大小
  534. new_width = int(image.width * current_zoom)
  535. new_height = int(image.height * current_zoom)
  536. resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
  537. # 计算选中的bbox
  538. selected_boxes = []
  539. if st.session_state.selected_text and st.session_state.selected_text in self.validator.text_bbox_mapping:
  540. info_list = self.validator.text_bbox_mapping[st.session_state.selected_text]
  541. for info in info_list:
  542. if 'bbox' in info:
  543. bbox = info['bbox']
  544. selected_box = [int(coord * current_zoom) for coord in bbox]
  545. selected_boxes.append(selected_box)
  546. # 收集所有框
  547. all_boxes = []
  548. if self.show_all_boxes:
  549. for text, info_list in self.validator.text_bbox_mapping.items():
  550. for info in info_list:
  551. bbox = info['bbox']
  552. if len(bbox) >= 4:
  553. scaled_bbox = [coord * current_zoom for coord in bbox]
  554. all_boxes.append(scaled_bbox)
  555. return resized_image, all_boxes, selected_boxes
  556. def _add_bboxes_to_plot_batch(self, fig: go.Figure, bboxes: List[List[int]],
  557. image_height: int,
  558. line_color: str = "blue",
  559. line_width: int = 2,
  560. fill_color: str = "rgba(0, 100, 200, 0.2)"):
  561. """
  562. 批量添加边界框(性能优化版)
  563. """
  564. if not bboxes or len(bboxes) == 0:
  565. return
  566. # 🎯 关键优化:构建 shapes 列表,一次性添加
  567. shapes = []
  568. for bbox in bboxes:
  569. if len(bbox) < 4:
  570. continue
  571. x1, y1, x2, y2 = bbox[:4]
  572. # 转换坐标
  573. plot_x1 = x1
  574. plot_x2 = x2
  575. plot_y1 = image_height - y2
  576. plot_y2 = image_height - y1
  577. shapes.append(dict(
  578. type="rect",
  579. x0=plot_x1, y0=plot_y1,
  580. x1=plot_x2, y1=plot_y2,
  581. line=dict(color=line_color, width=line_width),
  582. fillcolor=fill_color,
  583. ))
  584. # 🎯 一次性更新所有形状
  585. fig.update_layout(shapes=fig.layout.shapes + tuple(shapes))
  586. def _add_bboxes_as_scatter(self, fig: go.Figure, bboxes: List[List[int]],
  587. image_height: int,
  588. line_color: str = "blue",
  589. line_width: int = 2,
  590. name: str = "boxes"):
  591. """
  592. 使用 Scatter 绘制边界框(极致性能优化)
  593. """
  594. if not bboxes or len(bboxes) == 0:
  595. return
  596. # 🎯 收集所有矩形的边框线坐标
  597. x_coords = []
  598. y_coords = []
  599. for bbox in bboxes:
  600. if len(bbox) < 4:
  601. continue
  602. x1, y1, x2, y2 = bbox[:4]
  603. # 转换坐标
  604. plot_y1 = image_height - y2
  605. plot_y2 = image_height - y1
  606. # 绘制矩形:5个点(闭合)
  607. x_coords.extend([x1, x2, x2, x1, x1, None]) # None用于断开线段
  608. y_coords.extend([plot_y1, plot_y1, plot_y2, plot_y2, plot_y1, None])
  609. # 🎯 一次性添加所有边框
  610. fig.add_trace(go.Scatter(
  611. x=x_coords,
  612. y=y_coords,
  613. mode='lines',
  614. line=dict(color=line_color, width=line_width),
  615. name=name,
  616. showlegend=False,
  617. hoverinfo='skip'
  618. ))
  619. def create_resized_interactive_plot(self, image: Image.Image, selected_boxes: List[List[int]],
  620. zoom_level: float, all_boxes: List[List[int]]) -> go.Figure:
  621. """创建可调整大小的交互式图片 - 修复容器溢出问题"""
  622. fig = go.Figure()
  623. # 添加图片 - Plotly坐标系,原点在左下角
  624. fig.add_layout_image(
  625. dict(
  626. source=image,
  627. xref="x", yref="y",
  628. x=0, y=image.height, # 图片左下角在Plotly坐标系中的位置
  629. sizex=image.width,
  630. sizey=image.height,
  631. sizing="stretch",
  632. opacity=1.0,
  633. layer="below",
  634. yanchor="top" # 确保图片顶部对齐
  635. )
  636. )
  637. # 显示所有bbox(淡蓝色)
  638. if all_boxes:
  639. self._add_bboxes_as_scatter(
  640. fig=fig,
  641. bboxes=all_boxes,
  642. image_height=image.height,
  643. line_color="rgba(0, 100, 200, 0.8)",
  644. line_width=2,
  645. name="all_boxes"
  646. )
  647. # 高亮显示选中的bbox(红色)
  648. if selected_boxes:
  649. self._add_bboxes_to_plot_batch(
  650. fig=fig,
  651. bboxes=selected_boxes,
  652. image_height=image.height,
  653. line_color="red",
  654. line_width=2,
  655. fill_color="rgba(255, 0, 0, 0.3)"
  656. )
  657. # 修复:优化显示尺寸计算
  658. max_display_width = 1500
  659. max_display_height = 1000
  660. # 计算合适的显示尺寸,保持宽高比
  661. aspect_ratio = image.width / image.height
  662. if self.fit_to_container:
  663. # 自适应容器模式
  664. if aspect_ratio > 1: # 宽图
  665. display_width = min(max_display_width, image.width)
  666. display_height = int(display_width / aspect_ratio)
  667. else: # 高图
  668. display_height = min(max_display_height, image.height)
  669. display_width = int(display_height * aspect_ratio)
  670. # 确保不会太小
  671. display_width = max(display_width, 800)
  672. display_height = max(display_height, 600)
  673. else:
  674. # 固定尺寸模式,但仍要考虑容器限制
  675. display_width = min(image.width, max_display_width)
  676. display_height = min(image.height, max_display_height)
  677. # 设置布局 - 关键修改
  678. fig.update_layout(
  679. width=display_width,
  680. height=display_height,
  681. margin=dict(l=0, r=0, t=0, b=0),
  682. showlegend=False,
  683. plot_bgcolor='white',
  684. dragmode="pan",
  685. # 关键:让图表自适应容器
  686. # autosize=True, # 启用自动调整大小
  687. xaxis=dict(
  688. visible=False,
  689. range=[0, image.width],
  690. constrain="domain",
  691. fixedrange=False,
  692. autorange=False,
  693. showgrid=False,
  694. zeroline=False,
  695. ),
  696. # 修复:Y轴设置,确保范围正确
  697. yaxis=dict(
  698. visible=False,
  699. range=[0, image.height], # 确保Y轴范围从0到图片高度
  700. constrain="domain",
  701. scaleanchor="x",
  702. scaleratio=1,
  703. fixedrange=False,
  704. autorange=False,
  705. showgrid=False,
  706. zeroline=False
  707. )
  708. )
  709. return fig