迭代阈值图像分割.py

"""
Created on 2020/12/29 16:00.

@Author: yubaby@anne
@Email: [email protected]
"""


import time
import gdal
from gdalconst import *
import numpy as np


def tif_read(tifpath):
    image = gdal.Open(tifpath, GA_ReadOnly)
    im_clos = image.RasterXSize
    im_rows = image.RasterYSize
    band_array = image.ReadAsArray(0, 0, im_clos, im_rows).astype(np.float)
    del image
    return band_array


def tif_write(im_data, path):
    if 'int8' in im_data.dtype.name:
        datatype = gdal.GDT_Byte
    elif 'int16' in im_data.dtype.name:
        datatype = gdal.GDT_UInt16
    else:
        datatype = gdal.GDT_Float32
    if len(im_data.shape) == 3:
        im_bands, im_height, im_width = im_data.shape
    elif len(im_data.shape) == 2:
        im_data = np.array([im_data])
        im_bands, im_height, im_width = im_data.shape
    driver = gdal.GetDriverByName("GTiff")
    dataset = driver.Create(path, int(im_width), int(im_height), int(im_bands), datatype)
    for i in range(im_bands):
        dataset.GetRasterBand(i + 1).WriteArray(im_data[i])
    del dataset


def iteration_threshold(F):
    """
    （1）设置初始阈值t1。
    当目标与背景的面积相当时，可以将初始阈值t1置为整幅图像的平均灰度；
    当目标与背景的面积相差较大时，更好的选择是将初始阈值t1置为最大灰度值与最小灰度值的中间值。
    """
    t1 = np.mean(F)

    '''
    （2）根据t1将图像F分割为F1和F2两部分，
    其中F1包含所有灰度值小于t1的像素，F2包含所有大于t1的像素，分别求出F1和F2的平均灰度值μ1和μ2。
    '''
    # F = F.flatten()
    F1 = F[F <= t1]
    F2 = F[F > t1]
    u1 = np.mean(F1)
    u2 = np.mean(F2)

    '''
    （3）计算新的阈值t2=(μ1+μ2)/2。
    '''
    t2 = (u1 + u2) / 2

    '''
    （4）指定常数t0（很小的正数），
    如果|t2-t1|<=t0，即迭代过程中前后两次阈值很接近时（或者说μ1和μ2不再变化），终止迭代；
    否则令t1=t2，重复步骤（2）、（3）、（4）。
    设定常数t0的目的是为了加快迭代速度，如果不关心迭代速度，则可以设置为0。
    '''
    t0 = 0
    count = 0
    while abs(t2 - t1) != t0:  # PS：t0=0时，|t2-t1|不可能<0，所以，条件变为|t2-t1|=0，即t1=t2时停止迭代
        t1 = t2
        F1 = F[F <= t1]
        F2 = F[F > t1]
        u1 = np.mean(F1)
        u2 = np.mean(F2)
        t2 = (u1 + u2) / 2
        count += 1
        print('迭代次数：', count)

    return t2


def segmentation(F, T):
    F[F >= T] = 1
    F[F < T] = 0
    return F


if __name__ == '__main__':
    start = time.time()

    path_input = 'J:\BaiduNetdiskDownload\GF3\GF3_KAS_FSI_013781_E121.4_N37.6_20190323_L1A_HHHV_L10003900117\zone\\tif\\'
    path_output = path_input + 'output\\'

    list_tif = ['lee_2019_C2_HH',  # 可用数据：基本能目视看出目标地物
                'lee_2019_C2_HV',
                'lee_2019_DB_HH',
                'lee_2019_DB_HV',
                'mean_2019_C2_HH',
                'mean_2019_C2_HV']

    for tif in list_tif:
        img = tif_read(path_input + tif + '.tif')
        thr = iteration_threshold(img)
        result = segmentation(img, thr)
        tif_write(result, path_output + tif + '-' + str(thr) + '.tif')

    print('It takes', time.time() - start, "seconds.")