基于python实现高分二号遥感影像水体提取与水质反演（黑臭水体与水体富营养化）

水体提取函数——NDWI

水体提取函数water.py

import numpy as np
from osgeo import gdal, gdal_array  # 导入读取遥感影像的库
import cv2
from numba import jit
import math


def water_extract_NDWI(fileName):
    in_ds = gdal.Open(fileName)  # 打开样本文件
    # 读取tif影像信息
    xsize = in_ds.RasterXSize  # 获取行列数
    ysize = in_ds.RasterYSize
    bands = in_ds.RasterCount  # 获取波段数
    geotransform = in_ds.GetGeoTransform() # 获取仿射矩阵信息
    projection = in_ds.GetProjectionRef()   # 获取投影信息
    block_data = in_ds.ReadAsArray(0,0,xsize,ysize).astype(np.float32) # 获取影像信息
    # block_data = in_ds.ReadAsArray()
    # band1 = in_ds.GetRasterBand(1)
    # datatype = band1.DataType
    print("波段数为：", bands)
    print("行数为：", xsize)
    print("列数为：", ysize)

    # 获取GF-2号多个波段
    B = block_data[0, :, :]
    G = block_data[1, :, :]
    R = block_data[2, :, :]
    NIR = block_data[3, :, :]

    np.seterr(divide='ignore', invalid='ignore')    #忽略除法中的NaN值


    # 计算归一化差异水体指数NDWI
    NDWI = (G - NIR) * 1.0 / (G + NIR)

    # 生成tif遥感影像
    writetif(1, xsize, ysize, projection, geotransform, NDWI, "NDWI.tif")
    print("NDWI图像生成成功！")

    # 根据阈值划分水体与分水体
    threshold = 0.18
    NDWI = go_fast_NDWI(NDWI, threshold)     #使用numba加快for循环运行速度
    # for row in range(NDWI.shape[0]):
    #     for col in range(NDWI.shape[1]):
    #         if NDWI[row, col] >= threshold:
    #             NDWI[row, col] = 1
    #         else:
    #             NDWI[row, col] = 0


    # 最后对图像进行开运算进行去噪，即先腐蚀后膨胀
    kernel = np.ones((5, 5), np.uint32)
    opening = cv2.morphologyEx(NDWI, cv2.MORPH_OPEN, kernel)

    # 生成掩膜二值图——这种保存方法没有地理信息
    # gdal_array.SaveArray(opening.astype(gdal_array.numpy.uint32),
    #                      'NDWI_mask.tif', format="GTIFF", prototype='')
    # print("NDWI二值化掩膜图像生成成功！")

    # 生成掩膜二值图——采用gdal保，包含地理信息等
    writetif(1, xsize, ysize, projection, geotransform, opening, "NDWI_mask1.tif")
    print("NDWI二值化掩膜图像生成成功！")


    # # 生成掩膜二值图——采用opencv简单划分
    # retval, NDWI_mask = cv2.threshold(NDWI, threshold, 255, cv2.THRESH_BINARY)
    # cv2.imwrite('NDWI_mask.jpg', NDWI_mask)
    # 保存为jpg图片
    cv2.imwrite('NDWI_mask.jpg', NDWI)
    print("NDWI二值化掩膜图像生成成功！")


def writetif(im_bands, xsize, ysize, projection, geotransform, img, outimgname):
    driver = gdal.GetDriverByName('GTiff')
    dataset = driver.Create(outimgname, xsize, ysize, im_bands, gdal.GDT_Float32)  # 1为波段数量
    dataset.SetProjection(projection)  # 写入投影
    dataset.SetGeoTransform(geotransform)  # 写入仿射变换参数
    if im_bands == 1:
        dataset.GetRasterBand(1).WriteArray(img)  # 写入数组数据
    else:
        for i in range(im_bands):
            dataset.GetRasterBand(i + 1).WriteArray(img[i])
    del dataset



@jit(nopython=True)
def go_fast_NDWI(NDWI, threshold):
    for row in range(NDWI.shape[0]):
        for col in range(NDWI.shape[1]):
            if NDWI[row, col] >= threshold:
                NDWI[row, col] = 1
            else:
                NDWI[row, col] = 0
    return NDWI

主函数main.py

from water import water_extract_NDWI
from contours import draw_contours, river_end
from buildshp import raster2shp, raster2vector
from water_quality import water_quality_test


if __name__ == '__main__':

    # 读取原始影像进行水体提取
    water_extract_NDWI('GF2.tif')
    # 读取3波段影像3bd.tif和水体二值图影像NDWI_mask.jpg
    # 提取所有水体轮廓NDWI_water.jpg，进行滤波NDWI_river.jpg，几何约束NDWI_end.jpg，最终获取河流二值图NDWI_river_end.jpg
    draw_contours('3bd.tif', 'NDWI_mask.jpg', 'NDWI_water.jpg', 'NDWI_river.jpg', 'NDWI_end.jpg', 'NDWI_river_end.jpg')    # 绘制轮廓
    # 二值图转shp文件
    raster2shp('NDWI_river_end.jpg', 'NDWI.shp')

    # 进行水质反演
    water_quality_test("river.tif")

基于几何约束提取河流

contours.py

import cv2
import numpy as np
from osgeo import gdal
from numba import jit
import math
from PIL import Image


def cnt_area(cnt):
  area = cv2.contourArea(cnt)
  return area

def cnt_length(cnt):
  length = cv2.arcLength(cnt, True)
  return length


@jit(nopython=True)
def go_fast_river(im, block_data, river_end, river_end_mask):
    for row in range(river_end_mask.shape[0]):
        for col in range(river_end_mask.shape[1]):
            b, g, r = im[row, col]
            if b > 128 and g > 128 and r > 128:
                # river_end[:, row, col] = block_data[:, row, col]
                river_end_mask[row, col] = 1
            else:
                river_end[:, row, col] = 0
                river_end_mask[row, col] = 0
    return river_end, river_end_mask


def draw_contours(threebandsimg, mask_jpg, outimg1, outimg2, outimg3, outimg4):
    # drawContours需三波段图像，其它格式的图像会报错，因此读取3波段图像 (可使用ENVI另外输出）
    image = cv2.imread(threebandsimg)
    img = image.copy()
    # 轮廓提取
    thresh = cv2.imread(mask_jpg, cv2.CV_8UC1)  #图像需设置为8UC1格式
    img2, contours, hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    print("轮廓的数量为:%d" % (len(contours)))
    # 第三个参数传递值为-1，绘制所有轮廓
    cv2.drawContours(img, contours, -1, (0, 0, 255), -1)
    # cv.drawContours(img, contours, 3, (0, 255, 0), 3)
    # cnt = contours[3]
    # cv.drawContours(img, [cnt], 0, (0, 255, 0), 3)
    # 显示轮廓
    cv2.namedWindow('drawContours', 0)
    cv2.imshow('drawContours', img)
    # 保存包含轮廓的图片
    cv2.imwrite(outimg1, img)
    # 绘制面积前n的轮廓
    rivers = image.copy()
    # len(contours[0]), len中存放的是轮廓中的角点
    # 根据轮廓中的点数对轮廓进行排序
    # contours.sort(key=len, reverse=True)
    # 根据轮廓中的面积对轮廓进行排序
    # contours.sort(key=cnt_area, reverse=True)
    # 根据轮廓的长度对轮廓进行排序
    contours.sort(key=cnt_length, reverse=True)

    contours_max = contours[0:3000]  #显示面积前n的轮廓
    cv2.drawContours(rivers, contours_max, -1, (0, 0, 255), -1)     #最后一个参数-1为填充，其它为线宽
    cv2.namedWindow('The first n contours with the largest area', 0)
    cv2.imshow('The first n contours with the largest area', rivers)
    cv2.imwrite(outimg2, rivers)



    contours_end = []
    for c in contours_max:
        # 找到边界坐标
        x, y, w, h = cv2.boundingRect(c)  # 计算点集最外面的矩形边界
        # cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 2)    # 横平竖直的矩形

        # 找面积最小的矩形
        rect = cv2.minAreaRect(c)
        # 得到最小矩形的坐标
        box = cv2.boxPoints(rect)
        # 标准化坐标到整数
        box = np.int0(box)
        # 最小矩形的长宽
        width, height = rect[1]
        # 最小矩形的角度:[-90,0)
        angle = rect[2]
        # 画出边界
        # cv2.drawContours(image, [box], 0, (0, 255, 0), 3)
        # 计算轮廓周长
        length = cv2.arcLength(c, True)

        # 统计最小外接矩形面积与水域面积，便于区分出河流
        area_water = cv2.contourArea(c)
        area_MBR = w * h
        # 如果水体面积与其最小外接矩形之比小于0.45，或者长宽比小于<0.6，且水体长度大于300，则认为该水体为河流
        if (area_water * 1.0 / area_MBR < 0.45 or (width / height < 0.6 or height / width < 0.6)) and length > 300:
            contours_end.append(c)

        # # 计算最小封闭圆的中心和半径
        # (x, y), radius = cv2.minEnclosingCircle(c)
        # # 换成整数integer
        # center = (int(x), int(y))
        # radius = int(radius)
        # # 画圆
        # cv2.circle(image, center, radius, (0, 255, 0), 2)

    print("最终轮廓的数量为:%d" % (len(contours_end)))
    # 显示最终轮廓结果
    end = image.copy()
    cv2.drawContours(end, contours_end, -1, (0, 0, 255), -1)
    cv2.namedWindow('The first n contours with the largest area', 0)
    cv2.imshow('The first n contours with the largest area', end)
    cv2.imwrite(outimg3, end)
    # cv2.waitKey()

    # 创建一个全黑的图像，方便后续输入二值图，生成shp文件
    img_black = image.copy()
    img_black[:, :, 0:3] = 0
    # 生成最终的二值图
    cv2.drawContours(img_black, contours_end, -1, (255, 255, 255), -1)
    cv2.imwrite(outimg4, img_black)
    print("river_end图像生成成功！")


def river_end(river_jpg, ori_tif, outimg1, outimg2):
    in_ds = gdal.Open(ori_tif)  # 打开样本文件
    # 读取tif影像信息
    xsize = in_ds.RasterXSize  # 获取行列数
    ysize = in_ds.RasterYSize
    bands = in_ds.RasterCount  # 获取波段数
    geotransform = in_ds.GetGeoTransform()  # 获取仿射矩阵信息
    projection = in_ds.GetProjectionRef()  # 获取投影信息
    block_data = in_ds.ReadAsArray(0, 0, xsize, ysize).astype(np.float32)  # 获取影像信息

    im = cv2.imread(river_jpg)

    river_end = block_data
    river_end_mask = block_data[0, :, :]  # 随便给定一个初值
    # 获取筛选后的河流tif影像以及其二值tif影像
    river_end, river_end_mask = go_fast_river(im, block_data, river_end, river_end_mask)

    # 生成tif遥感影像
    driver = gdal.GetDriverByName('GTiff')  # 数据类型必须有，因为要计算需要多大内存空间
    im_bands = 4  # 设置输出影像的波段数
    dataset = driver.Create(outimg1, xsize, ysize, im_bands, gdal.GDT_Float32)  # 1为波段数量
    dataset.SetProjection(projection)  # 写入投影
    dataset.SetGeoTransform(geotransform)  # 写入仿射变换参数
    if im_bands == 1:
        dataset.GetRasterBand(1).WriteArray(river_end)  # 写入数组数据
    else:
        for i in range(im_bands):
            dataset.GetRasterBand(i + 1).WriteArray(river_end[i])
    del dataset
    print("river_end图像生成成功！")

    im_bands1 = 1  # 设置输出影像的波段数
    dataset1 = driver.Create(outimg2, xsize, ysize, im_bands1, gdal.GDT_Float32)  # 1为波段数量
    dataset1.SetProjection(projection)  # 写入投影
    dataset1.SetGeoTransform(geotransform)  # 写入仿射变换参数
    if im_bands1 == 1:
        dataset1.GetRasterBand(1).WriteArray(river_end_mask)  # 写入数组数据
    else:
        for i in range(im_bands1):
            dataset1.GetRasterBand(i + 1).WriteArray(river_end_mask[i])
    del dataset1
    print("river_end_mask图像生成成功！")

生成shp，方便后续裁剪水体

buildshp.py 该函数参考

from osgeo import gdal
from osgeo import ogr, osr                  # 导入处理shp文件的库
import os

def raster2shp(src, tgt):
    """
    函数输入的是一个二值影像，利用这个二值影像，创建shp文件
    """
    # src = "extracted_img.tif"
    # 输出的shapefile文件名称
    # tgt = "extract.shp"
    # 图层名称
    tgtLayer = "extract"
    # 打开输入的栅格文件
    srcDS = gdal.Open(src)
    # 获取第一个波段
    band = srcDS.GetRasterBand(1)
    # 让gdal库使用该波段作为遮罩层
    mask = band
    # 创建输出的shapefile文件
    driver = ogr.GetDriverByName("ESRI Shapefile")
    shp = driver.CreateDataSource(tgt)
    # 拷贝空间索引
    srs = osr.SpatialReference()
    srs.ImportFromWkt(srcDS.GetProjectionRef())
    layer = shp.CreateLayer(tgtLayer, srs=srs)
    # 创建dbf文件
    fd = ogr.FieldDefn("DN", ogr.OFTInteger)
    layer.CreateField(fd)
    dst_field = 0
    # 从图片中自动提取特征
    extract = gdal.Polygonize(band, mask, layer, dst_field, [], None)



def raster2vector(raster_path, vecter_path, ignore_values=None):
    # 读取路径中的栅格数据
    raster = gdal.Open(raster_path)
    # in_band 为想要转为矢量的波段,一般需要进行转矢量的栅格都是单波段分类结果
    # 若栅格为多波段,需要提前转换为单波段
    band = raster.GetRasterBand(1)

    # 读取栅格的投影信息,为后面生成的矢量赋予相同的投影信息
    prj = osr.SpatialReference()
    prj.ImportFromWkt(raster.GetProjection())

    drv = ogr.GetDriverByName("ESRI Shapefile")
    # 若文件已经存在,删除
    if os.path.exists(vecter_path):
        drv.DeleteDataSource(vecter_path)

    # 创建目标文件
    polygon = drv.CreateDataSource(vecter_path)
    # 创建面图层
    poly_layer = polygon.CreateLayer(vecter_path[:-4], srs=prj, geom_type=ogr.wkbMultiPolygon)
    # 添加浮点型字段,用来存储栅格的像素值
    field = ogr.FieldDefn("class", ogr.OFTReal)
    poly_layer.CreateField(field)

    # FPolygonize将每个像元转成一个矩形，然后将相似的像元进行合并
    # 设置矢量图层中保存像元值的字段序号为0
    gdal.FPolygonize(band, None, poly_layer, 0)

    # 删除ignore_value链表中的类别要素
    if ignore_values is not None:
        for feature in poly_layer:
            class_value = feature.GetField('class')
            for ignore_value in ignore_values:
                if class_value == ignore_value:
                    # 通过FID删除要素
                    poly_layer.DeleteFeature(feature.GetFID())
                    break

    polygon.SyncToDisk()
    polygon = None
    print("创建矢量shp文件成功")

水质反演

water_quality.py

import numpy as np
import math
from osgeo import gdal, gdal_array  # 导入读取遥感影像的库
from numba import jit
import math

def water_quality_test(fileName):
    in_ds = gdal.Open(fileName)  # 打开样本文件
    # 读取tif影像信息
    xsize = in_ds.RasterXSize  # 获取行列数
    ysize = in_ds.RasterYSize
    bands = in_ds.RasterCount  # 获取波段数
    geotransform = in_ds.GetGeoTransform() # 获取仿射矩阵信息
    projection = in_ds.GetProjectionRef()   # 获取投影信息
    block_data = in_ds.ReadAsArray(0, 0, xsize, ysize).astype(np.float32) # 获取影像信息
    # block_data = in_ds.ReadAsArray()
    # band1 = in_ds.GetRasterBand(1)
    # datatype = band1.DataType
    print("波段数为：", bands)
    print("行数为：", xsize)
    print("列数为：", ysize)

    # 获取GF-2号多个波段
    B = block_data[0, :, :]
    G = block_data[1, :, :]
    R = block_data[2, :, :]
    NIR = block_data[3, :, :]

    np.seterr(divide='ignore', invalid='ignore')

    # 化学水质参数计算
    # from 广州市黑臭水体评价模型构建及污染染溯源研究
    COD = (R / G - 0.5777) / 0.007  # 化学需氧量浓度，单位是mg/L

    # from 广州市黑臭水体评价模型构建及污染染溯源研究
    TP = 4.9703 * np.power((R / G), 11.618)  # 总磷的浓度，单位为mg/L

    # from 基于实测数据的鄱阳湖总氮、 总磷遥感反演模型研究
    TN = 2.2632 * (NIR / G) * (NIR / G) + 1.1392 * (NIR / G) + 0.7451

    # from 广州市黑臭水体评价模型构建及污染染溯源研究
    # 计算结果多为负数，不合理
    # NH3N = (R / G - 0.661) / 0.07       # 氨氮的浓度，单位为mg/L

    # from 基于Landsat-8 OLI影像的术河临沂段氮磷污染物反演
    # 取值范围为：-0.09到1.48 不合理
    # NH3N = 3.1749 * R / G - 1.9909  # 氨氮的浓度，单位为mg/L

    NH3N = 4.519 * ((G - R) / (G + R)) * ((G - R) / (G + R))\
           - 2.1 * (G - R) / (G + R) + 0.47      # 氨氮的浓度，单位为mg/L

    # from 平原水库微污染水溶解氧含量模型反演与验证
    # DO = 15.73229 - 30.80257 * (R + G) / 10000    # 溶解氧的浓度，单位为mg/L
    # from 基于自动监测和Sentinel-2影像的钦州湾溶解氧反演模型研究
    DO = 771.854 * (1 / R) - 1.476 * (R * R) / (B * B) + 6.435

    # 保存影像
    gdal_array.SaveArray(COD.astype(gdal_array.numpy.uint32),
                         './quality/COD.tif', format="GTIFF", prototype='')
    print("COD图像生成成功！")
    writetif(1, xsize, ysize, projection, geotransform, COD, "./quality/COD1.tif")
    print("COD1图像生成成功！")
    gdal_array.SaveArray(TP.astype(gdal_array.numpy.uint32),
                         './quality/TP.tif', format="GTIFF", prototype='')
    print("TP图像生成成功！")
    writetif(1, xsize, ysize, projection, geotransform, TP, "./quality/TP1.tif")
    print("TP1图像生成成功！")
    gdal_array.SaveArray(TN.astype(gdal_array.numpy.uint32),
                        './quality/TN.tif', format="GTIFF", prototype='')
    print("TN图像生成成功！")
    writetif(1, xsize, ysize, projection, geotransform, TN, "./quality/TN1.tif")
    print("TN1图像生成成功！")
    gdal_array.SaveArray(NH3N.astype(gdal_array.numpy.uint32),
                         './quality/NH3N.tif', format="GTIFF", prototype='')
    print("NH3N图像生成成功！")
    writetif(1, xsize, ysize, projection, geotransform, NH3N, "./quality/NH3N1.tif")
    print("NH3N1图像生成成功！")
    gdal_array.SaveArray(DO.astype(gdal_array.numpy.uint32),
                         './quality/DO.tif', format="GTIFF", prototype='')
    print("DO图像生成成功！")
    writetif(1, xsize, ysize, projection, geotransform, DO, "./quality/DO1.tif")
    print("DO1图像生成成功！")

    # 水质等级划分
    watergrades_all = R / R    # 随便给定一个初值
    watergrades_all = go_fast_waterquality_all(COD, TP, NH3N, DO, watergrades_all)
    # 保存影像
    gdal_array.SaveArray(watergrades_all.astype(gdal_array.numpy.uint32),
                         './quality/watergrades_all.tif', format="GTIFF", prototype='')
    print("watergrades_all图像生成成功！")
    writetif(1, xsize, ysize, projection, geotransform, watergrades_all, "./quality/watergrades_all1.tif")
    print("watergrades_all1图像生成成功！")



    # 营养状态指数计算
    # from 基于GF-1影像的洞庭湖区水体水质遥感监测
    chla = 4.089 * (NIR / R) * (NIR / R) - 0.746 * (NIR / R) + 29.7331  # chla 为叶绿素a浓度，单位为mg/m3
    TSS = 119.62 * np.power((R / G), 6.0823)  # tss为总悬浮物浓度，单位为mg / L
    sd = 284.15 * np.power(TSS, -0.67)  # sd 为透明度，单位是cm

    # 保存影像
    gdal_array.SaveArray(chla.astype(gdal_array.numpy.uint32),
                         './quality/chla.tif', format="GTIFF", prototype='')
    print("chla图像生成成功！")
    writetif(1, xsize, ysize, projection, geotransform, chla, "./quality/chla1.tif")
    print("chla1图像生成成功！")
    gdal_array.SaveArray(TSS.astype(gdal_array.numpy.uint32),
                         './quality/TSS.tif', format="GTIFF", prototype='')
    print("TSS图像生成成功！")
    writetif(1, xsize, ysize, projection, geotransform, TSS, "./quality/TSS1.tif")
    print("TSS1图像生成成功！")
    gdal_array.SaveArray(sd.astype(gdal_array.numpy.uint32),
                         './quality/sd.tif', format="GTIFF", prototype='')
    print("sd图像生成成功！")
    writetif(1, xsize, ysize, projection, geotransform, sd, "./quality/sd1.tif")
    print("sd1图像生成成功！")


    TLI_chla = 25 + 10.86 * np.log(chla)
    TLI_sd = 51.18 - 19.4 * np.log(sd)
    TLI_TP = 94.36 + 16.24 * np.log(TP)
    TLI_TN = 54.53 + 16.94 + np.log(TN)
    # TLI_CODMn = 1.09 + 26.61 * np.log(CODMn)

    TLI = 0.3261 * TLI_chla + 0.2301 * TLI_TP + 0.2192 * TLI_TN + 0.2246 * TLI_sd
    TLIgrades = G / G #随便给定一个初值
    go_fast_TLI(TLI, TLIgrades)
    # 保存影像
    gdal_array.SaveArray(TLI.astype(gdal_array.numpy.uint32),
                         './quality/TLI.tif', format="GTIFF", prototype='')
    print("TLI图像生成成功！")
    writetif(1, xsize, ysize, projection, geotransform, TLI, "./quality/TLI1.tif")
    print("TLI1图像生成成功！")

    gdal_array.SaveArray(TLIgrades.astype(gdal_array.numpy.uint32),
                         './quality/TLIgrades.tif', format="GTIFF", prototype='')
    print("TLIgrades图像生成成功！")
    writetif(1, xsize, ysize, projection, geotransform, TLIgrades, "./quality/TLIgrades1.tif")
    print("TLIgrades1图像生成成功！")


    NDBWI = (G - R) / (G + R)   # 黑臭水体指数
    BOI = (G - R) / (B + G + R) 
    # 保存影像
    gdal_array.SaveArray(NDBWI.astype(gdal_array.numpy.uint32),
                         './quality/NDBWI.tif', format="GTIFF", prototype='')
    print("NDBWI像生成成功！")
    writetif(1, xsize, ysize, projection, geotransform, NDBWI, "./quality/NDBWI1.tif")
    print("NDBWI1图像生成成功！")

    writetif(1, xsize, ysize, projection, geotransform, BOI, "./quality/BOI.tif")
    print("BOI图像生成成功！")

    NDBWI_grades = NDBWI
    NDBWI_grades = go_fast_NDBWI(NDBWI, NDBWI_grades)
    BOIgrades = BOI
    BOIgrades = go_fast_BOI(BOI, BOIgrades)
    # 保存影像
    gdal_array.SaveArray(NDBWI_grades.astype(gdal_array.numpy.uint32),
                         './quality/NDBWI_grades.tif', format="GTIFF", prototype='')
    print("NDBWI_grades像生成成功！")
    writetif(1, xsize, ysize, projection, geotransform, NDBWI_grades, "./quality/NDBWI_grades1.tif")
    print("NDBWI_grades1图像生成成功！")
    writetif(1, xsize, ysize, projection, geotransform, BOIgrades, "./quality/BOI_grades.tif")
    print("BOI_grades图像生成成功！")


def writetif(im_bands, xsize, ysize, projection, geotransform, img, outimgname):
    driver = gdal.GetDriverByName('GTiff')
    dataset = driver.Create(outimgname, xsize, ysize, im_bands, gdal.GDT_Float32)  # 1为波段数量
    dataset.SetProjection(projection)  # 写入投影
    dataset.SetGeoTransform(geotransform)  # 写入仿射变换参数
    if im_bands == 1:
        dataset.GetRasterBand(1).WriteArray(img)  # 写入数组数据
    else:
        for i in range(im_bands):
            dataset.GetRasterBand(i + 1).WriteArray(img[i])
    del dataset



# 水质等级划分
@jit(nopython=True)
def go_fast_waterquality_all(COD, TP, NH3N, DO, watergrades_all):
    for row in range(COD.shape[0]):
        for col in range(COD.shape[1]):
            if COD[row, col] > 40 or TP[row, col] > 0.4 or NH3N[row, col] > 2.0 or DO[row, col] < 2:
                watergrades_all[row, col] = 6
            elif COD[row, col] > 30 or TP[row, col] > 0.3 or NH3N[row, col] > 1.5 or DO[row, col] < 3:
                watergrades_all[row, col] = 5
            elif COD[row, col] > 20 or TP[row, col] > 0.2 or NH3N[row, col] > 1.0 or DO[row, col] < 5:
                watergrades_all[row, col] = 4
            elif COD[row, col] > 15 or TP[row, col] > 0.1 or NH3N[row, col] > 0.5 or DO[row, col] < 6:
                watergrades_all[row, col] = 3
            elif COD[row, col] > 15 or TP[row, col] > 0.02 or NH3N[row, col] > 0.15 or DO[row, col] < 7.5:
                watergrades_all[row, col] = 2
            elif COD[row, col] <= 15 and TP[row, col] <= 0.02 and NH3N[row, col] <= 0.15 and DO[row, col] >= 7.5:
                watergrades_all[row, col] = 1
            # else:
            #     watergrades_all[row, col] = 0

    return watergrades_all


# 水质营养状态指数划分
@jit(nopython=True)
def go_fast_TLI(TLI, TLIgrades):
    for row in range(TLI.shape[0]):
        for col in range(TLI.shape[1]):
            if TLI[row, col] < 30:
                TLIgrades[row, col] = 1
            elif (TLI[row, col] >= 30) and (TLI[row, col] < 50):
                TLIgrades[row, col] = 2
            elif (TLI[row, col] >= 50) and (TLI[row, col] < 60):
                TLIgrades[row, col] = 3
            elif (TLI[row, col] >= 60) and (TLI[row, col] < 70):
                TLIgrades[row, col] = 4
            elif TLI[row, col] >= 70:
                TLIgrades[row, col] = 5
            else:
                TLIgrades[row, col] = 0
    return TLI, TLIgrades



# 黑臭水体划分
@jit(nopython=True)
def go_fast_NDBWI(NDBWI, NDBWI_grades):
    for row in range(NDBWI.shape[0]):
        for col in range(NDBWI.shape[1]):
            if NDBWI[row, col] >= 0.115 or NDBWI[row, col] <= -0.05:
                NDBWI_grades[row, col] = 1
                # print(NDBWI_grades[row, col])
            elif NDBWI[row, col] < 0.115 and NDBWI[row, col] > -0.05:
                NDBWI_grades[row, col] = 2
            else:
                NDBWI_grades[row, col] = 0
    return NDBWI_grades

# 黑臭水体划分
@jit(nopython=True)
def go_fast_BOI(BOI, BOIgrades):
    for row in range(BOI.shape[0]):
        for col in range(BOI.shape[1]):
            if BOI[row, col] >= 0.065:
                BOIgrades[row, col] = 1
                # print(BOIgrades[row, col])
            elif BOI[row, col] < 0.065:
                BOIgrades[row, col] = 2
            else:
                BOIgrades[row, col] = 0
    return BOIgrades

最终结果

1、水体提取

2、河流几何筛选前后对比

3、水质反演结果

水库营养等级划分

黑臭水体划分

水质等级划分

代码有一段时间了，也没重新测试，有问题欢迎大家留言或者私戳我。