acgeospatial / satellite_imagery_python Goto Github PK

def k_means_segmentation(folder):
    images = []
    for filename in os.listdir(folder):
        img_ds = gdal.Open(os.path.join(folder,filename), gdal.GA_ReadOnly)
        ds = gdal.Open(os.path.join(folder, filename))
        band = ds.GetRasterBand(2)
        arr = band.ReadAsArray()
        [cols, rows] = arr.shape

        img = np.zeros((img_ds.RasterYSize, img_ds.RasterXSize, img_ds.RasterCount),
               gdal_array.GDALTypeCodeToNumericTypeCode(img_ds.GetRasterBand(1).DataType))
        for b in range(img.shape[2]):
            img[:, :, b] = img_ds.GetRasterBand(b + 1).ReadAsArray()
        images.append(img)
        
    for i in images:
        new_shape = (i.shape[0] * i.shape[1], i.shape[2])
        X = i[:, :, :8].reshape(new_shape)
        k_means = cluster.KMeans(n_clusters=5)
        k_means.fit(X)

        X_cluster = k_means.labels_
        X_cluster = X_cluster.reshape(i[:, :, 0].shape)
        
        format = "GTiff"
        driver = gdal.GetDriverByName(format)
        
        filename = "/testing_results/k_means.tif"
        outDataRaster = driver.Create(filename, rows, cols, 1, gdal.GDT_Byte)
        outDataRaster.SetGeoTransform(ds.GetGeoTransform())##sets same geotransform as input
        outDataRaster.SetProjection(ds.GetProjection())##sets same projection as input
        
        outDataRaster.GetRasterBand(1).WriteArray(X_cluster)
        
        outDataRaster.FlushCache() ## remove from memory
        del outDataRaster