Compare commits

...

2 Commits

Author SHA1 Message Date
c6e42c3d2f perf: Goodman 水体像素原地校正 + 行段跳跃大幅提速大尺度影像
核心优化 (消除 ~95% 的无效计算):

1. 原地校正 (in-place on water pixels only):
   旧: corrected = R - R_750 + A + B*diff (全图 86M 像素)
        np.where(water, corrected, R) (再分配 329MB)
   新: R[water] = R[water] - R_750[water] + A + B*diff[water]
        (仅水体像素, 零额外分配, 无 np.where)
   效果: 水体占 5% 时, 浮点运算减少 20×, 中间数组消除

2. 行段跳跃 (water row ranges):
   预计算含水行段, 纯陆地行直接跳过不做任何计算
   水域 < 50% 时自动启用 (_find_water_row_ranges)

3. SIMD 友好路径 (无掩膜时):
   np.subtract/add(..., out=R) 替代表达式
   避免临时中间数组, 利用 NumPy SIMD 向量化

预期效果 (6522×13215×150, 水体 10%):
  每波段耗时: ~62s → 估计 ~25-35s (计算部分加速 ~20×)
  IO 仍然占主导 (~15-20s 读+写 329MB)
2026-07-07 14:18:37 +08:00
d920863a0c 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 确保数据及时落盘
2026-07-07 14:10:47 +08:00
2 changed files with 155 additions and 50 deletions

View File

@ -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):
@ -190,46 +212,85 @@ 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 校正处理(性能优化版)。
优化策略:
1. 原地校正: 直接在 R 的水体像素上修改,零副本
2. 仅水体计算: 校正公式仅在 water==True 的像素上执行
3. 无 np.where: 消除创建 329MB 中间数组的 where 操作
4. 行段跳跃: 水域 < 50% 时按含水行段分块处理,跳过纯陆地行
Args:
out_dataset: 提供→流式写入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
# 获取水域掩膜
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):
# 读取当前波段(只加载一个波段到内存)
del R_640 # 释放不再需要的 R_640
# 水域掩膜 + 水体占比检测(控制行段跳跃模式)
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_band = R - R_750 + self.A + self.B * diff_640_750
np.maximum(corrected_band, 0, out=corrected_band)
# 如果存在水域掩膜,只对水域区域应用校正
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
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(R)
out_band.FlushCache()
del R
else:
corrected_bands.append(R)
# 清理
del R_750, diff_640_750
return corrected_bands
def _get_corrected_bands_gdal_mem(self):
@ -339,32 +400,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:

View File

@ -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