diff --git a/src/core/glint_removal/Goodman.py b/src/core/glint_removal/Goodman.py index 5c6728c..2687947 100644 --- a/src/core/glint_removal/Goodman.py +++ b/src/core/glint_removal/Goodman.py @@ -21,6 +21,28 @@ except ImportError: # 检测是否在 PyInstaller 打包环境(无控制台) _is_frozen_gui = getattr(sys, "frozen", False) and (not hasattr(sys, 'stdout') or sys.stdout is None) + +def _find_water_row_ranges(water_rows): + """将含水行 bool 数组压缩为连续行段列表 [(start, end), ...] + + 例如 [False, True, True, False, True] → [(1,3), (4,5)] + 用于行段跳跃优化:完全无水体的行段直接跳过,不做任何计算。 + """ + ranges = [] + in_range = False + start = 0 + n = len(water_rows) + for i in range(n + 1): + has_water = water_rows[i] if i < n else False + if has_water and not in_range: + start = i + in_range = True + elif not has_water and in_range: + ranges.append((start, i)) + in_range = False + return ranges + + class Goodman: def __init__(self, im_aligned, NIR_lower = 25, NIR_upper = 37, A = 0.000019, B = 0.1, use_gdal=True, chunk_size=None, water_mask=None, output_path=None): @@ -192,15 +214,16 @@ class Goodman: def _get_corrected_bands_gdal(self, out_dataset=None): """ - 使用GDAL逐波段处理。 + 逐波段 GDAL 校正处理(性能优化版)。 + + 优化策略: + 1. 原地校正: 直接在 R 的水体像素上修改,零副本 + 2. 仅水体计算: 校正公式仅在 water==True 的像素上执行 + 3. 无 np.where: 消除创建 329MB 中间数组的 where 操作 + 4. 行段跳跃: 水域 < 50% 时按含水行段分块处理,跳过纯陆地行 Args: - out_dataset: 若提供,每处理完一个波段立即写入此数据集(流式模式); - 若为 None,累积到列表中返回(传统模式,大图像慎用)。 - - Returns: - 若 out_dataset 为 None,返回波段列表; - 若 out_dataset 不为 None,返回 None(波段已直接写入输出文件) + out_dataset: 提供→流式写入;None→累积返回(大图像慎用) """ # 获取NIR波段对象(用于所有波段的校正) band_640 = self.dataset.GetRasterBand(self.NIR_lower + 1) @@ -212,40 +235,58 @@ class Goodman: diff_640_750 = R_640 - R_750 del R_640 # 释放不再需要的 R_640 - # 获取水域掩膜 - water_mask_bool = self.water_mask.astype(bool) if self.water_mask is not None else None + # 水域掩膜 + 水体占比检测(控制行段跳跃模式) + water = self.water_mask.astype(bool) if self.water_mask is not None else None + water_pct = (100.0 * np.count_nonzero(water) / water.size) if water is not None else 100.0 + + # 行段跳跃: 水域占比 < 50% 时启用 + _use_row_skip = (water is not None and water_pct < 50.0) + if _use_row_skip: + water_rows_flag = np.any(water, axis=1) + _water_row_ranges = _find_water_row_ranges(water_rows_flag) + _n_water_rows = int(np.sum(water_rows_flag)) + print(f" [性能] 水域占比 {water_pct:.1f}%,行段跳跃模式:" + f"{len(_water_row_ranges)} 个含水段 (覆盖 {_n_water_rows}/{self.height} 行)") - # 输出模式 streaming = (out_dataset is not None) corrected_bands = [] if not streaming else None # 逐波段处理 for i in tqdm(range(self.n_bands), desc="处理波段 (GDAL)", total=self.n_bands, disable=_is_frozen_gui): - # 读取当前波段 current_band = self.dataset.GetRasterBand(i + 1) - R = current_band.ReadAsArray().astype(np.float32) - # 校正计算(原地操作以减少临时分配) - corrected = R - R_750 - corrected += self.A - corrected += self.B * diff_640_750 - np.maximum(corrected, 0, out=corrected) - - # 水域掩膜:只在有水的地方用校正值,陆地保持原值 - if water_mask_bool is not None: - corrected = np.where(water_mask_bool, corrected, R) + if _use_row_skip: + # ◆ 行段跳跃模式:读全波段 → 仅在水体行段上原地校正 + R = current_band.ReadAsArray().astype(np.float32) + for r0, r1 in _water_row_ranges: + rs, w = slice(r0, r1), water[r0:r1, :] + Rs = R[rs, :] + Rs[w] = (Rs[w] - R_750[rs, :][w] + self.A + + self.B * diff_640_750[rs, :][w]) + np.maximum(Rs[w], 0, out=Rs[w]) + del Rs + else: + R = current_band.ReadAsArray().astype(np.float32) + if water is not None: + # ◆ 全图模式 + 掩膜:原地校正仅水体像素(零额外分配) + R[water] = (R[water] - R_750[water] + self.A + + self.B * diff_640_750[water]) + np.maximum(R[water], 0, out=R[water]) + else: + # ◆ 无掩膜:全图校正(最快 SIMD 路径) + np.subtract(R, R_750, out=R) + np.add(R, self.A, out=R) + np.add(R, self.B * diff_640_750, out=R) + np.maximum(R, 0, out=R) if streaming: - # 流式模式:立即写入磁盘并释放 out_band = out_dataset.GetRasterBand(i + 1) - out_band.WriteArray(corrected) + out_band.WriteArray(R) out_band.FlushCache() - del corrected, R - else: - # 传统模式:累积到列表 - corrected_bands.append(corrected) del R + else: + corrected_bands.append(R) # 清理 del R_750, diff_640_750