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()