diff --git a/src/utils/extract_water_area.py b/src/utils/extract_water_area.py index 0331567..be829a8 100644 --- a/src/utils/extract_water_area.py +++ b/src/utils/extract_water_area.py @@ -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