def clear_buffer(water_buf):
global water_buffer_value, no_data_value, dataset_ol
for cell in water_buf:
data_value = cu.get_raster_int_value(dataset_ol, cell[0], cell[1])
if data_value == water_buffer_value:
cu.set_raster_int_value(dataset_ol, cell[0], cell[1],
no_data_value)
def get_surface_route(surface_route_start):
global acc_ds, dir_ds, route_ds, slope_ds, water_value, river_th, surface_route_value
# 提取各水体的坡面流路
for route_start_point in surface_route_start:
# 初始化流路追踪数组
judge_route = [route_start_point]
while len(judge_route) > 0:
current_point = judge_route.pop()
xoff = current_point[0]
yoff = current_point[1]
# 获取此点其汇流累积量
acc_off = cu.off_transform(xoff, yoff, slope_ds, acc_ds)
acc_value = cu.get_raster_float_value(acc_ds, acc_off[0], acc_off[1])
# 获取此点在结果数据的值
ol_value = cu.get_raster_int_value(slope_ds, xoff, yoff)
# 若流向的点不为水体且不为河道则继续
if ol_value != water_value and acc_value < river_th:
# 记录此点为坡面流路
cu.set_raster_int_value(route_ds, xoff, yoff, surface_route_value)
# 获取此点流向
dir_value = cu.get_raster_int_value(dir_ds, xoff, yoff)
# 获取其流向的点
to_point = cu.get_to_point(xoff, yoff, dir_value)
if len(to_point) > 0:
to_p_a_off = cu.off_transform(to_point[0], to_point[1], slope_ds, acc_ds)
if cu.in_data(to_p_a_off[0], to_p_a_off[1], acc_ds.RasterXSize, acc_ds.RasterYSize):
# 加入判断数组
judge_route.append(to_point)
def outlet_basins(outlet, result_ds, s_id, dir_ds):
# 创建追踪数组
trace_array = [outlet]
# 若需要继续追踪
while len(trace_array) > 0:
# 取出当前像元
cur_cell = trace_array.pop()
cur_x = cur_cell[0]
cur_y = cur_cell[1]
# 对当前像元赋值
cu.set_raster_int_value(result_ds, cur_x, cur_y, s_id)
# 获取周边像元索引
n_8_cells = cu.get_8_dir(cur_x, cur_y)
# 遍历8个像元
for n_cell in n_8_cells:
n_x = n_cell[0]
n_y = n_cell[1]
# 若在数据范围内
if cu.in_data(n_x, n_y, dir_ds.RasterXSize, dir_ds.RasterYSize):
# 获得当前点流向值
dir_value = cu.get_raster_int_value(dir_ds, n_x, n_y)
# 获得下游点索引
to_point = cu.get_to_point_128(n_x, n_y, dir_value)
# 若为当前点的上游点
if to_point == cur_cell:
# 将此像元加入到流域追踪数组
trace_array.append(n_cell)
def surface_merge(upstream_inflow, new_slope_surface_id):
global slope_ds, dir_ds, water_buffer_value, water_value
# 定义合并后的坡面分组
slope_surface_unit = []
# 遍历上游流入点对坡面与水体相交处的入流口分组
for upstream_point in upstream_inflow:
# 判断坡面入流点已参与合并
not_in_unit = 1
for index in range(0, len(slope_surface_unit), 1):
if upstream_point in slope_surface_unit[index]:
not_in_unit = 0
break
# 若坡面入流点未参与合并
if not_in_unit:
# 创建新坡面集合并添加新坡面的首个水体入流点
slope_surface = [upstream_point]
# 获取新坡面集合的入流点集
slope_surface = get_new_slope_surface(upstream_point[0], upstream_point[1], upstream_inflow, slope_surface)
# 将新坡面放入合并后的坡面集合中
slope_surface_unit.append(slope_surface)
# 获取合并后新坡面的id
new_slope_surface_id += 1
# 以各合并坡面的入口点集开始合并坡面
for index in range(0, len(slope_surface_unit), 1):
# print(slope_surface_unit[index])
# 获取合并后新坡面与水体的交界点
s_s_inflows = slope_surface_unit[index]
for inflow_point in s_s_inflows:
# 初始化需要更新的像元集
to_update_points = [inflow_point]
# 循环直到待更新数组为空
while len(to_update_points) > 0:
# 取出第一个点
to_update_point = to_update_points.pop()
# 更新此点id
cu.set_raster_int_value(slope_ds, to_update_point[0], to_update_point[1], new_slope_surface_id)
# 获取此点周边需要更新的像元集
# 获取邻近像元
neighbor_points = cu.get_8_dir(to_update_point[0], to_update_point[1])
# 判断各像元是否需要更新
for point in neighbor_points:
point_x = point[0]
point_y = point[1]
judge_in_data = cu.in_data(point_x, point_y, slope_ds.RasterXSize, slope_ds.RasterYSize)
if judge_in_data:
# 若不为水体则继续
ol_data = cu.get_raster_int_value(slope_ds, point_x, point_y)
if ol_data != water_value:
# 若是当前点的上游点
dir_point = cu.off_transform(point_x, point_y, slope_ds, dir_ds)
dir_value = cu.get_raster_int_value(dir_ds, dir_point[0], dir_point[1])
to_point = cu.get_to_point(point_x, point_y, dir_value)
if to_point == to_update_point:
# 加入待更新数组
to_update_points.append(point)
return new_slope_surface_id
def dir_reclassify(old_tif_path, updated_tif_path, final_points_txt=None):
update_value_2array = [1, 2, 4, 8, 16, 32, 64, 128]
new_value_array = [1, 8, 7, 6, 5, 4, 3, 2]
final_value = -1
old_ds = gdal.Open(old_tif_path)
print("Create Classified Raster")
# 创建坡面提取结果数据
no_data_value = 0
file_format = "GTiff"
full_geotransform = old_ds.GetGeoTransform()
driver = gdal.GetDriverByName(file_format)
result_data_path = updated_tif_path
copy_ds = driver.Create(result_data_path, old_ds.RasterXSize, old_ds.RasterYSize, 1, gdal.GDT_Byte, options=['COMPRESS=DEFLATE'])
copy_ds.SetGeoTransform(full_geotransform)
copy_ds.SetProjection(old_ds.GetProjection())
copy_ds.GetRasterBand(1).SetNoDataValue(no_data_value)
txt_flag = 0
final_f = None
if final_points_txt is not None:
print("Record Final Points.")
txt_flag = 1
if os.path.exists(final_points_txt):
os.remove(final_points_txt)
if txt_flag:
final_f = open(final_points_txt, "a")
print("Direction Raster Classify...")
for j in range(copy_ds.RasterYSize):
for i in range(copy_ds.RasterXSize):
data_value = cu.get_raster_un_int_value(old_ds, i, j)
if data_value in update_value_2array:
index = update_value_2array.index(data_value)
new_value = new_value_array[index]
cu.set_raster_int_value(copy_ds, i, j, new_value)
else:
data_value = cu.get_raster_int_value(old_ds, i, j)
# 如果是内流点则随意赋值流向
if data_value == final_value and txt_flag:
# 拟为内流终点赋值流向
# cu.set_raster_int_value(copy_ds, i, j, 1)
# 记录内流区终点到txt
final_record_item = [i, j]
final_record_str = ','.join(str(k) for k in final_record_item)
final_f.write(final_record_str + '\n')
else:
cu.set_raster_int_value(copy_ds, i, j, no_data_value)
old_ds = None
copy_ds = None
final_f = None
def buffer_search(o_res_xoff, o_res_yoff, water_ol_buf):
global lake_ds, slope_ds, river_th, water_channel
# 存储需要遍历的水体像元
water_points = [[o_res_xoff, o_res_yoff]]
# 循环方式遍历水体像元
while len(water_points) > 0:
current_point = water_points.pop()
res_dir_array = cu.get_8_dir(current_point[0], current_point[1])
for index in range(0, 8, 1):
# 水体数据中的索引
res_xoff = res_dir_array[index][0]
res_yoff = res_dir_array[index][1]
# 结果数据的该点索引
ol_point_off = cu.off_transform(res_xoff, res_yoff, lake_ds,
slope_ds)
re_xoff = ol_point_off[0]
re_yoff = ol_point_off[1]
# 若在数据内
if cu.in_data(re_xoff, re_yoff, slope_ds.RasterXSize,
slope_ds.RasterYSize):
# 判断是否记录在外边界数组
in_ol_array = ol_point_off in water_ol_buf
# 判断是否在水体数据范围内
in_water_data = cu.in_data(res_xoff, res_yoff,
lake_ds.RasterXSize,
lake_ds.RasterYSize)
# 若在水体数据中进一步判断是否为水体内像元
if in_water_data:
data_value = cu.get_raster_int_value(
lake_ds, res_xoff, res_yoff)
not_in_water = data_value != water_value
else:
not_in_water = True
# 若不是水体像元且未记录则记录，即记录新边界点
if not in_ol_array and not_in_water:
# 添加缓冲区数组记录
water_ol_buf.append([re_xoff, re_yoff])
# 标记为水体外边界
cu.set_raster_int_value(slope_ds, re_xoff, re_yoff,
water_buffer_value)
# 判断是否在结果数据中记录水体
ol_data_value = cu.get_raster_int_value(
slope_ds, re_xoff, re_yoff)
recode_water = ol_data_value == water_value
# 若是水体且未记录，则继续搜索
if not recode_water and not not_in_water:
cu.set_raster_int_value(slope_ds, re_xoff, re_yoff,
water_value)
water_points.append([res_xoff, res_yoff])
def add_final_to_river(dir_tif,
final_points_txt,
river_tif,
acc_tif,
river_th=None):
dir_ds = gdal.Open(dir_tif, 1)
rt_ds = gdal.Open(river_tif, 1)
acc_ds = gdal.Open(acc_tif)
# 将内流区域终点更新到河流数据
with open(final_points_txt, 'r') as final_file:
for line in final_file.readlines():
final_info_str = line.strip('\n')
final_info = final_info_str.split(',')
# 将final像元的x,y索引
f_x_coord = float(final_info[0])
f_y_coord = float(final_info[1])
final_off = cu.coord_to_off([f_x_coord, f_y_coord], rt_ds)
final_xoff = final_off[0]
final_yoff = final_off[1]
# 得到在河流中对应的索引
r_off = cu.off_transform(final_xoff, final_yoff, dir_ds, rt_ds)
# 若在数据内
if cu.in_data(r_off[0], r_off[1], rt_ds.RasterXSize,
rt_ds.RasterYSize):
# 寻找内流终点周边最小汇流累积像元（最后形成虚拟河系）
ne_cells = cu.get_8_dir(final_xoff, final_yoff)
min_acc_point = []
min_acc = cu.get_raster_float_value(acc_ds, final_xoff,
final_yoff)
for point in ne_cells:
acc_value = cu.get_raster_float_value(
acc_ds, point[0], point[1])
if acc_value < min_acc:
min_acc = acc_value
min_acc_point = point
# 若最小累积量像元不成为河系
if river_th is None or min_acc < river_th:
# 则赋值内流终点流向该方向
dir_value = 1
if river_th is not None:
dir_value = cu.dir_between_points(
[final_xoff, final_yoff], min_acc_point)
cu.set_raster_int_value(dir_ds, final_xoff, final_yoff,
dir_value)
# 记录Final point到河流
cu.set_raster_int_value(rt_ds, r_off[0], r_off[1], 1)
def get_trace_from_record(refer_tif, seaside_txt, final_txt, trace_tif):
ref_ds = gdal.Open(refer_tif)
print("Create Trace file...")
file_format = "GTiff"
driver = gdal.GetDriverByName(file_format)
full_geotransform = ref_ds.GetGeoTransform()
trace_ds = driver.Create(trace_tif, ref_ds.RasterXSize, ref_ds.RasterYSize,
1, gdal.GDT_Byte)
trace_ds.SetGeoTransform(full_geotransform)
trace_ds.SetProjection(ref_ds.GetProjection())
trace_ds.GetRasterBand(1).SetNoDataValue(0)
# 将内流区域终点更新到河流数据
print("Add Final to Trace file...")
with open(final_txt, 'r') as final_file:
for line in final_file.readlines():
final_info_str = line.strip('\n')
final_info = final_info_str.split(',')
# 将final像元的x,y索引
f_x_coord = float(final_info[0])
f_y_coord = float(final_info[1])
off_xy = cu.coord_to_off([f_x_coord, f_y_coord], trace_ds)
if cu.in_data(off_xy[0], off_xy[1], trace_ds.RasterXSize,
trace_ds.RasterYSize):
cu.set_raster_int_value(trace_ds, off_xy[0], off_xy[1], 1)
# 将内流区域终点更新到河流数据
print("Add Seaside to Trace file...")
with open(seaside_txt, 'r') as final_file:
for line in final_file.readlines():
final_info_str = line.strip('\n')
final_info = final_info_str.split(',')
# 将seaside像元的x,y索引
f_x_coord = float(final_info[0])
f_y_coord = float(final_info[1])
off_xy = cu.coord_to_off([f_x_coord, f_y_coord], trace_ds)
if cu.in_data(off_xy[0], off_xy[1], trace_ds.RasterXSize,
trace_ds.RasterYSize):
cu.set_raster_int_value(trace_ds, off_xy[0], off_xy[1], 1)
print('over.')
ref_ds = None
trace_ds = None
def get_no_inflow_cells(dir_tif, no_inflow_tif):
dir_ds = gdal.Open(dir_tif)
# 创建无流入点数据
print("Create No inflow file...")
file_format = "GTiff"
driver = gdal.GetDriverByName(file_format)
full_geotransform = dir_ds.GetGeoTransform()
ni_ds = driver.Create(no_inflow_tif,
dir_ds.RasterXSize,
dir_ds.RasterYSize,
1,
gdal.GDT_Int32,
options=['COMPRESS=DEFLATE'])
ni_ds.SetGeoTransform(full_geotransform)
ni_ds.SetProjection(dir_ds.GetProjection())
ni_ds.GetRasterBand(1).SetNoDataValue(0)
# 遍历基础数据记录无流入点
print("Get No Inflow Data...")
for i in range(dir_ds.RasterYSize):
for j in range(dir_ds.RasterXSize):
c_point = [j, i]
# 邻近像元索引
n_points = cu.get_8_dir(c_point[0], c_point[1])
# 初始化流入标记
inflow_flag = 0
for n_point in n_points:
# 判断是否在数据内
in_data = cu.in_data(n_point[0], n_point[1],
dir_ds.RasterXSize, dir_ds.RasterYSize)
if in_data:
# 判断是否流向原像元点
n_dir = cu.get_raster_int_value(dir_ds, n_point[0],
n_point[1])
to_point = cu.get_to_point(n_point[0], n_point[1], n_dir)
if to_point == c_point:
inflow_flag = 1
break
# 若无流入则记录
if not inflow_flag:
cu.set_raster_int_value(ni_ds, c_point[0], c_point[1], 1)
def spill_point_dir(inner_ras_indexes, dem_ds):
global temp_ds
# 内边界的高程
inner_dem = []
for ras_index in inner_ras_indexes:
dem_value = cu.get_raster_float_value(dem_ds, ras_index[0],
ras_index[1])
inner_dem.append(dem_value)
# 获得最低位置的索引
spill_pt_index = list(map(inner_dem.index, heapq.nsmallest(1,
inner_dem)))[0]
spill_pt = inner_ras_indexes[spill_pt_index]
# 获得溢出方向
spill_dir = get_spill_dir(inner_ras_indexes, spill_pt_index)
# 得到接收点索引
reception_pt = cu.get_to_point(spill_pt[0], spill_pt[1], spill_dir)
cu.set_raster_int_value(temp_ds, spill_pt[0], spill_pt[1], spill_dir)
return spill_pt, spill_dir, reception_pt
def recode_from_river(watershed_id, river_x, river_y, watershed_tif_path, dir_tif_path, water_tif_path, rivers_index):
global water_value
# 创建数据集
watershed_ds = gdal.OpenEx(watershed_tif_path, 1)
dir_ds = gdal.Open(dir_tif_path)
water_ds = gdal.Open(water_tif_path)
# 初始化需要赋值的像元集合
update_cells = [[river_x, river_y]]
# 更新区域内像元值
# print(">>> update cell:", end='')
while len(update_cells) > 0:
# 取出要更新的像元索引
update_cell = update_cells.pop()
# print(update_cell, end='')
# 更新像元值
cu.set_raster_int_value(watershed_ds, update_cell[0], update_cell[1], watershed_id)
# print('update: ', update_cell, '->', watershed_id)
# 得到邻近像元集合
neighbor_cells = cu.get_8_dir(update_cell[0], update_cell[1])
# 搜索上游像元
for neighbor_cell in neighbor_cells:
n_x = neighbor_cell[0]
n_y = neighbor_cell[1]
# 判断邻近点是否在数据内
if cu.in_data(n_x, n_y, watershed_ds.RasterXSize, watershed_ds.RasterYSize):
# 若不为河段并不为湖泊/水库（即若为子流域）
water_off = cu.off_transform(n_x, n_y, watershed_ds, water_ds)
in_water = cu.is_water_cell(water_ds, water_off[0], water_off[1], water_value)
if neighbor_cell not in rivers_index and not in_water:
dir_value = cu.get_raster_int_value(dir_ds, n_x, n_y)
to_point = cu.get_to_point(n_x, n_y, dir_value)
# 若为上游点
if to_point == update_cell:
# 加入更新数组
update_cells.append(neighbor_cell)
watershed_ds = None
dir_ds = None
def get_data_boundary_cells(data_tif, boundary_tif):
data_ds = gdal.Open(data_tif)
# 创建边界点数据
print("Create Boundary file...")
file_format = "GTiff"
driver = gdal.GetDriverByName(file_format)
full_geotransform = data_ds.GetGeoTransform()
b_ds = driver.Create(boundary_tif, data_ds.RasterXSize, data_ds.RasterYSize, 1, gdal.GDT_Byte, options=['COMPRESS=DEFLATE'])
b_ds.SetGeoTransform(full_geotransform)
b_ds.SetProjection(data_ds.GetProjection())
b_ds.GetRasterBand(1).SetNoDataValue(0)
# 遍历基础数据记录边界点
print("Get Boundary Cells...")
# nodata值
no_data = int(data_ds.GetRasterBand(1).GetNoDataValue())
for i in range(data_ds.RasterYSize):
# 前个像元点
f_point_value = None
for j in range(data_ds.RasterXSize):
# 若为首尾行
if i == 0 or i == data_ds.RasterYSize - 1:
c_point_value = cu.get_raster_int_value(data_ds, j, i)
# 若为数据点则记录
if c_point_value != no_data:
cu.set_raster_int_value(b_ds, j, i, 1)
else:
# 若为第一个点或为行最后一个点
if j == 0 or j == data_ds.RasterXSize - 1:
c_point = [j, i]
f_point_value = cu.get_raster_int_value(data_ds, c_point[0], c_point[1])
# 若为数据点则记录
if f_point_value != no_data:
cu.set_raster_int_value(b_ds, c_point[0], c_point[1], 1)
# 若不为第一个点
else:
c_point = [j, i]
# 获取当前点的值
c_point_value = cu.get_raster_int_value(data_ds, c_point[0], c_point[1])
# 若当前点为数据点
if c_point_value != no_data:
# 获取周边邻接像元
n_points = cu.get_8_dir(c_point[0], c_point[1])
# 若周边存在无数据像元则记录
for n_point in n_points:
n_point_value = cu.get_raster_int_value(data_ds, n_point[0], n_point[1])
if n_point_value == no_data:
cu.set_raster_int_value(b_ds, c_point[0], c_point[1], 1)
break
def get_slope_surface(lake_tif_path, dir_tif_path, acc_tif_path, slope_tif_path, route_tif_path, river_threshold, no_data=None):
print("Extract Start")
global lake_ds, dir_ds, acc_ds, slope_ds, route_ds, river_th, water_ol_bufs, no_data_value
river_th = river_threshold
lake_ds = gdal.Open(lake_tif_path)
dir_ds = gdal.Open(dir_tif_path)
acc_ds = gdal.Open(acc_tif_path)
if no_data is not None:
no_data_value = no_data
else:
no_data_value = int(lake_ds.GetRasterBand(1).GetNoDataValue())
dir_geotransform = dir_ds.GetGeoTransform()
full_geotransform = dir_geotransform
# 创建坡面提取结果数据
file_format = "GTiff"
driver = gdal.GetDriverByName(file_format)
slope_ds = driver.Create(slope_tif_path, dir_ds.RasterXSize, dir_ds.RasterYSize, 1, gdal.GDT_Int32, options=['COMPRESS=DEFLATE'])
slope_ds.SetGeoTransform(full_geotransform)
slope_ds.SetProjection(dir_ds.GetProjection())
slope_ds.GetRasterBand(1).SetNoDataValue(no_data_value)
# 创建坡面流路结果数据
route_ds = driver.Create(route_tif_path, dir_ds.RasterXSize, dir_ds.RasterYSize, 1, gdal.GDT_Int16, options=['COMPRESS=DEFLATE'])
route_ds.SetGeoTransform(full_geotransform)
route_ds.SetProjection(dir_ds.GetProjection())
route_ds.GetRasterBand(1).SetNoDataValue(no_data_value)
print("Get lakes' outline...")
for i in range(lake_ds.RasterYSize):
for j in range(lake_ds.RasterXSize):
# 获取水体数据某点的值
scan_value = cu.get_raster_int_value(lake_ds, j, i)
# 若是水体内的像元
if scan_value == water_value:
# 计算水体数据的x,y坐标
off_point = cu.off_transform(j, i, lake_ds, slope_ds)
re_xoff = off_point[0]
re_yoff = off_point[1]
# 若在数据集内
if cu.in_data(re_xoff, re_yoff, slope_ds.RasterXSize, slope_ds.RasterYSize):
# 获取结果数据的值
re_data_value = cu.get_raster_int_value(slope_ds, re_xoff, re_yoff)
# 若未记录则继续
if re_data_value != water_value:
# 新建数组用于记录此水体外边界
water_ol_buf = []
# 结果数据记录水体
cu.set_raster_int_value(slope_ds, re_xoff, re_yoff, water_value)
# 使用3*3区域生成水体外包围像元集
buffer_search(j, i, water_ol_buf)
# 若有新的水体外边界则记录到集合
if len(water_ol_buf) > 0:
# 将此水体的外边界记录到集合中
water_ol_bufs.append(water_ol_buf)
print("Sort the extraction of slope surface...")
# 对水体的坡面提取排序（由小到大）
water_ol_bufs_ordered = []
surface_route_start = water_order(water_ol_bufs, water_ol_bufs_ordered)
print("Search slope surface inflow points...")
# 搜索水体外边界上坡面像元集合
# 坡面id初始化
# 上游流入点的集合
upstream_inflows = []
for water_ol_buf in water_ol_bufs_ordered:
slope_surface_inflows(water_ol_buf, upstream_inflows)
print("Merge lakes' slope surface...")
# 合并坡面入流口相邻的坡面
s_new_id = 0
for upstream_inflow in upstream_inflows:
s_new_id = surface_merge(upstream_inflow, s_new_id)
print("Record slope surface route...")
# 提取坡面流路
get_surface_route(surface_route_start)
print("Clear marks...")
# 清除边界线标记
for water_ol_buf in water_ol_bufs_ordered:
clear_buffer(water_ol_buf)
lake_ds = None
dir_ds = None
acc_ds = None
slope_ds = None
route_ds = None
print("Extract End")
def clear_buffer(water_buf):
global water_buffer_value, no_data_value, slope_ds
for cell in water_buf:
data_value = cu.get_raster_int_value(slope_ds, cell[0], cell[1])
if data_value == water_buffer_value:
cu.set_raster_int_value(slope_ds, cell[0], cell[1], no_data_value)
def watershed_erase_area(dem_tif_path, dir_tif_path, acc_tif_path,
stream_tif_path, water_s_s_tif_path, result_path):
print("Erasing Mask Area from Dataset")
# 获取数据集
dem_old_ds = gdal.Open(dem_tif_path)
dir_old_ds = gdal.Open(dir_tif_path)
acc_old_ds = gdal.Open(acc_tif_path)
stream_old_ds = gdal.Open(stream_tif_path)
mask_ds = gdal.Open(water_s_s_tif_path)
# 获取无数据标识
dem_no_data = get_nodata_value(dem_old_ds)
dir_no_data = get_nodata_value(dir_old_ds)
acc_no_data = get_nodata_value(acc_old_ds)
stream_no_data = get_nodata_value(stream_old_ds)
mask_no_data = get_nodata_value(mask_ds)
# 新数据路径
dem_new_path = result_path + "/dem_erase.tif"
dir_new_path = result_path + "/dir_erase.tif"
acc_new_path = result_path + "/acc_erase.tif"
stream_new_path = result_path + "/stream_erase.tif"
# 创建结果数据
file_format = "GTiff"
driver = gdal.GetDriverByName(file_format)
dem_new_ds = driver.CreateCopy(dem_new_path, dem_old_ds)
dir_new_ds = driver.CreateCopy(dir_new_path, dir_old_ds)
acc_new_ds = driver.CreateCopy(acc_new_path, acc_old_ds)
stream_new_ds = driver.CreateCopy(stream_new_path, stream_old_ds)
# 剔除掩膜区域处理
for j in range(mask_ds.RasterYSize):
for i in range(mask_ds.RasterXSize):
mask_data = cu.get_raster_int_value(mask_ds, i, j)
if mask_data != mask_no_data:
# 处理DEM
result_point = cu.off_transform(i, j, mask_ds, dem_new_ds)
if cu.in_data(result_point[0], result_point[1],
dem_new_ds.RasterXSize, dem_new_ds.RasterYSize):
cu.set_raster_float_value(dem_new_ds, result_point[0],
result_point[1], dem_no_data)
# 处理流向
result_point = cu.off_transform(i, j, mask_ds, dir_new_ds)
if cu.in_data(result_point[0], result_point[1],
dir_new_ds.RasterXSize, dir_new_ds.RasterYSize):
cu.set_raster_int_value(dir_new_ds, result_point[0],
result_point[1], int(dir_no_data))
# 处理汇流累积量
result_point = cu.off_transform(i, j, mask_ds, acc_new_ds)
if cu.in_data(result_point[0], result_point[1],
acc_new_ds.RasterXSize, acc_new_ds.RasterYSize):
cu.set_raster_float_value(acc_new_ds, result_point[0],
result_point[1], acc_no_data)
# 处理河网
result_point = cu.off_transform(i, j, mask_ds, stream_new_ds)
if cu.in_data(result_point[0], result_point[1],
stream_new_ds.RasterXSize,
stream_new_ds.RasterYSize):
cu.set_raster_int_value(stream_new_ds,
result_point[0], result_point[1],
int(stream_no_data))
dem_old_ds = None
dir_old_ds = None
acc_old_ds = None
stream_old_ds = None
dem_new_ds = None
dir_new_ds = None
acc_new_ds = None
stream_new_ds = None
mask_ds = None
return [dem_new_path, dir_new_path, acc_new_path, stream_new_path]
def get_trace_points(dir_tif, flag_tif, trace_tif, seaside_txt=None, final_txt=None):
coastline_value = 0
final_value = -1
dir_ds = gdal.Open(dir_tif)
f_ds = gdal.Open(flag_tif)
# 创建trace start数据
print("Create Trace file...")
file_format = "GTiff"
driver = gdal.GetDriverByName(file_format)
full_geotransform = dir_ds.GetGeoTransform()
trace_ds = driver.Create(trace_tif, dir_ds.RasterXSize, dir_ds.RasterYSize, 1, gdal.GDT_Byte, options=['COMPRESS=DEFLATE'])
trace_ds.SetGeoTransform(full_geotransform)
trace_ds.SetProjection(dir_ds.GetProjection())
trace_ds.GetRasterBand(1).SetNoDataValue(0)
# 创建记录文件
seaside_flag = 0
seaside_f = None
if seaside_txt is not None:
print("Record Final Points.")
seaside_flag = 1
if os.path.exists(seaside_txt):
os.remove(seaside_txt)
if seaside_flag:
seaside_f = open(seaside_txt, "a")
final_flag = 0
final_f = None
if final_txt is not None:
print("Record Final Points.")
final_flag = 1
if os.path.exists(final_txt):
os.remove(final_txt)
if final_flag:
final_f = open(final_txt, "a")
# 遍历基础数据记录seaside
print("Get Trace Data...")
for i in range(f_ds.RasterYSize):
for j in range(f_ds.RasterXSize):
# 获取流向的值
flag_value = cu.get_raster_int_value(f_ds, j, i)
# 如果是coastline
if flag_value == coastline_value:
# 获取周边像元索引
neibor_cells = cu.get_8_dir(j, i)
# 遍历
for n_cell in neibor_cells:
# 判断是否在数据内
in_data = cu.in_data(n_cell[0], n_cell[1], dir_ds.RasterXSize, dir_ds.RasterYSize)
# 若在数据内
if in_data:
# 获取流向值
dir_value = cu.get_raster_int_value(dir_ds, n_cell[0], n_cell[1])
# 获取流向的像元索引
n_to_point = cu.get_to_point(n_cell[0], n_cell[1], dir_value)
# 若为当前上游像元
if n_to_point == [j, i]:
# 记录为seaside
cu.set_raster_int_value(trace_ds, n_cell[0], n_cell[1], 1)
# 记录到文件
if seaside_flag:
# 获得栅格像元左上角坐标
seaside_record_item = cu.off_to_coord([j, i], f_ds)
seaside_record_str = ','.join(str(k) for k in seaside_record_item)
seaside_f.write(seaside_record_str + '\n')
# 如果是final point
elif flag_value == final_value:
# 记录为final points
cu.set_raster_int_value(trace_ds, j, i, 1)
# 记录到文件
if final_flag:
final_record_item = cu.off_to_coord([j, i], f_ds)
final_record_str = ','.join(str(k) for k in final_record_item)
final_f.write(final_record_str + '\n')
dir_ds = None
f_ds = None
trace_ds = None
seaside_f = None
final_f = None
import gdal
if __name__ == '__main__':
workspace = r'G:\Graduation\Program\Data\41\endorheic_area0\multi_basins'
for i in range(0, 8, 1):
shp = workspace + '/' + str(i) + '/basin' + str(i) + '.shp'
geoj = workspace + '/' + str(i) + '/basin' + str(i) + '.geojson'
cu.shp_to_geojson(shp, geoj)
dem_tif_path = workspace + '/' + str(i) + '/dem' + str(i) + '.tif'
s_pt, s_dir, reception_coord = sbe.basin_spill(geoj, dem_tif_path)
if len(s_pt) > 0:
temp_path = workspace + '/' + str(i) + '/outlet' + str(i) + '.tif'
dem_ds = gdal.Open(dem_tif_path)
file_format = "GTiff"
driver = gdal.GetDriverByName(file_format)
temp_ds = driver.Create(temp_path,
dem_ds.RasterXSize,
dem_ds.RasterYSize,
1,
gdal.GDT_Int16,
options=['COMPRESS=DEFLATE'])
temp_ds.SetGeoTransform(dem_ds.GetGeoTransform())
temp_ds.SetProjection(dem_ds.GetProjection())
temp_ds.GetRasterBand(1).SetNoDataValue(-1)
cu.set_raster_int_value(temp_ds, s_pt[0], s_pt[1], s_dir)
temp_ds = None
dem_ds = None
print('raster index: ', s_pt, ' direction: ', s_dir,
'next point location: ', reception_coord)
