diff --git a/src/utils/extract_water_area.py b/src/utils/extract_water_area.py index 1c6711b..0331567 100644 --- a/src/utils/extract_water_area.py +++ b/src/utils/extract_water_area.py @@ -163,16 +163,17 @@ def extract_water(ndwi, threshold=0.3, data_ignore_value=0): return water_region -def ndwi(file_path, ndwi_threshold=0.4, output_path=None, data_ignore_value=0): +def ndwi(file_path, ndwi_threshold=0.4, output_path=None, + data_ignore_value=0, sieve_threshold=20): 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() # 地图投影信息 + 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 @@ -182,10 +183,19 @@ def ndwi(file_path, ndwi_threshold=0.4, output_path=None, data_ignore_value=0): ndwi = calculate_NDWI(green_band_number, nir_band_number, file_path) - water_binary = extract_water(ndwi, threshold=ndwi_threshold) # 0.4 + water_binary = extract_water(ndwi, threshold=ndwi_threshold) write_bands(file_path, output_path, water_binary) + # ★ 去除小碎斑:SieveFilter 消除面积 < sieve_threshold 像素的孤立斑块 + if sieve_threshold > 0: + ds = gdal.Open(output_path, gdal.GA_Update) + if ds is not None: + srcband = ds.GetRasterBand(1) + gdal.SieveFilter(srcband, None, srcband, sieve_threshold) + ds.FlushCache() + ds = None + return output_path