python 基于gdal，richdem，pysheds实现 实现洼填、D8流向，汇流累计量计算，河网连接，分水岭及其水文分析与斜坡单元生成

python gdal实现水文分析算法及其斜坡单元生成

实现洼填、D8流向，汇流累计量计算，河网连接，分水岭

# utf-8
import richdem as rdre
from River import *
from pysheds.grid import Grid
import time
from time import time,sleep
import numpy as np
from osgeo import gdal, osr,ogr

# DEM反转
def DEM_inverse(inputDEM):
    raster = gdal.Open(inputDEM)

    # 行数
    rows = raster.RasterYSize
    # 列数
    cols = raster.RasterXSize
    Nodatavalue = getNoDataValue(inputDEM)
    # print(Nodatavalue)

    # 转arry
    ds_array = raster.ReadAsArray()
    maxvalue = ds_array.max()

    # 结果arry
    result_array = np.zeros(ds_array.shape)

    # 转二维list
    source = ds_array.tolist()
    result = result_array.tolist()
    for row in range(rows):
        for col in range(cols):
            if source[row][col] != Nodatavalue:
                result[row][col] = maxvalue - source[row][col]
            else:
                result[row][col] = source[row][col]
    return np.array(result)

# 数组转tif
def ArrayToTif(input, output, resultArray):
    # 数组转tif
    '''1.tif
    input:输入栅格（用于拷贝坐标信息）
    output: 根据numpy格式转出tif格式
    resultArray:数组转tif
    '''
    raster = gdal.Open(input)
    output_ds = gdal.GetDriverByName('GTiff').CreateCopy(output, raster)
    outband = output_ds.GetRasterBand(1)
    try:
        outband.WriteArray(resultArray)
        outband.SetNoDataValue(getNoDataValue(input))
        return 1
    except:
        print("[INFO] 写入失败")
        return -1

def isNoData(v, NoDataValue, eps=1e-6):
    """Check if a value is NoData"""
    return abs(v - NoDataValue) < eps

# 河网水系提取（正DEM提取山谷线，反DEM提取山脊线）
def stream_compute(acc_path, dir_path, result_path, threshold):
    """Execute stream computation"""
    sourceacc_ds = gdal.Open(acc_path)
    sourcedir_ds = gdal.Open(dir_path)
    sourceacc = sourceacc_ds.ReadAsArray()
    sourcedir = sourcedir_ds.ReadAsArray()

    if sourceacc is None or sourcedir is None:
        return False

    NoDataValue = sourceacc_ds.GetRasterBand(1).GetNoDataValue()
    result = np.full_like(sourceacc, NoDataValue)
    rows, cols = sourceacc.shape

    for row in range(1, rows - 1):
        for col in range(1, cols - 1):
            accvalue = sourceacc[row, col]
            if isNoData(accvalue, NoDataValue) or accvalue < threshold:
                result[row, col] = NoDataValue
            else:
                dir_ = -1.0
                flow_in_count = 0

                dir_ = sourcedir[row + 1, col - 1]
                accvalue = sourceacc[row + 1, col - 1]
                if accvalue >= threshold and dir_ == 128:
                    flow_in_count += 1

                dir_ = sourcedir[row + 1, col]
                accvalue = sourceacc[row + 1, col]
                if accvalue >= threshold and dir_ == 64:
                    flow_in_count += 1

                dir_ = sourcedir[row + 1, col + 1]
                accvalue = sourceacc[row + 1, col + 1]
                if accvalue >= threshold and dir_ == 32:
                    flow_in_count += 1

                dir_ = sourcedir[row, col - 1]
                accvalue = sourceacc[row, col - 1]
                if accvalue >= threshold and dir_ == 1:
                    flow_in_count += 1

                dir_ = sourcedir[row, col + 1]
                accvalue = sourceacc[row, col + 1]
                if accvalue >= threshold and dir_ == 16:
                    flow_in_count += 1

                dir_ = sourcedir[row - 1, col - 1]
                accvalue = sourceacc[row - 1, col - 1]
                if accvalue >= threshold and dir_ == 2:
                    flow_in_count += 1

                dir_ = sourcedir[row - 1, col]
                accvalue = sourceacc[row - 1, col]
                if accvalue >= threshold and dir_ == 4:
                    flow_in_count += 1

                dir_ = sourcedir[row - 1, col + 1]
                accvalue = sourceacc[row - 1, col + 1]
                if accvalue >= threshold and dir_ == 8:
                    flow_in_count += 1

                if flow_in_count != 1:
                    result[row, col] = -1.0
                else:
                    result[row, col] = -2.0

    strcode = 1
    for row in range(rows):
        for col in range(cols):
            if result[row, col] != -1.0 or isNoData(result[row, col], NoDataValue):
                continue

            result[row, col] = strcode * 1.0
            _row, _col = row, col
            while True:
                dir_ = sourcedir[_row, _col]
                if dir_ == 1:
                    _col += 1
                if dir_ == 2:
                    _row += 1
                    _col += 1
                if dir_ == 4:
                    _row += 1
                if dir_ == 8:
                    _row += 1
                    _col -= 1
                if dir_ == 16:
                    _col -= 1
                if dir_ == 32:
                    _row -= 1
                    _col -= 1
                if dir_ == 64:
                    _row -= 1
                if dir_ == 128:
                    _row -= 1
                    _col += 1

                val = result[_row, _col]
                if isNoData(val, NoDataValue) or val == -1.0:
                    break
                if val == -2.0:
                    result[_row, _col] = strcode * 1.0
                else:
                    break

            strcode += 1

    driver = gdal.GetDriverByName('GTiff')
    outRaster = driver.Create(result_path, cols, rows, 1, gdal.GDT_Float32)
    outRaster.SetGeoTransform(sourceacc_ds.GetGeoTransform())

    outband = outRaster.GetRasterBand(1)
    outband.WriteArray(result)
    outband.SetNoDataValue(NoDataValue)
    outRaster.SetProjection(sourceacc_ds.GetProjection())
    outband.FlushCache()

    return True
#栅格转多边形shp
def RasteerToShp(input, output, nodata):
    '''
    栅格转多边形
    :param input: 输入栅格
    :param output: 输出shp文件
    :param nodata: 忽略nodata值
    :return: NULL
    '''

    src_ds = gdal.Open(input)

    srcband = src_ds.GetRasterBand(1)

    # 获取栅格数据集的地理参考信息
    geo_ref = src_ds.GetGeoTransform()
    proj_ref = src_ds.GetProjection()

    # 获取栅格数据集的坐标系

    srs = osr.SpatialReference()
    srs.ImportFromWkt(proj_ref)

    dst_layername = "Value"
    drv = ogr.GetDriverByName("ESRI Shapefile")
    dst_ds = drv.CreateDataSource(output)
    dst_layer = dst_ds.CreateLayer(dst_layername, srs=srs)

    dst_fieldname = 'DN'
    fd = ogr.FieldDefn(dst_fieldname, ogr.OFTInteger)
    dst_layer.CreateField(fd)
    dst_field = 0
    gdal.Polygonize(srcband, None, dst_layer, 0, ["noData=" + str(nodata)], callback=None)

    for feature in dst_layer:
        if feature.GetField("DN") == nodata:
            # print(feature)
            dst_layer.DeleteFeature(feature.GetFID())

        if feature.GetField("DN") == -2147483648:
            # print(feature)
            dst_layer.DeleteFeature(feature.GetFID())


# 分水岭生产，转矢量可以直接生成斜坡单元
def setWatershedCell(data_dir, data_stm, source_rlt, row, col, NoDataValue, eps):
    """Recursively set watershed cell"""
    stm_value = source_rlt[row, col]

    neighbors = [(0, -1, 1.0), (-1, -1, 2.0), (-1, 0, 4.0), (-1, 1, 8.0),
                 (0, 1, 16.0), (1, 1, 32.0), (1, 0, 64.0), (1, -1, 128.0)]

    for nrow, ncol, nvalue in neighbors:
        nrow += row
        ncol += col
        if 0 <= nrow < data_dir.shape[0] and 0 <= ncol < data_dir.shape[1]:
            rlt_value = source_rlt[nrow, ncol]
            if data_dir[nrow, ncol] == nvalue and isNoData(rlt_value, NoDataValue, eps) and isNoData(
                    data_stm[nrow, ncol], NoDataValue, eps):
                source_rlt[nrow, ncol] = stm_value
                setWatershedCell(data_dir, data_stm, source_rlt, nrow, ncol, NoDataValue, eps)


def watershed_compute(source_from1, source_from2, result_to):
    """Execute watershed computation"""
    # Read sources
    l1 = gdal.Open(source_from1)
    l2 = gdal.Open(source_from2)
    sourcedir = l1.ReadAsArray()
    sourcestm = l2.ReadAsArray()

    if sourcedir is None or sourcestm is None:
        return False
    NoDataValue = l2.GetRasterBand(1).GetNoDataValue()
    result = np.full_like(sourcedir, NoDataValue)
    rows, cols = sourcedir.shape

    for row in range(rows):
        for col in range(cols):
            stmvalue = sourcestm[row, col]
            if not isNoData(stmvalue, NoDataValue, eps=1e-6):
                result[row, col] = stmvalue
                setWatershedCell(sourcedir, sourcestm, result, row, col, NoDataValue, eps=1e-6)

    # Save results
    driver = gdal.GetDriverByName('GTiff')
    outRaster = driver.Create(result_to, cols, rows, 1, gdal.GDT_Float32)
    outRaster.SetGeoTransform(l1.GetGeoTransform())

    outband = outRaster.GetRasterBand(1)
    outband.WriteArray(result)
    outband.SetNoDataValue(NoDataValue)
    outRasterSRS = osr.SpatialReference()
    outRaster.SetProjection(l1.GetProjection())
    outband.FlushCache()

    # Return success status
    return True

主函数 斜坡单元及其水文分析所有函数使用

斜坡单元生成原理为正反DEM分水岭shp进行Union，下面例子为正向DEM的水文分析

# utf-8
import richdem as rdre
from River import *
from pysheds.grid import Grid
import time
from time import time,sleep
import numpy as np
from osgeo import gdal, osr,ogr

'''

    :param input: 输入原始的DEM数据(要求nodata值，例如您的DEM为float32，nodata为-np.float32)
    :param filepah: 目标服务器下缓存目录
    :param output: 矢量目录
    :param steam_thold: 河网阈值
   

 '''
inputdem = "./dem.tif"
filepah = "./"
steam_thold =1000
DemFillPath = filepah + 'DEM_FII.tif'
# D8流向计算
FlowPath = filepah + "DEM_FLOW.tif"
# pyshed 流向计算
grid = Grid.from_raster(inputdem )
dem = grid.read_raster(inputdem )
pit_filled_dem = grid.fill_pits(dem)
# Fill depressions in DEM
flooded_dem = grid.fill_depressions(pit_filled_dem)
ArrayToTif(inputdem , DemFillPath, flooded_dem)
# Resolve flats in DEM
inflated_dem = grid.resolve_flats(flooded_dem)
# Determine D8 flow directions from DEM
# ----------------------
nodata22 = getNoDataValue(inputdem )
# Determine D8 flow directions from DEM
# ----------------------
# Specify directional mapping
dirmap = (64, 128, 1, 2, 4, 8, 16, 32)
# Compute flow directions
# -------------------------------------
fdir = grid.flowdir(inflated_dem, dirmap=dirmap, nodata_in=nodata22)
# Calculate flow accumulation
ArrayToTif(inputdem , FlowPath, fdir)
accumPath = filepah + "DEM_ACC.tif"

# 流量计算(基于pyshed)
acc = grid.accumulation(fdir, nodata_in=nodata22)
# tmp2 = np.where(acc == 0, nodata22, acc)
ArrayToTif(FlowPath, accumPath, acc)

#  河网
# print("河网")
SteamPath = filepah + "Stream" + "_" + str(steam_thold) + ".tif"
stream_compute(accumPath, FlowPath, SteamPath, steam_thold)

# 分水岭
# print("分水岭")
Watershed_Path = filepah + "Watershed" + "_" + str(steam_thold) + ".tif"
watershed_compute(FlowPath,SteamPath ,Watershed_Path)
# 分水岭转矢量shp
shpfile3 = filepah + "Watershed" + "_Finshed" + str(steam_thold) + ".shp"
RasteerToShp(Watershed_Path,getNoDataValue(Watershed_Path))

#DEM反转
InvserseDEM = DEM_inverse(inputDEM)
ArrayToTif(inputDEM, filepath+ "Inverse_DEM.tif", InvserseDEM)