python中写出遥感影像、保存数组到geotiff文件（geotifwrite）

python读入遥感影像后进行数组和矩阵的运算，在计算完成后还需要导出为geotiff格式。
总共需要五个参数，其中：

“坐标系”、"仿射参数"与输入影像的相同；
“数据矩阵”、“数据类型”、“输出路径”自行设置。

坐标系projection、仿射参数geo_tansform在读取影像的时候获取：python利用gdal读遥感卫星影像

def tifwrite(savepath,data,geo_transform,projection,datatype):
    driver = gdal.GetDriverByName("GTiff")
    if len(data.shape) == 3:
         rows,cols,bands=data.shape   
    elif len(data.shape) == 2:
        rows,cols=data.shape
        bands = 1
    if datatype == "FLOAT32":
        dataset=driver.Create(savepath,cols,rows,bands,gdal.GDT_Float32)
    elif datatype == "UINT8":
        dataset = driver.Create(savepath,cols,rows,bands,gdal.GDT_Byte) 
    elif datatype == "UINT16":
        dataset = driver.Create(savepath,cols,rows,bands,gdal.GDT_UInt16)
    else:
        print("A datatype dose not support yet!")
    dataset.SetGeoTransform(geo_transform)
    dataset.SetProjection(projection)
    if bands == 1:
        dataset.GetRasterBand(1).WriteArray(data)
    else:
        for i in range(bands):
            dataset.GetRasterBand(i+1).WriteArray(data[:,:,i])
    dataset = None #关闭文件
    
if __name__ == "__main__":
	tifpath=r"D:/data.tif"
	savepath=r"D:/result.tif"
	tifffile=tifread(tifpath)  ##见另一篇文章：python利用gdal读取遥感卫星影像
	tifwrite(savepath,tifffile.dataarray,tiffile.geo_transform,tiffile.projection,datatype="FLOAT32")