@ -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 )
# 校正计算
co rrected _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 = cu rrent _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 :