143 lines
5.5 KiB
Python
143 lines
5.5 KiB
Python
from util import *
|
||
from osgeo import gdal
|
||
import argparse
|
||
|
||
|
||
def xml2shp():
|
||
pass
|
||
|
||
|
||
def rasterize_envi_xml(shp_filepath):
|
||
pass
|
||
|
||
|
||
@timeit
|
||
def rasterize_shp(shp_filepath, raster_fn_out, img_path, NoData_value=0):
|
||
dataset = gdal.Open(img_path)
|
||
im_width = dataset.RasterXSize
|
||
im_height = dataset.RasterYSize
|
||
geotransform = dataset.GetGeoTransform()
|
||
imgdata_in = dataset.GetRasterBand(1).ReadAsArray()
|
||
del dataset
|
||
|
||
# Open the data source and read in the extent
|
||
source_ds = gdal.OpenEx(shp_filepath)
|
||
# about 25 metres(ish) use 0.001 if you want roughly 100m
|
||
pixel_size = geotransform[1]
|
||
raster_fn_out_tmp = append2filename(raster_fn_out, "_tmp_delete")
|
||
gdal.Rasterize(raster_fn_out_tmp, source_ds, format='envi', outputType=gdal.GDT_Byte,
|
||
noData=NoData_value, initValues=NoData_value, xRes=pixel_size, yRes=-pixel_size, allTouched=True,
|
||
burnValues=1)
|
||
|
||
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()
|