#!/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("留空 → 工作目录/10_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, "10_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, "10_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() )