python实现遥感图像阈值分割

1.阈值分割

import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
from osgeo import gdal

GRAY_SCALE = 256

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 write_tiff(path,image_gray,out):
    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()

    driver = gdal.GetDriverByName('GTiff')
    outfile = out + "\\" + os.path.basename(path).split(".tif")[0] + "_mask.tif"  # 对输出文件命名
    out_dataset = driver.Create(outfile, xsize, ysize, 1, gdal.GDT_Float32)  # 创建一个一波段的数据框架
    out_band1 = out_dataset.GetRasterBand(1)
    out_band1.WriteArray(image_gray)

    out_dataset.SetGeoTransform(geotransform)  # 写入仿射变换
    out_dataset.SetProjection(projection)
if __name__ == '__main__':
    path = r"D:\data\实验数据\3\3.tif"
    out = r"D:\data\实验结果"

    #设置阈值
    thresh=40

    #tif转jpg并灰度化
    img = tif_jpg(path).astype(np.uint8)
    image_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    #高斯滤波
    img_blur = cv2.GaussianBlur(image_gray , (5, 5), 5)
    # 阈值提取
    img_blur[img_blur > thresh] = 255
    img_blur[img_blur<=  thresh] =1
    write_tiff(path, img_blur, out)

2.直方图双峰法阈值分割

import os
from osgeo import gdal
import numpy as np
import cv2
GRAY_SCALE = 256

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)
    G1 = (G / np.max(G) * 255)
    B1 = (B / np.max(B) * 255)
    data2 = cv2.merge([R1,G1,B1]).astype(np.int16)
    return data2

def calcGrayHist(image):
    '''
    统计像素值
    :param image:
    :return:
    '''
    # 灰度图像的高，宽
    rows, cols = image.shape
    # 存储灰度直方图
    grayHist = np.zeros([256], np.uint64)
    for r in range(rows):
        for c in range(cols):
            grayHist[image[r][c]] += 1
    return grayHist

#直方图全局阈值
def threshTwoPeaks(image):

    # 计算灰度直方图
    histogram = calcGrayHist(image)

    # 找到灰度直方图的最大峰值对应的灰度值
    maxLoc = np.where(histogram == np.max(histogram))
    firstPeak = maxLoc[0][0]

    # 寻找灰度直方图的第二个峰值对应的灰度值
    measureDists = np.zeros([256], np.float32)
    for k in range(256):
        measureDists[k] = pow(k - firstPeak, 3) * histogram[k]#GGGGGGGGGGGGGGGGGGGGGGGGGGGGGGG
    maxLoc2 = np.where(measureDists == np.max(measureDists))
    secondPeak = maxLoc2[0][0]

    if firstPeak > secondPeak:
        temp = histogram[int(secondPeak): int(firstPeak)]
        minLoc = np.where(temp == np.min(temp))
        thresh = secondPeak + minLoc[0][0]+ 1
    else:
        temp = histogram[int(firstPeak): int(secondPeak)]
        minLoc = np.where(temp == np.min(temp))
        thresh = firstPeak + minLoc[0][0]
    img = image.copy()
    img[img >= thresh] = 255
    img[img < thresh] = 0
    print("firstPeak:",firstPeak,",secondPeak:",secondPeak,",thresh:",thresh)
    return img

def write_tiff(path,image_gray,out):
    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()
    driver = gdal.GetDriverByName('GTiff')
    outfile = out + "\\" + os.path.basename(path).split(".tif")[0] + "_mask.tif"  # 对输出文件命名
    out_dataset = driver.Create(outfile, xsize, ysize, 1, gdal.GDT_Float32)  # 创建一个一波段的数据框架
    out_band1 = out_dataset.GetRasterBand(1)
    out_band1.WriteArray(image_gray)

    out_dataset.SetGeoTransform(geotransform)  # 写入仿射变换
    out_dataset.SetProjection(projection)
    return outfile

if __name__ == '__main__':
    mask = r"F:\algorithm\算法练习\3_cut.tif"
    out = r"F:\algorithm\算法练习"
    img=tif_jpg(mask).astype(np.uint8)
    gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    seg_data = threshTwoPeaks(gray_img)
    write_tiff(mask, seg_data , out)