fix: Goodman 流式处理 — 逐波段写入磁盘杜绝 OOM
问题: 6522×13215×150 大影像处理到第 142 波段时崩溃 'Unable to allocate 329. MiB for an array' 根因: _get_corrected_bands_gdal 将全部 150 个波段累积在 corrected_bands 列表中 (≈46 GB)，内存耗尽。 修复 (Goodman.py): - _get_corrected_bands_gdal(): 新增 out_dataset 参数， 流式模式下每处理完一个波段立即 WriteArray→FlushCache→del - 新增 _get_corrected_bands_streaming(): 创建输出文件后 调用流式处理，内存峰值 ≈ 3 波段 (NIR×2 + 当前) ≈ 1 GB - get_corrected_bands(): output_path 已设置时自动走流式模式 - 原 _get_corrected_bands_numpy() 和 _gdal_mem() 的 output_path=None 路径保持向后兼容 修复 (glint_removal_step.py): - corrected_bands 为 None 时跳过 _save_bands_as_image (流式模式已直接写入磁盘) 性能微优化: - corrected = R - R_750 → += self.A → += self.B*diff (原地) - del R_640 提早释放 - WriteArray + FlushCache 确保数据及时落盘
This commit is contained in:
duxin
@ -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:
|
||||
|
||||
@ -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
|
||||
|
||||
Reference in New Issue
Block a user