diff --git a/src/core/glint_removal/Goodman.py b/src/core/glint_removal/Goodman.py index 5a43910..5c6728c 100644 --- a/src/core/glint_removal/Goodman.py +++ b/src/core/glint_removal/Goodman.py @@ -190,46 +190,66 @@ class Goodman: corrected_bands.append(corrected_band) return corrected_bands - def _get_corrected_bands_gdal(self): + def _get_corrected_bands_gdal(self, out_dataset=None): """ - 使用GDAL逐波段处理，直接处理整个波段（不分块） - - 内存峰值 = NIR波段(2个) + 当前处理的波段(1个) + 已处理的波段(累积在列表中) + 使用GDAL逐波段处理。 + + Args: + out_dataset: 若提供，每处理完一个波段立即写入此数据集（流式模式）； + 若为 None，累积到列表中返回（传统模式，大图像慎用）。 + + Returns: + 若 out_dataset 为 None，返回波段列表； + 若 out_dataset 不为 None，返回 None（波段已直接写入输出文件） """ - corrected_bands = [] - # 获取NIR波段对象（用于所有波段的校正） - band_640 = self.dataset.GetRasterBand(self.NIR_lower + 1) # GDAL波段从1开始 + band_640 = self.dataset.GetRasterBand(self.NIR_lower + 1) band_750 = self.dataset.GetRasterBand(self.NIR_upper + 1) - - # 先读取NIR波段（用于所有波段的校正，会一直保存在内存中） + + # 读取NIR波段（全程保留在内存中） R_640 = band_640.ReadAsArray().astype(np.float32) R_750 = band_750.ReadAsArray().astype(np.float32) 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 - - # 逐波段处理：每次只读取和处理一个波段 - for i in tqdm(range(self.n_bands), desc="处理波段 (GDAL)", total=self.n_bands, disable=_is_frozen_gui): - # 读取当前波段（只加载一个波段到内存） + + # 输出模式 + 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_band = R - R_750 + self.A + self.B * diff_640_750 - np.maximum(corrected_band, 0, out=corrected_band) - - # 如果存在水域掩膜，只对水域区域应用校正 + + # 校正计算（原地操作以减少临时分配） + 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_band = np.where(water_mask_bool, corrected_band, R) - - # 添加到结果列表（corrected_band会保留在列表中） - corrected_bands.append(corrected_band) - - # 释放当前波段数据（显式删除有助于及时释放内存） - del R - + corrected = np.where(water_mask_bool, corrected, R) + + if streaming: + # 流式模式：立即写入磁盘并释放 + out_band = out_dataset.GetRasterBand(i + 1) + out_band.WriteArray(corrected) + out_band.FlushCache() + del corrected, R + else: + # 传统模式：累积到列表 + corrected_bands.append(corrected) + del R + + # 清理 + del R_750, diff_640_750 + return corrected_bands def _get_corrected_bands_gdal_mem(self): @@ -339,32 +359,75 @@ class Goodman: def get_corrected_bands(self): """ - 获取校正后的波段 - 根据输入类型和大小自动选择最优处理方法 - - :return: 校正后的波段列表 + 获取校正后的波段（自动选择最优处理模式） + + 内存优化: 当 output_path 已设置时，使用流式模式逐波段直接写入磁盘， + 避免在内存中累积全部波段（对 6522×13215×150 的大影像可节省 ~46 GB）。 + + :return: 校正后的波段列表（流式模式返回 None，波段已在输出文件中） """ - # 如果输入是文件路径，使用GDAL直接读取 + # ── 流式模式：output_path 已设置 → 逐波段处理+立即写入 → 零累积 ── + if self.output_path is not None: + return self._get_corrected_bands_streaming() + + # ── 传统模式：output_path 为空 → 返回波段列表 ── if self.is_file_path: if self.use_gdal: - corrected_bands = self._get_corrected_bands_gdal() + return self._get_corrected_bands_gdal(out_dataset=None) else: raise ValueError("输入为文件路径时，必须安装GDAL") else: - # 如果输入是numpy数组 if self.use_gdal and self.height * self.width * self.n_bands > 100000000: - # 大图像使用GDAL内存驱动逐波段处理 - corrected_bands = self._get_corrected_bands_gdal_mem() + return self._get_corrected_bands_gdal_mem() else: - # 小图像使用numpy直接处理 - corrected_bands = self._get_corrected_bands_numpy() - - # 如果提供了输出路径，保存结果 - if self.output_path is not None: - self._save_corrected_bands(corrected_bands) - - return corrected_bands + return self._get_corrected_bands_numpy() + def _get_corrected_bands_streaming(self): + """流式处理：逐波段校正并直接写入输出文件，不累积内存 + + 适用于大尺度影像（如 6522×13215×150）。 + 内存峰值 ≈ 3 个全波段数组（NIR×2 + 当前波段），而非全部 150 个波段。 + + :return: None（波段已在输出文件中） + """ + import os + + # ── 创建输出文件 ── + base_path, ext = os.path.splitext(self.output_path) + bsq_path = base_path + '.bsq' if ext.lower() != '.bsq' else self.output_path + output_dir = os.path.dirname(bsq_path) + if output_dir and not os.path.exists(output_dir): + os.makedirs(output_dir, exist_ok=True) + + driver = gdal.GetDriverByName('ENVI') + out_ds = driver.Create( + bsq_path, self.width, self.height, self.n_bands, gdal.GDT_Float32 + ) + if out_ds is None: + raise ValueError(f"无法创建输出文件: {bsq_path}") + + # ── 设置地理参考 ── + if self.is_file_path and self.dataset is not None: + out_ds.SetGeoTransform(self.dataset.GetGeoTransform()) + out_ds.SetProjection(self.dataset.GetProjection()) + + try: + # 逐波段处理 + 立即写入（波段在 _get_corrected_bands_gdal 的循环中 + # 由 WriteArray 写入 → FlushCache → del，不会累积） + self._get_corrected_bands_gdal(out_dataset=out_ds) + finally: + out_ds = None # 关闭文件，确保数据落盘 + + # ── 日志 ── + hdr_path = bsq_path + '.hdr' + if os.path.exists(hdr_path): + print(f"校正后的图像已保存至: {bsq_path} (BSQ格式, 流式写入)") + else: + print(f"校正后的图像已保存至: {bsq_path} (BSQ格式)") + print("警告: 未检测到.hdr文件，但GDAL应该已自动创建") + + return None + def __del__(self): """清理资源""" if self.dataset is not None and self.is_file_path: diff --git a/src/core/steps/glint_removal_step.py b/src/core/steps/glint_removal_step.py index daffacc..960568e 100644 --- a/src/core/steps/glint_removal_step.py +++ b/src/core/steps/glint_removal_step.py @@ -300,7 +300,8 @@ class GlintRemovalStep: ) corrected_bands = goodman.get_corrected_bands() - if not Path(hardcoded_bsq).exists(): + # 流式模式（大图）：波段已由 Goodman 直接写入磁盘，返回 None + if corrected_bands is not None and not Path(hardcoded_bsq).exists(): _save_bands_as_image(corrected_bands, hardcoded_bsq, geotransform, projection) _copy_hdr_info(img_path, hardcoded_bsq) del corrected_bands