Python熟练使用GDAL库进行栅格影像IO

GDAL库在Python读取栅格影像时具有广泛的应用，这里记录使用GDAL库读取栅格影像和写入栅格影像的函数，并且代码已经经过个人长达2年的完善更新，在此记录。

1. 导入包

import numpy as np
from osgeo import gdal

2. 读取栅格影像函数 (支持tiff影像,ENVI影像)

def readRaster(inputfile: str):
    dataset = gdal.Open(inputfile)
    if dataset == None:
        print(inputfile+" can't open !!!")      
    width = dataset.RasterXSize 
    height = dataset.RasterYSize 
    bands = dataset.RasterCount 
    data = dataset.ReadAsArray(0, 0, width, height)
    geotrans = dataset.GetGeoTransform()
    proj = dataset.GetProjection()
    metadata_dict = {}
    for metadata_domain in dataset.GetMetadataDomainList():
        metadata_dict[metadata_domain] = dataset.GetMetadata(metadata_domain)
    del dataset
    return data, geotrans, proj, metadata_dict

3. 存储栅格影像函数 (支持tiff影像,ENVI影像)

def writeRaster(data: np.ndarray, geotrans: tuple, proj: str, metadata_dict: dict, path: str, driver_name: str):
    if 'int8' in data.dtype.name:
        datatype = gdal.GDT_Byte
    elif 'uint8' in data.dtype.name:
        datatype = gdal.GDT_Byte
    elif 'int16' in data.dtype.name:
        datatype = gdal.GDT_Int16
    elif 'uint16' in data.dtype.name:
        datatype = gdal.GDT_UInt16
    elif 'int32' in data.dtype.name:
        datatype = gdal.GDT_Int32
    elif 'uint32' in data.dtype.name:
        datatype = gdal.GDT_UInt32
    elif 'float32' in data.dtype.name:
        datatype = gdal.GDT_Float32
    elif 'float64' in data.dtype.name:
        datatype = gdal.GDT_Float64
    elif 'complex64' in data.dtype.name:
        datatype = gdal.GDT_CFloat32
    elif 'complex128' in data.dtype.name:
        datatype = gdal.GDT_CFloat64
    else:
        raise Exception(f"Unsupported data type: {data.dtype}")
    if len(data.shape) == 3:
        bands, height, width = data.shape
    elif len(data.shape) == 2:
        data = np.array([data])
        bands, height, width = data.shape
    else:
        raise Exception("wrong input")
    # 创建文件
    driver = gdal.GetDriverByName(driver_name)
    if driver_name == 'ENVI':
        dataset = driver.Create(path, int(width), int(height), int(bands), datatype, options=["INTERLEAVE="+metadata_dict["ENVI"]["interleave"]])
    else:
        dataset = driver.Create(path, int(width), int(height), int(bands), datatype)
    if(dataset!= None):
        dataset.SetGeoTransform(geotrans)  # 写入仿射变换参数
        dataset.SetProjection(proj)  # 写入投影
        for doamin_key, domain_value in metadata_dict.items():
            for key, value in domain_value.items():  # 写入元数据
                dataset.SetMetadataItem(key, value, doamin_key)
    for i in range(bands):
        dataset.GetRasterBand(i+1).WriteArray(data[i])
    dataset.FlushCache()
    del dataset

4. 函数调用示例

if __name__ == "__main__":
    bin_file = r"D:\Python\IR_HSI_Gas_Simulation\L_off_sources\1.tif"
    data, geotrans, proj, metadata_dict = readRaster(bin_file)
    wavelengths_str = metadata_dict["ENVI"]["wavelength"].strip('{}').split(',')
    wavelength = np.array([float(str_item) for str_item in wavelengths_str])
    data_1 = (data - data.min()) / (data.max() - data.min())
    writeRaster(data_1, geotrans, proj, metadata_dict, "./SO2_25.dat", "ENVI")
    print("All Job Were Done!")

5. 补充ENVI头文件数据类型编号查找表

The type of data representation:

• 1 = Byte: 8-bit unsigned integer
• 2 = Integer: 16-bit signed integer
• 3 = Long: 32-bit signed integer
• 4 = Floating-point: 32-bit single-precision
• 5 = Double-precision: 64-bit double-precision floating-point
• 6 = Complex: Real-imaginary pair of single-precision floating-point
• 9 = Double-precision complex: Real-imaginary pair of double precision floating-point
• 12 = Unsigned integer: 16-bit
• 13 = Unsigned long integer: 32-bit
• 14 = 64-bit long integer (signed)
• 15 = 64-bit unsigned long integer (unsigned)

祝编程顺利！