mode_setup.py 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  1. """
  2. 模式设置(新建/编辑)
  3. """
  4. import streamlit as st
  5. from PIL import Image
  6. from typing import Dict, Tuple
  7. import sys
  8. from pathlib import Path
  9. # 添加父目录到路径
  10. sys.path.insert(0, str(Path(__file__).parent.parent))
  11. try:
  12. from ..table_line_generator import TableLineGenerator
  13. except ImportError:
  14. from table_line_generator import TableLineGenerator
  15. from .display_controls import create_display_settings_section, create_undo_redo_section
  16. from .analysis_controls import create_analysis_section
  17. def setup_new_annotation_mode(
  18. ocr_data: list,
  19. image: Image.Image,
  20. display_config: Dict
  21. ) -> Tuple:
  22. """
  23. 设置新建标注模式
  24. Returns:
  25. (generator, structure, undo_stack, line_width, display_mode, zoom_level, show_line_numbers)
  26. """
  27. # 🔑 获取当前工具类型
  28. tool = st.session_state.get('current_tool', 'ppstructv3')
  29. # 初始化生成器
  30. if 'generator' not in st.session_state:
  31. st.session_state.generator = TableLineGenerator(image, ocr_data)
  32. # 🆕 自动校正倾斜和旋转
  33. corrected_image, angle = st.session_state.generator.correct_skew()
  34. # 获取角度信息
  35. ocr_data_dict = st.session_state.generator.ocr_data
  36. image_rotation_angle = ocr_data_dict.get('image_rotation_angle', 0.0) if isinstance(ocr_data_dict, dict) else 0.0
  37. skew_angle = ocr_data_dict.get('skew_angle', 0.0) if isinstance(ocr_data_dict, dict) else 0.0
  38. if abs(skew_angle) > 0.1 or image_rotation_angle != 0:
  39. st.info(f"🔄 自动校正: 旋转={image_rotation_angle}°, 倾斜={skew_angle:.2f}°")
  40. # 更新 session_state 中的图片
  41. st.session_state.image = corrected_image
  42. # 分析控件
  43. structure = create_analysis_section(
  44. st.session_state.generator,
  45. tool=tool # 🔑 传入工具类型
  46. )
  47. # 显示控件
  48. line_width, display_mode, zoom_level, show_line_numbers = create_display_settings_section(
  49. display_config
  50. )
  51. # 撤销/重做
  52. undo_stack = []
  53. return (
  54. st.session_state.generator,
  55. structure,
  56. undo_stack,
  57. line_width,
  58. display_mode,
  59. zoom_level,
  60. show_line_numbers
  61. )
  62. def setup_edit_annotation_mode(
  63. structure: Dict,
  64. image: Image.Image,
  65. display_config: Dict
  66. ) -> Tuple:
  67. """
  68. 设置编辑标注模式
  69. Returns:
  70. (image, line_width, display_mode, zoom_level, show_line_numbers)
  71. """
  72. # 显示控件
  73. line_width, display_mode, zoom_level, show_line_numbers = create_display_settings_section(
  74. display_config
  75. )
  76. # 撤销/重做控件
  77. create_undo_redo_section()
  78. return (
  79. image,
  80. line_width,
  81. display_mode,
  82. zoom_level,
  83. show_line_numbers
  84. )