adjustments.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261
  1. """
  2. 手动调整功能
  3. """
  4. from .state_manager import save_state_for_undo
  5. from .drawing import clear_table_image_cache
  6. def create_adjustment_section(structure):
  7. """
  8. 创建手动调整区域
  9. """
  10. import streamlit as st
  11. st.divider()
  12. st.header("🛠️ 手动调整")
  13. horizontal_lines = structure.get('horizontal_lines', [])
  14. vertical_lines = structure.get('vertical_lines', [])
  15. adjusted = False
  16. # 行操作, 列操作
  17. adjustment_action = st.radio(
  18. "行&列操作",
  19. ["调整横线", "添加横线", "删除横线", "调整竖线", "添加竖线", "删除竖线", "微调旋转"],
  20. horizontal=True,
  21. index=None,
  22. label_visibility="collapsed",
  23. key="adjustment_action_radio"
  24. )
  25. if adjustment_action == "调整横线" and horizontal_lines:
  26. line_index = st.selectbox(
  27. "选择横线",
  28. range(len(horizontal_lines)),
  29. format_func=lambda i: f"R{i+1} (Y={horizontal_lines[i]})",
  30. key="adjust_h_select"
  31. )
  32. new_y = st.number_input(
  33. "新的Y坐标",
  34. min_value=0,
  35. value=horizontal_lines[line_index],
  36. step=1,
  37. key="adjust_h_value"
  38. )
  39. if st.button("✅ 应用横线调整"):
  40. if new_y != horizontal_lines[line_index]:
  41. save_state_for_undo(structure)
  42. structure['horizontal_lines'][line_index] = new_y
  43. # structure.setdefault('modified_h_lines', [])
  44. if line_index not in structure['modified_h_lines']:
  45. structure['modified_h_lines'].append(line_index)
  46. _update_row_intervals(structure)
  47. clear_table_image_cache()
  48. adjusted = True
  49. st.success(f"✅ R{line_index+1} 已更新")
  50. elif adjustment_action == "添加横线":
  51. new_h_y = st.number_input(
  52. "新横线的Y坐标",
  53. min_value=0,
  54. value=horizontal_lines[-1] + 50 if horizontal_lines else 100,
  55. step=1,
  56. key="add_h_value"
  57. )
  58. if st.button("➕ 确认添加横线"):
  59. save_state_for_undo(structure)
  60. structure['horizontal_lines'].append(new_h_y)
  61. structure['horizontal_lines'].sort()
  62. idx = structure['horizontal_lines'].index(new_h_y)
  63. # structure.setdefault('modified_h_lines', [])
  64. if idx not in structure['modified_h_lines']:
  65. structure['modified_h_lines'].append(idx)
  66. _update_row_intervals(structure)
  67. clear_table_image_cache()
  68. adjusted = True
  69. st.success(f"✅ 新增横线 Y={new_h_y}")
  70. elif adjustment_action == "删除横线" and len(horizontal_lines) > 2:
  71. to_delete = st.multiselect(
  72. "选择要删除的横线",
  73. range(len(horizontal_lines)),
  74. format_func=lambda i: f"R{i+1} (Y={horizontal_lines[i]})",
  75. key="del_h_select"
  76. )
  77. if to_delete and st.button("🗑️ 确认删除横线"):
  78. save_state_for_undo(structure)
  79. for idx in sorted(to_delete, reverse=True):
  80. del structure['horizontal_lines'][idx]
  81. structure['modified_h_lines'] = []
  82. _update_row_intervals(structure)
  83. clear_table_image_cache()
  84. adjusted = True
  85. st.success(f"✅ 已删除 {len(to_delete)} 条横线")
  86. elif adjustment_action == "调整竖线" and vertical_lines:
  87. line_index = st.selectbox(
  88. "选择竖线",
  89. range(len(vertical_lines)),
  90. format_func=lambda i: f"C{i+1} (X={vertical_lines[i]})",
  91. key="adjust_v_select"
  92. )
  93. new_x = st.number_input(
  94. "新的X坐标",
  95. min_value=0,
  96. value=vertical_lines[line_index],
  97. step=1,
  98. key="adjust_v_value"
  99. )
  100. if st.button("✅ 应用竖线调整"):
  101. if new_x != vertical_lines[line_index]:
  102. save_state_for_undo(structure)
  103. structure['vertical_lines'][line_index] = new_x
  104. # structure.setdefault('modified_v_lines', [])
  105. if line_index not in structure['modified_v_lines']:
  106. structure['modified_v_lines'].append(line_index)
  107. _update_column_intervals(structure)
  108. clear_table_image_cache()
  109. adjusted = True
  110. st.success(f"✅ C{line_index+1} 已更新")
  111. elif adjustment_action == "添加竖线":
  112. new_v_x = st.number_input(
  113. "新竖线的X坐标",
  114. min_value=0,
  115. value=vertical_lines[-1] + 100 if vertical_lines else 100,
  116. step=1,
  117. key="add_v_value"
  118. )
  119. if st.button("➕ 确认添加竖线"):
  120. structure['vertical_lines'].append(new_v_x)
  121. structure['vertical_lines'].sort()
  122. idx = structure['vertical_lines'].index(new_v_x)
  123. # structure.setdefault('modified_v_lines', [])
  124. if idx not in structure['modified_v_lines']:
  125. structure['modified_v_lines'].append(idx)
  126. _update_column_intervals(structure)
  127. clear_table_image_cache()
  128. adjusted = True
  129. st.success(f"✅ 新增竖线 X={new_v_x}")
  130. elif adjustment_action == "删除竖线" and len(vertical_lines) > 2:
  131. to_delete = st.multiselect(
  132. "选择要删除的竖线",
  133. range(len(vertical_lines)),
  134. format_func=lambda i: f"C{i+1} (X={vertical_lines[i]})",
  135. key="del_v_select"
  136. )
  137. if to_delete and st.button("🗑️ 确认删除竖线"):
  138. save_state_for_undo(structure)
  139. for idx in sorted(to_delete, reverse=True):
  140. del structure['vertical_lines'][idx]
  141. structure['modified_v_lines'] = []
  142. _update_column_intervals(structure)
  143. clear_table_image_cache()
  144. adjusted = True
  145. st.success(f"✅ 已删除 {len(to_delete)} 条竖线")
  146. elif adjustment_action == "微调旋转":
  147. st.info("📐 手动修正图片的倾斜角度。系统会根据设定的倾斜角度,进行**反向旋转**来校正图片。")
  148. if 'generator' in st.session_state:
  149. generator = st.session_state.generator
  150. current_skew = generator.ocr_data.get('skew_angle', 0.0)
  151. original_skew = generator.ocr_data.get('original_skew_angle', 0.0)
  152. # 计算实际执行的旋转角度
  153. applied_rotation = -current_skew
  154. st.markdown(f"""
  155. * **原始识别倾斜:** `{original_skew:.2f}°` (OCR模型检测到的原始角度)
  156. * **当前设定倾斜:** `{current_skew:.2f}°` (系统认为图片歪了多少)
  157. * **实际执行旋转:** `{applied_rotation:.2f}°` (系统实际执行的**反向**校正旋转)
  158. """)
  159. col1, col2 = st.columns([1, 1], width=300, gap="small")
  160. with col1:
  161. delta_angle = st.number_input(
  162. "调整设定值 (度)",
  163. min_value=-10.0,
  164. max_value=10.0,
  165. value=0.0,
  166. step=0.05,
  167. format="%.2f",
  168. help="增加设定值 -> 系统认为图片逆时针歪得更多 -> 执行更大幅度的顺时针旋转",
  169. label_visibility="collapsed",
  170. key="rotate_delta_input"
  171. )
  172. with col2:
  173. if st.button("🔄 应用调整", key="apply_rotate_btn"):
  174. if delta_angle != 0:
  175. # 更新 skew_angle
  176. generator.ocr_data['skew_angle'] = current_skew + delta_angle
  177. with st.spinner("正在重新校正图片..."):
  178. # 强制重新校正
  179. corrected_image, _ = generator.correct_skew(force=True)
  180. st.session_state.image = corrected_image
  181. # 清除缓存
  182. clear_table_image_cache()
  183. adjusted = True
  184. # 标记需要重新分析结构
  185. st.session_state.need_reanalysis = True
  186. st.success(f"✅ 已应用旋转 {delta_angle}°")
  187. st.rerun()
  188. else:
  189. st.warning("⚠️ 无法获取 TableLineGenerator 实例")
  190. return adjusted
  191. def _update_row_intervals(structure):
  192. """根据横线坐标更新行区间"""
  193. horizontal_lines = structure.get('horizontal_lines', [])
  194. rows = []
  195. for i in range(len(horizontal_lines) - 1):
  196. rows.append({
  197. 'y_start': horizontal_lines[i],
  198. 'y_end': horizontal_lines[i + 1],
  199. 'bboxes': []
  200. })
  201. structure['rows'] = rows
  202. # 更新表格边界框
  203. if 'table_bbox' in structure:
  204. vertical_lines = structure.get('vertical_lines', [])
  205. structure['table_bbox'] = [
  206. vertical_lines[0] if vertical_lines else 0,
  207. horizontal_lines[0],
  208. vertical_lines[-1] if vertical_lines else 0,
  209. horizontal_lines[-1]
  210. ]
  211. def _update_column_intervals(structure):
  212. """根据竖线坐标更新列区间"""
  213. vertical_lines = structure.get('vertical_lines', [])
  214. columns = []
  215. for i in range(len(vertical_lines) - 1):
  216. columns.append({
  217. 'x_start': vertical_lines[i],
  218. 'x_end': vertical_lines[i + 1]
  219. })
  220. structure['columns'] = columns
  221. # 更新列宽
  222. col_widths = [col['x_end'] - col['x_start'] for col in columns]
  223. structure['col_widths'] = col_widths
  224. # 更新表格边界框
  225. if 'table_bbox' in structure:
  226. horizontal_lines = structure.get('horizontal_lines', [])
  227. structure['table_bbox'] = [
  228. vertical_lines[0],
  229. horizontal_lines[0] if horizontal_lines else 0,
  230. vertical_lines[-1],
  231. horizontal_lines[-1] if horizontal_lines else 0
  232. ]