python实现分水岭算法分割遥感图像

1. 定义

　　分水岭算法(watershed algorithm)可以将图像中的边缘转化为“山脉”，将均匀区域转化为“山谷”，在这方面有助于分割目标。
　　分水岭算法：是一种基于拓扑理论的数学形态学的分割方法。把图像看作是测地学上的拓扑地貌，图像中的每一个点像素值的灰度值表示该点的海拔高度，每一个局部极小值及其影响的区域称为“集水盆”，集水盆的边界可以看成分水岭。在每一个局部极小值表面刺穿一个小孔，然后把整个模型慢慢浸入水中，随着浸入的加深，每一个局部极小值的影响域慢慢的向外扩展，在两个集水盆汇合处构建大坝，形成分水岭。

迭代标注过程：

1. 排序过程：对每个像素的灰度级进行从低到高的排序
2. 淹没过程：对每一个局部最小值在h阶高度的影响域采用先进先出结构判断及标注。

2.实现算法:watershed()函数

　　这些标记的值可以使用findContours（）函数和drawContours（）函数由二进制的掩模检索出来

3.程序代码：

import numpy as np
import cv2
from osgeo import gdal, gdal_array
import shapefile
try:
    import Image
    import ImageDraw
except:
    from PIL import Image, ImageDraw

def tif_jpg(rasterfile):
    in_ds = gdal.Open(rasterfile)  # 打开样本文件
    xsize = in_ds.RasterXSize  # 获取行列数
    ysize = in_ds.RasterYSize
    bands = in_ds.RasterCount
    block_data = in_ds.ReadAsArray(0, 0, xsize, ysize).astype(np.float32)
    B = block_data[0, :, :]

    G = block_data[ 1,:, :]
    R = block_data[2,:, :]
    R1 =  (R/np.max(R)*255).astype(np.int16)
    G1 = (G / np.max(G) * 255).astype(np.int16)
    B1 = (B / np.max(B) * 255).astype(np.int16)
    data2 = cv2.merge([R1,G1,B1])
    return data2

def watershed(path,out):
    print("分水岭分割")
    in_ds = gdal.Open(path)  # 打开样本文件
    xsize = in_ds.RasterXSize  # 获取行列数
    ysize = in_ds.RasterYSize
    bands = in_ds.RasterCount
    geotransform = in_ds.GetGeoTransform()
    projection = in_ds.GetProjectionRef()
    #tif转jpg     非255通道转换为255通道
    img=tif_jpg(path).astype(np.uint8)
    # 转换为灰度图片
    gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # canny边缘检测 函数返回一副二值图，其中包含检测出的边缘。
    canny = cv2.Canny(gray_img, 80,120)
    # 寻找图像轮廓 返回修改后的图像 图像的轮廓  以及它们的层次
    # canny, contours, hierarchy = cv2.findContours(canny, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    contours, hierarchy = cv2.findContours(canny, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    # 32位有符号整数类型，
    marks = np.zeros(img.shape[:2], np.int32)
    # findContours检测到的轮廓
    imageContours = np.zeros(img.shape[:2], np.uint8)
    # 轮廓颜色
    compCount = 0
    index = 0
    # 绘制每一个轮廓
    for index in range(len(contours)):
        # 对marks进行标记，对不同区域的轮廓使用不同的亮度绘制，相当于设置注水点，有多少个轮廓，就有多少个轮廓
        # 图像上不同线条的灰度值是不同的，底部略暗，越往上灰度越高
        marks = cv2.drawContours(marks, contours, index, (index, index, index), 1, 8, hierarchy)
        # 绘制轮廓，亮度一样
        imageContours = cv2.drawContours(imageContours, contours, index, (255, 255, 255), 1, 8, hierarchy)
    # 查看 使用线性变换转换输入数组元素成8位无符号整型。
    markerShows = cv2.convertScaleAbs(marks)
    # cv2.imshow('imageContours',imageContours)
    # 使用分水岭算法
    marks = cv2.watershed(img, marks)
    driver = gdal.GetDriverByName('GTiff')
    outfile_lake = out + "\\" + "watershed_cut.tif"
    out_dataset = driver.Create(outfile_lake, xsize, ysize, 1, gdal.GDT_Float32)
    out_band1 = out_dataset.GetRasterBand(1)
    out_band1.WriteArray(marks)
    out_dataset.SetGeoTransform(geotransform)  # 写入仿射变换
    out_dataset.SetProjection(projection)
    return outfile_lake

if __name__ == "__main__":
    path = r"D:\data\实验数据\fenlei2.tif"
    out = r"D:\data\实验结果\分割结果"
    watershed(path, out)