400 lines
15 KiB
Python
400 lines
15 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
Step5_5 面板 - 水质指数计算
|
||
"""
|
||
|
||
import os
|
||
from pathlib import Path
|
||
from typing import Dict, List, Union
|
||
|
||
import pandas as pd
|
||
from PyQt5.QtWidgets import (
|
||
QWidget, QVBoxLayout, QGroupBox, QFormLayout, QGridLayout,
|
||
QHBoxLayout, QLabel, QLineEdit, QComboBox, QCheckBox,
|
||
QPushButton, QMessageBox,
|
||
)
|
||
from PyQt5.QtCore import Qt
|
||
|
||
from src.gui.components.custom_widgets import FileSelectWidget
|
||
from src.gui.styles import ModernStylesheet
|
||
|
||
|
||
class Step5_5Panel(QWidget):
|
||
"""步骤5.5：水质指数计算"""
|
||
|
||
def __init__(self, parent=None):
|
||
super().__init__(parent)
|
||
self.index_checkboxes: Dict[str, QCheckBox] = {}
|
||
self.csv_columns = [] # 存储CSV文件列名
|
||
self.init_ui()
|
||
|
||
def init_ui(self):
|
||
main_layout = QVBoxLayout()
|
||
|
||
# 标题
|
||
|
||
|
||
# 数据文件选择
|
||
data_group = QGroupBox("数据文件")
|
||
data_layout = QVBoxLayout()
|
||
|
||
# 训练数据CSV文件选择
|
||
self.training_data_widget = FileSelectWidget("训练数据CSV文件:", "CSV Files (*.csv)")
|
||
data_layout.addWidget(self.training_data_widget)
|
||
|
||
# 公式CSV文件选择
|
||
self.formula_csv_widget = FileSelectWidget("公式CSV文件:", "CSV Files (*.csv)")
|
||
data_layout.addWidget(self.formula_csv_widget)
|
||
|
||
# 刷新公式按钮
|
||
refresh_layout = QHBoxLayout()
|
||
self.refresh_button = QPushButton("刷新公式列表")
|
||
self.refresh_button.clicked.connect(self.refresh_formulas)
|
||
refresh_layout.addWidget(self.refresh_button)
|
||
refresh_layout.addStretch()
|
||
data_layout.addLayout(refresh_layout)
|
||
|
||
data_group.setLayout(data_layout)
|
||
main_layout.addWidget(data_group)
|
||
|
||
# 公式选择区域
|
||
self.formula_group = QGroupBox("选择要计算的公式")
|
||
formula_outer_layout = QVBoxLayout()
|
||
|
||
# 按钮控制区域
|
||
button_layout = QHBoxLayout()
|
||
self.select_all_btn = QPushButton("全选")
|
||
self.select_all_btn.clicked.connect(self.select_all_formulas)
|
||
self.deselect_all_btn = QPushButton("清空")
|
||
self.deselect_all_btn.clicked.connect(self.deselect_all_formulas)
|
||
button_layout.addWidget(self.select_all_btn)
|
||
button_layout.addWidget(self.deselect_all_btn)
|
||
button_layout.addStretch()
|
||
|
||
formula_outer_layout.addLayout(button_layout)
|
||
|
||
# 公式勾选框网格布局
|
||
self.formula_layout = QGridLayout()
|
||
formula_outer_layout.addLayout(self.formula_layout)
|
||
|
||
self.formula_group.setLayout(formula_outer_layout)
|
||
main_layout.addWidget(self.formula_group)
|
||
|
||
# 输出文件设置
|
||
output_group = QGroupBox("输出设置")
|
||
output_layout = QVBoxLayout()
|
||
|
||
self.output_file_widget = FileSelectWidget(
|
||
"输出文件:", "CSV Files (*.csv)", mode="save"
|
||
)
|
||
output_layout.addWidget(self.output_file_widget)
|
||
|
||
output_group.setLayout(output_layout)
|
||
main_layout.addWidget(output_group)
|
||
|
||
# 启用选项
|
||
self.enable_checkbox = QCheckBox("启用此步骤")
|
||
self.enable_checkbox.setChecked(True)
|
||
main_layout.addWidget(self.enable_checkbox)
|
||
|
||
# 独立运行按钮
|
||
self.run_button = QPushButton("独立运行此步骤")
|
||
self.run_button.setStyleSheet(ModernStylesheet.get_button_stylesheet('success'))
|
||
self.run_button.clicked.connect(self.run_step)
|
||
main_layout.addWidget(self.run_button)
|
||
|
||
# 公式编辑区域
|
||
formula_edit_group = QGroupBox("添加自定义公式")
|
||
formula_edit_layout = QFormLayout()
|
||
|
||
self.formula_name_edit = QLineEdit()
|
||
|
||
# 公式类别下拉选择框
|
||
self.formula_category_combo = QComboBox()
|
||
self.formula_category_combo.addItems([
|
||
"chlorophyll_a",
|
||
"Phycocyanin (BGA_PC)",
|
||
"Total Nitrogen (TN)",
|
||
"Total Phosphorus (TP)",
|
||
"Orthophosphate",
|
||
"COD",
|
||
"BOD",
|
||
"TOC",
|
||
"Dissolved Oxygen (DO)",
|
||
"E. coli",
|
||
"Total Coliforms",
|
||
"Turbidity",
|
||
"Total Suspended Solids (TSS)",
|
||
"Color",
|
||
"pH",
|
||
"Temperature",
|
||
"Conductivity",
|
||
"Total Dissolved Solids (TDS)"
|
||
])
|
||
self.formula_category_combo.setEditable(True) # 允许用户输入自定义类别
|
||
|
||
self.formula_expression_edit = QLineEdit()
|
||
self.formula_reference_edit = QLineEdit()
|
||
|
||
formula_edit_layout.addRow("公式名称:", self.formula_name_edit)
|
||
formula_edit_layout.addRow("公式类别:", self.formula_category_combo)
|
||
formula_edit_layout.addRow("公式表达式:", self.formula_expression_edit)
|
||
formula_edit_layout.addRow("参考文献:", self.formula_reference_edit)
|
||
|
||
add_button = QPushButton("添加公式")
|
||
add_button.clicked.connect(self.add_custom_formula)
|
||
formula_edit_layout.addRow(add_button)
|
||
|
||
formula_edit_group.setLayout(formula_edit_layout)
|
||
main_layout.addWidget(formula_edit_group)
|
||
|
||
main_layout.addStretch()
|
||
self.setLayout(main_layout)
|
||
|
||
# 自动加载内置公式文件
|
||
formula_csv_path = (
|
||
Path(__file__).resolve().parent.parent / "model" / "waterindex.csv"
|
||
)
|
||
if formula_csv_path.is_file():
|
||
self.formula_csv_widget.set_path(str(formula_csv_path))
|
||
self.refresh_formulas()
|
||
|
||
def refresh_formulas(self):
|
||
"""刷新公式列表"""
|
||
formula_csv_path = self.formula_csv_widget.get_path()
|
||
if not formula_csv_path or not os.path.exists(formula_csv_path):
|
||
QMessageBox.warning(self, "警告", "请先选择有效的公式CSV文件")
|
||
return
|
||
|
||
try:
|
||
# 清除现有的勾选框
|
||
for checkbox in self.index_checkboxes.values():
|
||
self.formula_layout.removeWidget(checkbox)
|
||
checkbox.deleteLater()
|
||
self.index_checkboxes.clear()
|
||
|
||
# 读取公式CSV文件
|
||
df = pd.read_csv(formula_csv_path)
|
||
if df.empty or 'Formula_Name' not in df.columns:
|
||
QMessageBox.warning(self, "警告", "公式CSV文件格式不正确")
|
||
return
|
||
|
||
# 获取所有公式名称（跳过第一行）
|
||
formula_names = df['Formula_Name'].tolist()[1:]
|
||
|
||
# 创建3列布局的勾选框
|
||
row, col = 0, 0
|
||
for formula_name in formula_names:
|
||
if pd.isna(formula_name) or not formula_name.strip():
|
||
continue
|
||
|
||
checkbox = QCheckBox(formula_name.strip())
|
||
checkbox.setChecked(True)
|
||
self.index_checkboxes[formula_name.strip()] = checkbox
|
||
self.formula_layout.addWidget(checkbox, row, col)
|
||
|
||
col += 1
|
||
if col >= 3: # 每行3列
|
||
col = 0
|
||
row += 1
|
||
|
||
except Exception as e:
|
||
QMessageBox.critical(self, "错误", f"读取公式文件失败: {str(e)}")
|
||
|
||
def add_custom_formula(self):
|
||
"""添加自定义公式到公式CSV文件"""
|
||
formula_csv_path = self.formula_csv_widget.get_path()
|
||
if not formula_csv_path:
|
||
QMessageBox.warning(self, "警告", "请先选择公式CSV文件")
|
||
return
|
||
|
||
formula_name = self.formula_name_edit.text().strip()
|
||
formula_category = self.formula_category_combo.currentText().strip()
|
||
formula_expression = self.formula_expression_edit.text().strip()
|
||
formula_reference = self.formula_reference_edit.text().strip()
|
||
|
||
if not all([formula_name, formula_category, formula_expression]):
|
||
QMessageBox.warning(self, "警告", "请填写公式名称、类别和表达式")
|
||
return
|
||
|
||
try:
|
||
# 读取现有公式文件或创建新文件
|
||
if os.path.exists(formula_csv_path):
|
||
df = pd.read_csv(formula_csv_path)
|
||
else:
|
||
df = pd.DataFrame(columns=['Formula_Name', 'Category', 'Formula', 'Reference'])
|
||
|
||
# 添加新公式
|
||
new_row = pd.DataFrame({
|
||
'Formula_Name': [formula_name],
|
||
'Category': [formula_category],
|
||
'Formula': [formula_expression],
|
||
'Reference': [formula_reference]
|
||
})
|
||
df = pd.concat([df, new_row], ignore_index=True)
|
||
|
||
# 保存文件
|
||
df.to_csv(formula_csv_path, index=False, encoding='utf-8')
|
||
|
||
# 清空输入框
|
||
self.formula_name_edit.clear()
|
||
self.formula_category_combo.setCurrentIndex(0) # 重置到第一个选项
|
||
self.formula_expression_edit.clear()
|
||
self.formula_reference_edit.clear()
|
||
|
||
# 刷新公式列表
|
||
self.refresh_formulas()
|
||
|
||
QMessageBox.information(self, "成功", "公式添加成功")
|
||
|
||
except Exception as e:
|
||
QMessageBox.critical(self, "错误", f"添加公式失败: {str(e)}")
|
||
|
||
def get_config(self) -> Dict[str, Union[List[str], str, bool]]:
|
||
"""获取配置"""
|
||
selected = [
|
||
name for name, checkbox in self.index_checkboxes.items()
|
||
if checkbox.isChecked()
|
||
]
|
||
output_path = self.output_file_widget.get_path()
|
||
return {
|
||
'training_spectra_path': self.training_data_widget.get_path() or None,
|
||
'formula_csv_file': self.formula_csv_widget.get_path() or None,
|
||
'formula_names': selected,
|
||
'output_file': output_path or None,
|
||
'enabled': self.enable_checkbox.isChecked()
|
||
}
|
||
|
||
def set_config(self, config):
|
||
"""设置配置"""
|
||
if 'training_spectra_path' in config:
|
||
self.training_data_widget.set_path(config['training_spectra_path'])
|
||
|
||
if 'formula_csv_file' in config:
|
||
self.formula_csv_widget.set_path(config['formula_csv_file'])
|
||
self.refresh_formulas()
|
||
|
||
if 'formula_names' in config:
|
||
selected_formulas = set(config['formula_names'])
|
||
for name, checkbox in self.index_checkboxes.items():
|
||
checkbox.setChecked(name in selected_formulas)
|
||
|
||
if 'output_file' in config and config['output_file']:
|
||
self.output_file_widget.set_path(config['output_file'])
|
||
elif 'output_filename' in config and config['output_filename']:
|
||
self.output_file_widget.set_path(config['output_filename'])
|
||
|
||
if 'enabled' in config:
|
||
self.enable_checkbox.setChecked(config['enabled'])
|
||
|
||
def update_from_config(self, work_dir=None, pipeline=None):
|
||
"""从全局配置自动填充训练数据和输出路径
|
||
|
||
Args:
|
||
work_dir: 工作目录路径
|
||
pipeline: Pipeline 实例（未使用，保留接口兼容性）
|
||
"""
|
||
if work_dir:
|
||
self.work_dir = work_dir
|
||
elif hasattr(self, 'work_dir') and self.work_dir:
|
||
pass
|
||
else:
|
||
self.work_dir = None
|
||
|
||
# 1. 自动填入训练数据路径（从 Step5 的输出中获取）
|
||
# 优先级：直接 widget > pipeline.step_outputs 回退
|
||
main_window = self.window()
|
||
if hasattr(main_window, 'step5_panel'):
|
||
# 优先直接从 Step5 的输出 widget 读取（已运行的最新输出）
|
||
step5_output = main_window.step5_panel.output_file.get_path()
|
||
if step5_output:
|
||
# 若为相对路径，使用 work_dir 合成为绝对路径
|
||
if not os.path.isabs(step5_output):
|
||
step5_output = os.path.join(self.work_dir or '', step5_output).replace('\\', '/')
|
||
self.training_data_widget.set_path(step5_output)
|
||
else:
|
||
# 退而求其次，使用 Step5 的输入 CSV
|
||
step5_csv = main_window.step5_panel.csv_file.get_path()
|
||
if step5_csv:
|
||
# 若为相对路径，使用 work_dir 合成为绝对路径
|
||
if not os.path.isabs(step5_csv):
|
||
step5_csv = os.path.join(self.work_dir or '', step5_csv).replace('\\', '/')
|
||
self.training_data_widget.set_path(step5_csv)
|
||
|
||
# 如果上述都没找到，尝试从 pipeline.step_outputs 回退
|
||
if not self.training_data_widget.get_path() and pipeline and hasattr(pipeline, 'step_outputs'):
|
||
step5_outputs = getattr(pipeline, 'step_outputs', {}).get('step5', {})
|
||
training_path = step5_outputs.get('training_spectra')
|
||
if training_path:
|
||
self.training_data_widget.set_path(training_path)
|
||
|
||
# 2. 自动填入输出文件的绝对路径
|
||
if self.work_dir:
|
||
output_abs = os.path.join(self.work_dir, "6_water_quality_indices",
|
||
"training_spectra_indices.csv").replace('\\', '/')
|
||
self.output_file_widget.set_path(output_abs)
|
||
|
||
def is_enabled(self) -> bool:
|
||
return self.enable_checkbox.isChecked()
|
||
|
||
def select_all_formulas(self):
|
||
"""全选所有公式"""
|
||
for checkbox in self.index_checkboxes.values():
|
||
checkbox.setChecked(True)
|
||
|
||
def deselect_all_formulas(self):
|
||
"""清空所有公式"""
|
||
for checkbox in self.index_checkboxes.values():
|
||
checkbox.setChecked(False)
|
||
|
||
def run_step(self):
|
||
"""独立运行步骤5.5：计算水质指数。
|
||
|
||
动态根据输入 CSV 文件名生成输出文件名，自动填入 output_file_widget。
|
||
例如：training_spectra.csv → training_spectra_indices.csv
|
||
sampling_spectra.csv → sampling_spectra_indices.csv
|
||
"""
|
||
# 验证输入
|
||
training_csv_path = self.training_data_widget.get_path()
|
||
formula_csv_path = self.formula_csv_widget.get_path()
|
||
|
||
if not training_csv_path:
|
||
QMessageBox.warning(self, "输入验证失败", "请选择训练数据CSV文件")
|
||
return
|
||
if not formula_csv_path:
|
||
QMessageBox.warning(self, "输入验证失败", "请选择公式CSV文件")
|
||
return
|
||
if not os.path.exists(training_csv_path):
|
||
QMessageBox.warning(self, "输入验证失败", "训练数据CSV文件不存在")
|
||
return
|
||
if not os.path.exists(formula_csv_path):
|
||
QMessageBox.warning(self, "输入验证失败", "公式CSV文件不存在")
|
||
return
|
||
|
||
# 动态生成输出文件：自动拼接 _indices 后缀
|
||
input_name = Path(training_csv_path).stem
|
||
dynamic_output = f"{input_name}_indices.csv"
|
||
|
||
# 合成完整绝对路径（优先使用 work_dir，其次从 training_csv_path 推导）
|
||
work_dir = getattr(self, 'work_dir', None)
|
||
if work_dir:
|
||
dynamic_output = os.path.join(
|
||
work_dir, "6_water_quality_indices", dynamic_output
|
||
).replace('\\', '/')
|
||
|
||
self.output_file_widget.set_path(dynamic_output)
|
||
|
||
# 获取配置
|
||
config = self.get_config()
|
||
|
||
# 调用GUI的run_single_step方法
|
||
parent = self.parent()
|
||
while parent and not hasattr(parent, 'run_single_step'):
|
||
parent = parent.parent()
|
||
|
||
if parent and hasattr(parent, 'run_single_step'):
|
||
parent.run_single_step('step5_5', {'step5_5': config})
|
||
else:
|
||
QMessageBox.critical(self, "错误", "无法找到父级GUI对象")
|