影像后处理：移除小物体和填补空洞

from skimage import morphology
import skimage
import cv2
import numpy as np
import gdal
from scipy import ndimage as ndi
import matplotlib.pyplot as plt


def read_img(filename):
    dataset=gdal.Open(filename)

    im_width = dataset.RasterXSize
    im_height = dataset.RasterYSize

    im_geotrans = dataset.GetGeoTransform()
    im_proj = dataset.GetProjection()
    im_data = dataset.ReadAsArray(0,0,im_width,im_height)

    del dataset
    return im_proj,im_geotrans,im_width, im_height,im_data


def write_img(filename, im_proj, im_geotrans, im_data):
    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
    else:
        im_bands, (im_height, im_width) = 1,im_data.shape

    driver = gdal.GetDriverByName("GTiff")
    dataset = driver.Create(filename, im_width, im_height, im_bands, datatype)

    dataset.SetGeoTransform(im_geotrans)
    dataset.SetProjection(im_proj)

    if im_bands == 1:
        dataset.GetRasterBand(1).WriteArray(im_data)
    else:
        for i in range(im_bands):
            dataset.GetRasterBand(i+1).WriteArray(im_data[i])

    del dataset


def imgProcess(imgPath, outPath, component_size=1200):
    im_proj, im_geotrans, im_width, im_height, im_data = read_img(imgPath)
    img = cv2.imread(imgPath, cv2.IMREAD_LOAD_GDAL)
    img = np.where(img > 1, 1, 0)
    labels = ndi.label(img, output=np.uint32)[0]
    # remove small objects
    unique, counts = np.unique(labels, return_counts=True)
    for (k, v) in dict(zip(unique, counts)).items():
        if v < component_size:
            print("remove small objects successfully")
            img[labels == k] = 0

    # remove small holes
    binary = np.array(backgournd, bool)
    binary = morphology.remove_small_holes(binary, min_size=1000, connectivity=8)
    binary = remove_small_objects(binary , 100, connectivity=8)
    write_img(outPath, im_proj, im_geotrans, binary)

if __name__=="__main__":
    imgPath = 'finnal.tif'
    outPath = imgPath.replace('finnal', 'finnal_v1')
    imgProcess(imgPath, outPath)