WQ_GUI/src/utils/extract_water_area.py

from src.utils.util import *
from osgeo import gdal, ogr
import argparse
gdal.UseExceptions()
ogr.UseExceptions()

def xml2shp():
    pass


def rasterize_envi_xml(shp_filepath):
    pass


@timeit
def rasterize_shp(shp_filepath, raster_fn_out, img_path, NoData_value=None):
    # ---------- 防御性处理：路径标准化 ----------
    shp_filepath = os.path.abspath(shp_filepath).replace('\\', '/')
    print(f"[DEBUG rasterize_shp] 标准化后的 SHP 路径: {shp_filepath}")

    # 检查伴随文件完整性
    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}")

    # 确保 GDAL/OGR 驱动已注册
    gdal.AllRegister()
    ogr.RegisterAll()

    # 检查 ESRI Shapefile 驱动
    driver = ogr.GetDriverByName("ESRI Shapefile")
    if driver is None:
        raise RuntimeError(
            "系统中未找到 ESRI Shapefile 驱动！请检查 GDAL 是否正确安装及是否包含 Shapefile 支持。"
        )
    print(f"[DEBUG rasterize_shp] ESRI Shapefile 驱动: {driver.GetName()}")

    # 打开参考影像获取尺寸信息
    dataset = gdal.Open(img_path)
    if dataset 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

    # ---------- 打开 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)

    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

    # 创建和输入影像相同行列号、相同分辨率的水域掩膜，方便后续使用
    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]

            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

            if imgdata_in[row, column] == 0:  # 当shp区域比影像区域大时，略过
                continue

            water_mask[row, column] = data_tmp[coor_pixel[1], coor_pixel[0]]

    write_bands(img_path, raster_fn_out, water_mask)

    os.remove(raster_fn_out_tmp)


def calculate_NDWI(green_bandnumber, nir_bandnumber, filename):
    dataset = gdal.Open(filename)  # 打开文件
    num_bands = dataset.RasterCount  # 栅格矩阵的波段数
    im_geotrans = dataset.GetGeoTransform()  # 仿射矩阵
    im_proj = dataset.GetProjection()  # 地图投影信息

    tmp = dataset.GetRasterBand(green_bandnumber + 1)  # 波段计数从1开始
    band_green = tmp.ReadAsArray().astype(np.int16)
    tmp = dataset.GetRasterBand(nir_bandnumber + 1)  # 波段计数从1开始
    band_nir = tmp.ReadAsArray().astype(np.int16)

    ndwi = (band_green - band_nir) / (band_green + band_nir)

    del dataset

    return ndwi


def extract_water(ndwi, threshold=0.3, data_ignore_value=0):
    water_region = np.where(ndwi > threshold, 1, data_ignore_value)

    return water_region


def ndwi(file_path, ndwi_threshold=0.4, output_path=None, data_ignore_value=0):
    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)

    water_binary = extract_water(ndwi, threshold=ndwi_threshold)  # 0.4

    write_bands(file_path, output_path, water_binary)

    return output_path


def main():
    parser = argparse.ArgumentParser(description="此程序用于提取水域区域，输出的水域栅格和输入的影像具有相同的行列数。")

    # parser.add_argument("--global_arg", type=str, help="A global argument for all modes", required=True)

    # 创建子命令解析器
    subparsers = parser.add_subparsers(dest="algorithm", required=True, help="Choose a mode")

    rasterize_shp_ = subparsers.add_parser("rasterize_shp", help="Mode 1 description")
    rasterize_shp_.add_argument('-i1', '--img_path', type=str, required=True, help='输入影像文件的路径')
    rasterize_shp_.add_argument('-i2', '--shp_path', type=str, required=True, help='输入shp文件的路径')
    rasterize_shp_.add_argument('-o', '--water_mask_outpath', required=True, type=str, help='输出水体掩膜文件的路径')
    rasterize_shp_.set_defaults(func=rasterize_shp)

    ndwi_ = subparsers.add_parser("ndwi", help="Mode 2 description")
    ndwi_.add_argument('-i1', '--img_path', type=str, required=True, help='输入影像文件的路径')
    ndwi_.add_argument('-i2', '--ndwi_threshold', type=float, required=True, help='输入ndwi水体阈值，大于此值的为水域')
    ndwi_.add_argument('-o', '--water_mask_outpath', required=True, type=str, help='输出水体掩膜文件的路径')
    ndwi_.set_defaults(func=ndwi)

    # 解析参数
    args = parser.parse_args()
    if args.algorithm == "rasterize_shp":
        args.func(args.shp_path, args.water_mask_outpath, args.img_path)
    elif args.algorithm == "ndwi":
        args.func(args.img_path, args.ndwi_threshold, args.water_mask_outpath)


# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    main()