duxingchen/WQ_GUI
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf: 水域掩膜改为分块生成 + GDAL 原生栅格化，消除全幅驻留

Browse Source 
ndwi() 走新增的 generate_water_mask_chunked：按块读取 float32 计算 NDWI 并直接写盘，替代 calculate_NDWI + extract_water 的全幅浮点流程（峰值内存 10GB+）。分母为 0 的像元置 NaN，保持与原 extract_water 一致的 ignore 语义。

rasterize_shp 重写为 GDAL RasterizeLayer 原生烧录，消灭 Python 双重循环。因直接烧入 ENVI driver 会静默失败，ENVI(.dat/.bil/.bsq) 输出先烧到临时 GTiff 载体再 CreateCopy 转格式。

掩膜输出 dtype 收紧为 GDT_Byte（二值省 4 倍空间），格式跟随扩展名：.tif/.tiff 写 GTiff+LZW，其余沿用项目 ENVI 惯例。

保留 calculate_NDWI / extract_water 兼容层并标注废弃，仅供旧脚本直接 import。
This commit is contained in:
duxin
2026-09-15 10:34:39 +08:00
parent abc8fe5694
commit 32832f122d
1 changed files with 251 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files

362
src/utils/extract_water_area.py
View File

@ -1,9 +1,47 @@
import os
import numpy as np
from src.utils.util import *
from osgeo import gdal, ogr
import argparse
gdal.UseExceptions()
ogr.UseExceptions()
# ============================================================
# 掩膜输出辅助：格式跟随输出扩展名
#
# 历史实现用 util.write_bands() 输出 ENVI(.dat/.bil/.bsq) + Float32。
# 本模块改造后掩膜都是 0/1 二值，dtype 收紧为 GDT_Byte（省 4×），
# 并保留项目内 .dat→ENVI 的格式惯例；只有当输出路径是 .tif/.tiff 时才
# 写 GTiff（带 LZW 压缩）。
# ============================================================
def _open_mask_output(out_path, cols, rows):
"""创建与 out_path 扩展名匹配的 1 波段 uint8 掩膜数据集。
Returns:
(gdal.Dataset, gdal.Band)
"""
ext = os.path.splitext(out_path)[1].lower()
if ext in ('.tif', '.tiff'):
driver = gdal.GetDriverByName('GTiff')
ds = driver.Create(out_path, cols, rows, 1, gdal.GDT_Byte,
options=['COMPRESS=LZW'])
else:
driver = gdal.GetDriverByName('ENVI')
ds = driver.Create(out_path, cols, rows, 1, gdal.GDT_Byte)
return ds, ds.GetRasterBand(1)
def _remove_raster_output(path):
"""移除已有栅格输出（兼容 GTiff / ENVI .dat+.hdr / 无扩展名 .hdr）。"""
base = os.path.splitext(path)[0]
candidates = [path, path + '.hdr', base + '.hdr', base + '.dat', base + '.aux.xml']
for c in dict.fromkeys(candidates):
if c and os.path.exists(c):
os.remove(c)
def xml2shp():
pass
@ -14,18 +52,34 @@ def rasterize_envi_xml(shp_filepath):
@timeit
def rasterize_shp(shp_filepath, raster_fn_out, img_path, NoData_value=None):
# ---------- 防御性处理：路径标准化 ----------
"""将 shapefile 栅格化并对齐到 img_path 的像元网格。
★ v3 重写：用 GDAL 原生 RasterizeLayer 直接把矢量烧录到影像网格，
彻底消灭原实现中 O(height*width) 的纯 Python 双重 for 循环 + 逐像素
坐标反算（6.7 亿像素时从小时级降到秒级）。
接口（签名/参数顺序/返回）与原 rasterize_shp 完全一致，三个调用方
（water_mask_step / mask_converter / CLI main）无需任何改动。
语义说明：
- 掩膜前景烧录值 = 1（uint8），背景 = NoData_value（默认 0）。
- ALL_TOUCHED=TRUE：保留原代码对边界像素的烧录倾向。
- 要求 shp 与影像处于同一投影坐标系（与原实现假设一致）。
"""
# ---------- 防御性处理 ----------
shp_filepath = os.path.abspath(shp_filepath).replace('\\', '/')
print(f"[DEBUG rasterize_shp] 标准化后的 SHP 路径: {shp_filepath}")
if not os.path.exists(shp_filepath):
raise FileNotFoundError(f"Shapefile 不存在: {shp_filepath}")
if not os.path.exists(img_path):
raise FileNotFoundError(f"参考影像不存在: {img_path}")
# 检查伴随文件完整性
shp_base = os.path.splitext(shp_filepath)[0]
for ext in ['.dbf', '.shx', '.prj']:
companion = shp_base + ext
if os.path.exists(companion):
print(f"[DEBUG rasterize_shp] 伴随文件存在: {companion}")
else:
print(f"[WARNING rasterize_shp] 伴随文件缺失: {companion}")
if not os.path.exists(shp_base + ext):
print(f"[WARNING rasterize_shp] 伴随文件缺失: {shp_base + ext}")
# 确保 GDAL/OGR 驱动已注册
gdal.AllRegister()
@ -39,107 +93,86 @@ def rasterize_shp(shp_filepath, raster_fn_out, img_path, NoData_value=None):
)
print(f"[DEBUG rasterize_shp] ESRI Shapefile 驱动: {driver.GetName()}")
# 打开参考影像获取尺寸信息
dataset = gdal.Open(img_path)
if dataset is None:
# ---------- 1. 读取参考影像的空间信息（决定输出网格） ----------
ref_ds = gdal.Open(img_path)
if ref_ds is None:
raise ValueError(f"无法打开参考影像文件: {img_path}")
im_width = dataset.RasterXSize
im_height = dataset.RasterYSize
geotransform = dataset.GetGeoTransform()
imgdata_in = dataset.GetRasterBand(1).ReadAsArray()
del dataset
cols = ref_ds.RasterXSize
rows = ref_ds.RasterYSize
geo_transform = ref_ds.GetGeoTransform()
projection = ref_ds.GetProjection()
ref_ds = None
# ---------- 打开 SHP 文件（双重尝试获取详细错误） ----------
source_ds = gdal.OpenEx(shp_filepath, gdal.OF_VECTOR)
if source_ds is None:
# gdal.OpenEx 失败，尝试 ogr.Open 获取更详细的错误信息
try:
ogr_ds = ogr.Open(shp_filepath)
except Exception as ogr_err:
raise RuntimeError(
f"GDAL/OGR 无法打开 SHP 文件（详细原因）:\n"
f" ogr.Open 抛出异常: {str(ogr_err)}\n"
f" 文件路径: {shp_filepath}\n"
f"常见原因:\n"
f" 1. 路径包含中文/空格/特殊字符（建议复制到纯英文路径下重试）\n"
f" 2. .dbf 或 .shx 伴随文件缺失或损坏\n"
f" 3. GDAL 未注册 ESRI Shapefile 驱动\n"
f" 4. 文件被其他程序锁定"
)
if ogr_ds is None:
raise RuntimeError(
f"ogr.Open 和 gdal.OpenEx 均返回 None，无法打开 SHP 文件。\n"
f"文件路径: {shp_filepath}\n"
f"请检查:\n"
f" 1. 所有伴随文件（.dbf/.shx/.prg）是否齐全\n"
f" 2. 文件是否被其他程序占用\n"
f" 3. 路径中是否存在不支持的字符"
)
# 检查图层数量，如果有多层，指定使用第一层
layer_count = source_ds.GetLayerCount()
layer_name = None
if layer_count > 1:
print(f"警告: shapefile包含{layer_count}个图层，将使用第一个图层进行栅格化")
# 获取第一个图层
layer = source_ds.GetLayer(0)
layer_name = layer.GetName()
# 计算像素分辨率（考虑旋转参数）
# 对于旋转影像（gt[2] 或 gt[4] != 0），像素实际分辨率需用勾股定理
pixel_size_x = np.sqrt(geotransform[1]**2 + geotransform[2]**2)
pixel_size_y = np.sqrt(geotransform[4]**2 + geotransform[5]**2)
raster_fn_out_tmp = append2filename(raster_fn_out, "_tmp_delete")
# 构建栅格化参数
rasterize_kwargs = {
'format': 'envi',
'outputType': gdal.GDT_Byte,
'noData': NoData_value,
'initValues': NoData_value,
'xRes': pixel_size_x,
'yRes': pixel_size_y,
'allTouched': True,
'burnValues': 1
}
# 如果有多层，指定使用第一层
if layer_name is not None:
rasterize_kwargs['layers'] = [layer_name]
# 执行栅格化
gdal.Rasterize(raster_fn_out_tmp, source_ds, **rasterize_kwargs)
print(f"[DEBUG rasterize_shp] 目标网格: {cols} x {rows}")
dataset_tmp = gdal.Open(raster_fn_out_tmp)
geotransform_tmp = dataset_tmp.GetGeoTransform()
inv_geotransform_tmp = gdal.InvGeoTransform(geotransform_tmp)
data_tmp = dataset_tmp.GetRasterBand(1).ReadAsArray()
del dataset_tmp
# ---------- 2. 创建掩膜载体（GTiff） ----------
out_dir = os.path.dirname(os.path.abspath(raster_fn_out))
if out_dir and not os.path.exists(out_dir):
os.makedirs(out_dir, exist_ok=True)
# 创建和输入影像相同行列号、相同分辨率的水域掩膜
# ★ v2: 移除 imgdata_in == 0 的误杀逻辑
# 高光谱影像第一波段(380nm)水体反射率接近0 → 大量水体像素被当成"影像外"跳过
# 像素坐标越界检查 (coor_pixel bounds) 已能防止 SHP 超出影像范围的问题
water_mask = np.zeros((im_height, im_width))
for row in range(im_height):
for column in range(im_width):
coor = gdal.ApplyGeoTransform(geotransform, column, row)
coor_pixel = gdal.ApplyGeoTransform(inv_geotransform_tmp, coor[0], coor[1])
coor_pixel = [int(num) for num in coor_pixel]
out_ext = os.path.splitext(raster_fn_out)[1].lower()
is_gtiff = out_ext in ('.tif', '.tiff')
if coor_pixel[0] < 0 or coor_pixel[0] >= data_tmp.shape[1]:
continue
if coor_pixel[1] < 0 or coor_pixel[1] >= data_tmp.shape[0]:
continue
# ★ GDAL RasterizeLayer 写入 GTiff 是可靠路径；直接烧入 ENVI driver
# 目标会静默失败（不抛异常但不写入）。因此 ENVI(.dat/.bil/.bsq) 输出
# 先烧到临时 GTiff 载体，再通过 CreateCopy 转成 ENVI（分块 I/O）。
burn_path = raster_fn_out if is_gtiff else raster_fn_out + '__raster_tmp.tif'
burn_ds = gdal.GetDriverByName('GTiff').Create(
burn_path, cols, rows, 1, gdal.GDT_Byte,
options=['COMPRESS=LZW'])
burn_ds.SetGeoTransform(geo_transform)
burn_ds.SetProjection(projection)
water_mask[row, column] = data_tmp[coor_pixel[1], coor_pixel[0]]
background = 0 if NoData_value is None else NoData_value
burn_band = burn_ds.GetRasterBand(1)
burn_band.Fill(background)
write_bands(img_path, raster_fn_out, water_mask)
# ---------- 3. 载入矢量图层 ----------
shp_ds = ogr.Open(shp_filepath)
if shp_ds is None:
burn_ds = None
raise RuntimeError(f"无法打开 Shapefile: {shp_filepath}")
os.remove(raster_fn_out_tmp)
if shp_ds.GetLayerCount() > 1:
print(f"警告: shapefile包含{shp_ds.GetLayerCount()}个图层，将使用第一个图层进行栅格化")
layer = shp_ds.GetLayer(0)
else:
layer = shp_ds.GetLayer()
if layer is None:
burn_ds = None
raise RuntimeError(f"Shapefile 中没有可用图层: {shp_filepath}")
# ---------- 4. GDAL 原生栅格化（C++ 实现，消灭 Python 双重循环） ----------
gdal.RasterizeLayer(burn_ds, [1], layer, burn_values=[1],
options=["ALL_TOUCHED=TRUE"])
burn_band.FlushCache()
layer = None
shp_ds = None
# ---------- 5. 落盘为最终格式 ----------
burn_ds = None
if not is_gtiff:
_remove_raster_output(raster_fn_out) # 清掉旧 .dat/.hdr，保证可重跑
final_ds = gdal.GetDriverByName('ENVI').CreateCopy(
raster_fn_out, gdal.Open(burn_path, gdal.GA_ReadOnly))
final_ds = None
if os.path.exists(burn_path):
os.remove(burn_path)
print(f"[DEBUG rasterize_shp] 完成: {raster_fn_out}")
# ============================================================
# 兼容层（已废弃）：calculate_NDWI / extract_water
#
# 这两个旧函数会把整幅影像分别读入内存（浮点全幅），内存峰值可达
# 10 GB+。新的 ndwi() 已改用分块的 generate_water_mask_chunked。
# 此处保留仅为向后兼容旧脚本的直接 import/调用，主流程不再使用。
# ============================================================
def calculate_NDWI(green_bandnumber, nir_bandnumber, filename):
"""[已废弃，勿用于大影像] 整幅读取计算 NDWI，返回 float32 全幅数组。"""
dataset = gdal.Open(filename) # 打开文件
num_bands = dataset.RasterCount # 栅格矩阵的波段数
im_geotrans = dataset.GetGeoTransform() # 仿射矩阵
@ -158,44 +191,151 @@ def calculate_NDWI(green_bandnumber, nir_bandnumber, filename):
def extract_water(ndwi, threshold=0.3, data_ignore_value=0):
"""[已废弃，勿用于大影像] 全幅阈值判断生成水体掩膜。"""
water_region = np.where(ndwi > threshold, 1, data_ignore_value)
return water_region
# ============================================================
# NDWI 分块掩膜生成（v3 核心，替代 calculate_NDWI + extract_water 全幅流程）
# ============================================================
def generate_water_mask_chunked(img_path, out_mask_path,
green_band_idx, nir_band_idx,
threshold=0.3, data_ignore_value=0,
block_size=2048):
"""
分块滑动窗口计算 NDWI 并直接写出 uint8 水体掩膜。
与原 calculate_NDWI + extract_water 相比：
- 不再整幅读入波段（原峰值 ≈ 3~4 个全幅 float32 + int64 中间量 >13 GB）
- 每块只驻留 block_size² 的数据，内存恒定在几十 MB
- 输出为 uint8（0/1），格式跟随 out_mask_path 扩展名
Args:
img_path: 输入影像路径
out_mask_path: 输出掩膜路径（.dat/.bil/.bsq → ENVI；.tif/.tiff → GTiff）
green_band_idx: 绿光波段序号（0-based）
nir_band_idx: 近红外波段序号（0-based）
threshold: NDWI 阈值，大于该值判为水体（1）
data_ignore_value: 非水体像素填充值（背景，默认 0）
block_size: 分块边长（像素）
Returns:
输出的掩膜文件路径
"""
ds_in = gdal.Open(img_path)
if ds_in is None:
raise ValueError(f"无法打开影像文件: {img_path}")
cols = ds_in.RasterXSize
rows = ds_in.RasterYSize
band_green = ds_in.GetRasterBand(green_band_idx + 1)
band_nir = ds_in.GetRasterBand(nir_band_idx + 1)
out_dir = os.path.dirname(os.path.abspath(out_mask_path))
if out_dir and not os.path.exists(out_dir):
os.makedirs(out_dir, exist_ok=True)
ds_out, band_out = _open_mask_output(out_mask_path, cols, rows)
ds_out.SetGeoTransform(ds_in.GetGeoTransform())
ds_out.SetProjection(ds_in.GetProjection())
n_blocks = 0
for y in range(0, rows, block_size):
y_size = min(block_size, rows - y)
for x in range(0, cols, block_size):
x_size = min(block_size, cols - x)
# 只读当前块（float32）
g_data = band_green.ReadAsArray(x, y, x_size, y_size).astype(np.float32)
n_data = band_nir.ReadAsArray(x, y, x_size, y_size).astype(np.float32)
# 块级 NDWI。
# ★ 分母为 0 的像元置 NaN：使 np.where(ndwi > threshold) 判为 False，
# 落入 data_ignore_value —— 与原 extract_water 对除零像元的语义一致。
with np.errstate(divide='ignore', invalid='ignore'):
denominator = g_data + n_data
ndwi = np.where(denominator == 0, np.nan,
(g_data - n_data) / denominator)
# 块级阈值判断，直接转 uint8
mask_block = np.where(ndwi > threshold,
np.uint8(1), np.uint8(data_ignore_value)).astype(np.uint8)
# 写回磁盘
band_out.WriteArray(mask_block, x, y)
del g_data, n_data, denominator, ndwi, mask_block
n_blocks += 1
print(f"[generate_water_mask_chunked] 完成: 影像 {cols}x{rows}, "
f"{n_blocks} 个分块, 阈值={threshold}, 输出={out_mask_path}")
band_out.FlushCache()
ds_out = None
ds_in = None
return out_mask_path
@timeit
def ndwi(file_path, ndwi_threshold=0.4, output_path=None,
data_ignore_value=0, sieve_threshold=20):
"""基于 NDWI 生成水体掩膜（分块实现，内存安全）。
接口与原 ndwi 完全一致。流程：
1) 按波长自动定位绿/近红外波段（find_band_number，0-based）
2) generate_water_mask_chunked 分块计算并写出 uint8 掩膜
3) 若 sieve_threshold>0：先在 GTiff 载体上用 gdal.SieveFilter 去除面积
< sieve_threshold 的孤立碎斑，再 CreateCopy 转为最终格式
（GDAL 的原地更新在 ENVI driver 上会静默失败，故需 GTiff 载体）。
"""
if output_path is None:
output_path = append2filename(file_path, "_waterarea")
dataset_in = gdal.Open(file_path)
im_width_in = dataset_in.RasterXSize
im_height_in = dataset_in.RasterYSize
num_bands_in = dataset_in.RasterCount
geotrans_in = dataset_in.GetGeoTransform()
proj_in = dataset_in.GetProjection()
del dataset_in
green_wave = 552.19
nir_wave = 809.2890
green_band_number = find_band_number(green_wave, file_path)
nir_band_number = find_band_number(nir_wave, file_path)
ndwi = calculate_NDWI(green_band_number, nir_band_number, file_path)
# ★ ENVI(.dat) 输出上 GDAL 的原地写（SieveFilter 等）会静默失败，
# 因此需要碎斑过滤时，先写到 GTiff 临时载体做 Sieve（Byte 可靠），
# 再 CreateCopy 转为最终格式。
out_ext = os.path.splitext(output_path)[1].lower()
final_gtiff = out_ext in ('.tif', '.tiff')
need_sieve = sieve_threshold > 0
work_path = output_path if (final_gtiff or not need_sieve) \
else output_path + '__sieve_tmp.tif'
water_binary = extract_water(ndwi, threshold=ndwi_threshold)
generate_water_mask_chunked(
img_path=file_path,
out_mask_path=work_path,
green_band_idx=green_band_number,
nir_band_idx=nir_band_number,
threshold=ndwi_threshold,
data_ignore_value=data_ignore_value,
)
write_bands(file_path, output_path, water_binary)
# ★ 去除小碎斑：SieveFilter 消除面积 < sieve_threshold 像素的孤立斑块
if sieve_threshold > 0:
ds = gdal.Open(output_path, gdal.GA_Update)
# ★ 去除小碎斑：SieveFilter 消除面积 < sieve_threshold 像素的孤立斑块。
# SieveFilter 是 GDAL 文件级流式算子，对大栅格内存安全。
if need_sieve:
ds = gdal.Open(work_path, gdal.GA_Update)
if ds is not None:
srcband = ds.GetRasterBand(1)
gdal.SieveFilter(srcband, None, srcband, sieve_threshold)
ds.FlushCache()
ds = None
if work_path != output_path:
_remove_raster_output(output_path) # 清掉旧 .dat/.hdr，保证可重跑
final_ds = gdal.GetDriverByName('ENVI').CreateCopy(
output_path, gdal.Open(work_path, gdal.GA_ReadOnly))
final_ds = None
if os.path.exists(work_path):
os.remove(work_path)
return output_path