from osgeo import gdal
import os

import numpy as np

def write_image(filename, img_proj, img_geotrans, img_data):
    # 判断栅格数据类型
    if ‘int8‘ in img_data.dtype.name:
        datatype = gdal.GDT_Byte
    elif ‘int16‘ in img_data.dtype.name:
        datatype = gdal.GDT_UInt16
    else:
        datatype = gdal.GDT_Float32

    # 判断数组维度
    if len(img_data.shape) == 3:
        img_bands, img_height, img_width = img_data.shape
    else:
        img_bands, (img_height, img_width) = 1, img_data.shape

    # 创建文件
    driver = gdal.GetDriverByName(‘GTiff‘)
    image = driver.Create(filename, img_width, img_height, img_bands, datatype)

    image.SetGeoTransform(img_geotrans)
    image.SetProjection(img_proj)

    if img_bands == 1:
        image.GetRasterBand(1).WriteArray(img_data)
    else:
        for i in range(img_bands):
            image.GetRasterBand(i + 1).WriteArray(img_data[i])

    del image  # 删除变量,保留数据

def divide_image(path,m,n,out):
    in_ds1 = gdal.Open(path)  # 读取原始图像文件信息
    xsize = in_ds1.RasterXSize
    ysize = in_ds1.RasterYSize
    bands = in_ds1.RasterCount
    geotransform = in_ds1.GetGeoTransform()
    projection = in_ds1.GetProjectionRef()
    data = in_ds1.ReadAsArray(0, 0, xsize, ysize).astype(np.float32)
    data1=data*0.0


    patch_ysize=int(ysize/n)
    patch_xsize = int(xsize / m)
    x_mod= xsize % m
    y_mod=ysize % n

    num=0
    for i in range(1,n+1):
        for j in range(1,m+1):
            num += 1
            outfile=out+"\\"+str(i)+‘_‘+str(j)+".tif"
            if i==n and j==m:
                div_image = data[:,(i - 1) * patch_ysize: i * patch_ysize+y_mod, (j - 1) * patch_xsize: j * patch_xsize + x_mod]
            elif i == n:
                div_image = data[:,(i - 1) * patch_ysize: i * patch_ysize + y_mod, (j - 1) * patch_xsize: j * patch_xsize]
            elif j == m:
                div_image = data[:,(i - 1) * patch_ysize: i * patch_ysize, (j - 1) * patch_xsize: j * patch_xsize + x_mod]
            else:
                div_image = data[:,(i - 1) * patch_ysize: i * patch_ysize, (j - 1) * patch_xsize: j * patch_xsize]
            write_image(outfile, projection, geotransform , div_image)
    return xsize, ysize, data1, projection, geotransform,x_mod,y_mod


def merge_image(m,n,xsize,ysize,data1,projection, geotransform ,out,x_mod,y_mod):
    patch_ysize = int(ysize / n)
    patch_xsize = int(xsize / m)
    for i in range(1,n+1):
        for j in range(1,m+1):
            cut_image = out + "\\" + str(i) + ‘_‘ + str(j) + ".tif"
            in_ds1 = gdal.Open(cut_image)  # 读取原始图像文件信息
            xsize = in_ds1.RasterXSize
            ysize = in_ds1.RasterYSize
            data = in_ds1.ReadAsArray(0, 0, xsize, ysize).astype(np.float32)
            if i == n and j == m:
                data1[:,(i - 1) * patch_ysize: i * patch_ysize+y_mod, (j - 1) * patch_xsize: j * patch_xsize + x_mod]=data
            elif i == n:
                data1[:,(i - 1) * patch_ysize: i * patch_ysize + y_mod, (j - 1) * patch_xsize: j * patch_xsize]=data
            elif j == m:
                data1[:,(i - 1) * patch_ysize: i * patch_ysize, (j - 1) * patch_xsize: j * patch_xsize + x_mod]=data
            else:
                data1[:,(i - 1) * patch_ysize: i * patch_ysize, (j - 1) * patch_xsize: j * patch_xsize]=data

    outfile1=out+‘\\‘+‘merge.tif‘
    write_image(outfile1, projection, geotransform ,data1)

if __name__ == ‘__main__‘:
    path = r"D:\test\\test.tif"
    out=r"D:\data\实验结果\divide_image"
    #m代表行的分块数,n代表列的分块数
    m = 8
    n = 4
    xsize, ysize, data1, projection, geotransform,x_mod,y_mod=divide_image(path,m,n,out)
merge_image(m, n, xsize, ysize, data1, projection, geotransform, out,x_mod,y_mod)