Files
WQ_GUI/tmp_watercolor_rescue.py

569 lines
22 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Step9 面板 - 水色指数反演(直接处理去耀斑 BSQ 影像)
将 waterindex.csv 中的公式直接应用于去耀斑高光谱影像,
输出各水质参数指数的 GeoTIFF 栅格图像。
"""
import os
import traceback
from pathlib import Path
from typing import Dict, List, Optional
from PyQt5.QtWidgets import (
QWidget, QVBoxLayout, QHBoxLayout, QGridLayout, QFormLayout,
QGroupBox, QLabel, QLineEdit, QComboBox, QCheckBox, QPushButton,
QFileDialog, QMessageBox, QListWidget, QListWidgetItem,
QAbstractItemView, QProgressBar, QTextEdit, QFrame,
QScrollArea, QSizePolicy,
)
from PyQt5.QtGui import QFont
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from src.gui.components.custom_widgets import FileSelectWidget
from src.gui.styles import ModernStylesheet
class WaterIndexWorker(QThread):
"""后台线程:执行水色指数反演"""
finished_ok = pyqtSignal(dict)
failed = pyqtSignal(str)
progress = pyqtSignal(str, float) # message, percent
log = pyqtSignal(str)
def __init__(
self,
bsq_path: str,
hdr_path: str,
output_dir: str,
selected_formulas: List[str],
waterindex_csv: str,
water_mask_path: Optional[str] = None,
work_dir: Optional[str] = None,
):
super().__init__()
self.bsq_path = bsq_path
self.hdr_path = hdr_path
self.output_dir = output_dir
self.selected_formulas = selected_formulas
self.waterindex_csv = waterindex_csv
self.water_mask_path = water_mask_path
self.work_dir = work_dir
def run(self):
try:
from src.core.algorithms.waterindex_inversion import WaterIndexProcessor
self.progress.emit("正在初始化水色指数处理器…", 2)
processor = WaterIndexProcessor(self.waterindex_csv)
self.progress.emit("正在读取影像元数据…", 5)
# 获取影像元数据
meta = processor.get_image_metadata(self.bsq_path, self.hdr_path)
if not meta:
self.failed.emit("无法读取影像元数据,请检查 BSQ 和 HDR 文件是否匹配")
return
n_bands = meta.get('bands', 0)
wv_range = meta.get('wavelength_range', '未知')
self.log.emit(
f"影像信息: {meta['width']}×{meta['height']} 像素, "
f"{n_bands} 波段, {wv_range}"
)
if self.water_mask_path:
self.log.emit(f"使用水域掩膜: {self.water_mask_path}")
# 使用 run_inversion 入口(含掩膜拦截链路)
results = processor.run_inversion(
deglint_img_path=self.bsq_path,
work_dir=self.work_dir or self.output_dir,
formula_csv_path=self.waterindex_csv,
selected_formulas=self.selected_formulas,
water_mask_path=self.water_mask_path,
callback=self._on_progress,
)
self.progress.emit(f"完成!共生成 {len(results)} 个指数图", 100)
self.finished_ok.emit(results)
except Exception as e:
self.failed.emit(f"{e}\n{traceback.format_exc()}")
def _on_progress(self, msg: str, pct: float):
self.progress.emit(msg, pct)
class Step11WaterColorPanel(QWidget):
"""步骤11水色指数反演直接处理 BSQ 影像)"""
def __init__(self, parent=None):
super().__init__(parent)
self._worker: Optional[WaterIndexWorker] = None
self._waterindex_csv = self._find_waterindex_csv()
self._categories: List[str] = []
self._all_formulas: List[Dict] = []
self._formula_list_widgets: Dict[str, QListWidgetItem] = {}
self.init_ui()
self._load_formulas()
def init_ui(self):
layout = QVBoxLayout()
# ---- 标题 ----
title = QLabel("步骤11水色指数反演高光谱影像直接处理")
title.setFont(QFont("Arial", 12, QFont.Bold))
layout.addWidget(title)
# ---- 说明 ----
hint = QLabel(
"将 waterindex.csv 中的公式直接应用于去耀斑高光谱影像BSQ"
"输出各水质参数指数的 GeoTIFF 栅格图像。"
"指数图可直接用于水质专题图生成。"
)
hint.setWordWrap(True)
hint.setStyleSheet(f"color: {ModernStylesheet.COLORS.get('text_secondary', '#666')};")
layout.addWidget(hint)
# ---- 输入影像选择 ----
input_group = QGroupBox("输入影像")
input_layout = QFormLayout()
self.bsq_file = FileSelectWidget(
"去耀斑 BSQ 影像:",
"BSQ Files (*.bsq);;DAT Files (*.dat);;All Files (*.*)"
)
self.bsq_file.line_edit.setPlaceholderText("选择去耀斑处理后的 BSQ 影像")
self.bsq_file.browse_btn.clicked.disconnect()
self.bsq_file.browse_btn.clicked.connect(self._browse_bsq)
input_layout.addRow("BSQ 影像:", self.bsq_file)
self.hdr_file = FileSelectWidget(
"ENVI 头文件:",
"HDR Files (*.hdr);;All Files (*.*)"
)
self.hdr_file.line_edit.setPlaceholderText("自动关联同路径 .hdr 文件")
self.hdr_file.browse_btn.clicked.disconnect()
self.hdr_file.browse_btn.clicked.connect(self._browse_hdr)
input_layout.addRow("HDR 文件:", self.hdr_file)
# 影像信息显示
self.meta_label = QLabel("未加载影像")
self.meta_label.setStyleSheet(
"background: #f0f0f0; padding: 4px 8px; border-radius: 4px; "
"font-size: 12px; color: #333;"
)
input_layout.addRow("影像信息:", self.meta_label)
input_group.setLayout(input_layout)
layout.addWidget(input_group)
# ---- 公式选择 ----
formula_group = QGroupBox("公式选择")
formula_layout = QGridLayout()
# 类别过滤
formula_layout.addWidget(QLabel("按类别筛选:"), 0, 0)
self.category_combo = QComboBox()
self.category_combo.currentTextChanged.connect(self._on_category_changed)
formula_layout.addWidget(self.category_combo, 0, 1, 1, 2)
# 全选/取消全选
select_btn_layout = QHBoxLayout()
self.select_all_btn = QPushButton("全选")
self.select_all_btn.setMaximumWidth(80)
self.select_all_btn.clicked.connect(self._select_all)
select_btn_layout.addWidget(self.select_all_btn)
self.deselect_all_btn = QPushButton("取消全选")
self.deselect_all_btn.setMaximumWidth(80)
self.deselect_all_btn.clicked.connect(self._deselect_all)
select_btn_layout.addWidget(self.deselect_all_btn)
select_btn_layout.addStretch()
formula_layout.addLayout(select_btn_layout, 0, 3)
# 公式列表
self.formula_list = QListWidget()
self.formula_list.setSelectionMode(QAbstractItemView.MultiSelection)
self.formula_list.setMinimumHeight(200)
self.formula_list.itemChanged.connect(self._on_item_changed)
formula_layout.addWidget(self.formula_list, 1, 0, 1, 4)
formula_group.setLayout(formula_layout)
layout.addWidget(formula_group)
# ---- 输出设置 ----
output_group = QGroupBox("输出设置")
output_layout = QFormLayout()
self.output_dir = FileSelectWidget(
"输出目录:",
"Directories"
)
self.output_dir.line_edit.setPlaceholderText("留空 → 工作目录/8_WaterIndex_Images")
self.output_dir.browse_btn.clicked.disconnect()
self.output_dir.browse_btn.clicked.connect(self._browse_output_dir)
output_layout.addRow("输出目录:", self.output_dir)
self.format_combo = QComboBox()
self.format_combo.addItems(["GTiff (GeoTIFF)", "ENVI", "PCI"])
self.format_combo.setCurrentIndex(0)
output_layout.addRow("输出格式:", self.format_combo)
output_group.setLayout(output_layout)
layout.addWidget(output_group)
# ---- 进度显示 ----
self.progress_bar = QProgressBar()
self.progress_bar.setMinimum(0)
self.progress_bar.setMaximum(100)
self.progress_bar.setValue(0)
self.progress_bar.setTextVisible(True)
layout.addWidget(self.progress_bar)
self.progress_label = QLabel("")
self.progress_label.setStyleSheet("font-size: 11px; color: #666;")
layout.addWidget(self.progress_label)
# ---- 启用 & 运行 ----
self.enable_checkbox = QCheckBox("启用此步骤")
self.enable_checkbox.setChecked(True)
layout.addWidget(self.enable_checkbox)
self.run_btn = QPushButton("▶ 执行水色指数反演")
self.run_btn.setStyleSheet(ModernStylesheet.get_button_stylesheet('success'))
self.run_btn.clicked.connect(self.run_step)
layout.addWidget(self.run_btn)
layout.addStretch()
self.setLayout(layout)
def _find_waterindex_csv(self) -> str:
"""查找 waterindex.csv 路径"""
candidates = [
Path(__file__).parent.parent.parent / "model" / "waterindex.csv",
Path(__file__).parent.parent.parent.parent / "src" / "gui" / "model" / "waterindex.csv",
]
for c in candidates:
if c.exists():
return str(c)
return ""
def _load_formulas(self):
"""加载 waterindex.csv 中的公式"""
if not self._waterindex_csv or not Path(self._waterindex_csv).exists():
self.meta_label.setText("⚠️ waterindex.csv 未找到")
return
import csv
self._all_formulas = []
try:
with open(self._waterindex_csv, 'r', encoding='utf-8-sig') as f:
reader = csv.DictReader(f)
self._all_formulas = list(reader)
except Exception as e:
self.meta_label.setText(f"⚠️ 加载公式失败: {e}")
return
# 提取所有类别
cats = set()
for f in self._all_formulas:
c = f.get('Category', '').strip()
if c:
cats.add(c)
self._categories = sorted(cats)
self.category_combo.clear()
self.category_combo.addItem("全部")
self.category_combo.addItems(self._categories)
self._populate_list("全部")
def _populate_list(self, category: str):
"""根据类别填充公式列表"""
self.formula_list.clear()
self._formula_list_widgets.clear()
formulas_to_show = (
[f for f in self._all_formulas if f.get('Category', '') == category]
if category != "全部"
else self._all_formulas
)
for f in formulas_to_show:
name = f.get('Formula_Name', '')
formula_str = f.get('Formula', '')
cat = f.get('Category', '')
ftype = f.get('Formula_Type', '')
item = QListWidgetItem()
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
item.setCheckState(Qt.Checked)
item.setData(Qt.UserRole, name)
item.setText(
f"{name} [{cat}] ({ftype})\n {formula_str}"
)
item.setToolTip(f"{name}\n{category}\n{formula_str}")
self.formula_list.addItem(item)
self._formula_list_widgets[name] = item
def _on_category_changed(self, category: str):
self._populate_list(category)
def _select_all(self):
for item in self.formula_list.selectedItems():
item.setCheckState(Qt.Checked)
# 也全选当前显示的
for i in range(self.formula_list.count()):
it = self.formula_list.item(i)
it.setCheckState(Qt.Checked)
def _deselect_all(self):
for i in range(self.formula_list.count()):
it = self.formula_list.item(i)
it.setCheckState(Qt.Unchecked)
def _on_item_changed(self, item: QListWidgetItem):
pass # 可扩展:实时统计选中数量
def _browse_bsq(self):
path, _ = QFileDialog.getOpenFileName(
self, "选择去耀斑 BSQ 影像",
"",
"BSQ Files (*.bsq);;DAT Files (*.dat);;All Files (*.*)"
)
if path:
self.bsq_file.set_path(path)
# 自动关联同路径 hdr
hdr = Path(path).with_suffix('.hdr')
if hdr.exists():
self.hdr_file.set_path(str(hdr))
self._load_metadata(path, str(hdr) if hdr.exists() else "")
def _browse_hdr(self):
path, _ = QFileDialog.getOpenFileName(
self, "选择 ENVI 头文件",
"",
"HDR Files (*.hdr);;All Files (*.*)"
)
if path:
self.hdr_file.set_path(path)
bsq_path = self.bsq_file.get_path()
if bsq_path:
self._load_metadata(bsq_path, path)
def _browse_output_dir(self):
d = QFileDialog.getExistingDirectory(self, "选择输出目录", "")
if d:
self.output_dir.set_path(d)
def _load_metadata(self, bsq_path: str, hdr_path: str):
"""加载并显示影像元数据"""
if not bsq_path or not Path(bsq_path).exists():
self.meta_label.setText("⚠️ 影像文件不存在")
return
if not hdr_path or not Path(hdr_path).exists():
self.meta_label.setText("⚠️ 头文件不存在")
return
try:
from src.core.algorithms.waterindex_inversion import WaterIndexProcessor
processor = WaterIndexProcessor(self._waterindex_csv)
meta = processor.get_image_metadata(bsq_path, hdr_path)
if meta:
self.meta_label.setText(
f"{meta['width']}×{meta['height']} | "
f"{meta['bands']} 波段 | {meta.get('wavelength_range', '未知')} | "
f"驱动: {meta['driver']}"
)
else:
self.meta_label.setText("⚠️ 无法读取元数据")
except Exception as e:
self.meta_label.setText(f"⚠️ 元数据读取失败: {e}")
def _get_selected_formula_names(self) -> List[str]:
names = []
for i in range(self.formula_list.count()):
item = self.formula_list.item(i)
if item.checkState() == Qt.Checked:
name = item.data(Qt.UserRole)
if name:
names.append(name)
return names
def _get_default_work_dir(self) -> str:
if hasattr(self, 'work_dir') and self.work_dir:
return str(self.work_dir)
mw = self.window()
if mw and hasattr(mw, 'work_dir') and mw.work_dir:
return str(mw.work_dir)
return ""
def get_config(self) -> dict:
bsq = self.bsq_file.get_path()
return {
'bsq_path': bsq,
'hdr_path': self.hdr_file.get_path(),
'deglint_img_path': bsq,
'output_dir': self.output_dir.get_path(),
'output_format': self.format_combo.currentText().split()[0],
'selected_formulas': self._get_selected_formula_names(),
}
def set_config(self, config: dict):
if config.get('bsq_path'):
self.bsq_file.set_path(config['bsq_path'])
if config.get('hdr_path'):
self.hdr_file.set_path(config['hdr_path'])
if config.get('output_dir'):
self.output_dir.set_path(config['output_dir'])
if 'selected_formulas' in config:
names = set(config['selected_formulas'])
for i in range(self.formula_list.count()):
item = self.formula_list.item(i)
name = item.data(Qt.UserRole)
item.setCheckState(Qt.Checked if name in names else Qt.Unchecked)
def update_from_config(self, work_dir=None, pipeline=None):
if work_dir:
self.work_dir = work_dir
elif hasattr(self, 'work_dir') and self.work_dir:
pass
else:
self.work_dir = None
main_window = self.window()
# 自动填入去耀斑影像
if main_window and hasattr(main_window, 'step3_panel'):
deglint_path = main_window.step3_panel.output_file.get_path()
if deglint_path and not self.bsq_file.get_path():
if not os.path.isabs(deglint_path):
deglint_path = os.path.join(self.work_dir or '', deglint_path).replace('\\', '/')
self.bsq_file.set_path(deglint_path)
hdr = Path(deglint_path).with_suffix('.hdr')
if hdr.exists():
self.hdr_file.set_path(str(hdr))
self._load_metadata(deglint_path, str(hdr))
# 自动填入输出目录
if self.work_dir:
out_dir = os.path.join(self.work_dir, "8_WaterIndex_Images").replace('\\', '/')
os.makedirs(out_dir, exist_ok=True)
if not self.output_dir.get_path():
self.output_dir.set_path(out_dir)
def run_step(self):
bsq_path = self.bsq_file.get_path().strip()
hdr_path = self.hdr_file.get_path().strip()
output_dir = self.output_dir.get_path().strip()
# 验证输入
if not bsq_path:
QMessageBox.warning(self, "输入错误", "请选择去耀斑 BSQ 影像!")
return
if not Path(bsq_path).exists():
QMessageBox.warning(self, "输入错误", f"BSQ 影像不存在:\n{bsq_path}")
return
if not hdr_path:
# 尝试自动查找
auto_hdr = Path(bsq_path).with_suffix('.hdr')
if auto_hdr.exists():
hdr_path = str(auto_hdr)
self.hdr_file.set_path(hdr_path)
else:
QMessageBox.warning(self, "输入错误", "请选择 ENVI 头文件!")
return
if not Path(hdr_path).exists():
QMessageBox.warning(self, "输入错误", f"HDR 文件不存在:\n{hdr_path}")
return
if not output_dir:
work_dir = self._get_default_work_dir()
output_dir = os.path.join(work_dir, "8_WaterIndex_Images").replace('\\', '/')
os.makedirs(output_dir, exist_ok=True)
self.output_dir.set_path(output_dir)
selected = self._get_selected_formula_names()
if not selected:
QMessageBox.warning(self, "输入错误", "请至少选择一个公式!")
return
if self._waterindex_csv and not Path(self._waterindex_csv).exists():
QMessageBox.warning(self, "配置错误", f"waterindex.csv 不存在:\n{self._waterindex_csv}")
return
# ── 自动扫描工作目录下的水域掩膜文件 ────────────────────────────
work_dir = self.work_dir or str(Path(bsq_path).parent)
mask_dir = os.path.join(work_dir, "1_water_mask")
water_mask_path: Optional[str] = None
if os.path.isdir(mask_dir):
# ★★★ glob 智能扫描:取任意 .dat 或 .tif 文件 ★★★
for pattern in ("*.dat", "*.tif", "*.TIF", "*.DT"):
candidates = sorted(Path(mask_dir).glob(pattern))
if candidates:
water_mask_path = str(candidates[0])
break
if water_mask_path:
print(f"[Step8] 自动找到水域掩膜: {water_mask_path}")
else:
print(f"[Step8] 未找到水域掩膜,跳过陆地剔除(陆地将保留在指数图中)")
# 开始后台处理
self.run_btn.setEnabled(False)
self.progress_bar.setValue(0)
self.progress_label.setText("")
self._worker = WaterIndexWorker(
bsq_path=bsq_path,
hdr_path=hdr_path,
output_dir=output_dir,
selected_formulas=selected,
waterindex_csv=self._waterindex_csv,
water_mask_path=water_mask_path,
work_dir=work_dir,
)
self._worker.progress.connect(self._on_progress)
self._worker.finished_ok.connect(self._on_finished)
self._worker.failed.connect(self._on_failed)
self._worker.log.connect(lambda m: self.progress_label.setText(m))
self._worker.start()
def _on_progress(self, msg: str, pct: float):
self.progress_bar.setValue(int(pct))
self.progress_label.setText(msg)
def _on_finished(self, results: Dict[str, str]):
self.run_btn.setEnabled(True)
n = len(results)
QMessageBox.information(
self, "执行成功",
f"水色指数反演完成!\n"
f"共生成 {n} 个指数图GeoTIFF\n\n"
f"输出目录: {self.output_dir.get_path()}"
)
main_window = self.window()
if main_window and hasattr(main_window, 'log_message'):
main_window.log_message(f"步骤8水色指数反演完成生成 {n} 个指数图", "info")
def _on_failed(self, err: str):
self.run_btn.setEnabled(True)
self.progress_bar.setValue(0)
QMessageBox.critical(self, "执行错误", f"水色指数反演失败:\n\n{err[:500]}")
def get_output_dir(self) -> str:
return self.output_dir.get_path().strip() or ""
def get_output_tif_paths(self) -> List[str]:
"""获取输出目录下的所有 GeoTIFF 文件路径"""
out_dir = self.get_output_dir()
if not out_dir or not os.path.isdir(out_dir):
return []
return sorted(
str(p) for p in Path(out_dir).glob("*.tif")
if p.is_file()
)