import sys, traceback
from osgeo import gdal

# 读写影像类
class ImageReaderWriter(object):

    #读图像文件
    @classmethod
    def read_img(cls, filename, xoff=0, yoff=0, im_width=None, im_height=None):
        try:
            dataset = gdal.Open(filename)  # 打开文件
            if im_width == None:
                im_width = dataset.RasterXSize  # 栅格矩阵的列数
            if im_height == None:
                im_height = dataset.RasterYSize  # 栅格矩阵的行数
            num_bands = dataset.RasterCount  # 栅格矩阵的波段数
            im_geotrans = dataset.GetGeoTransform()  # 仿射矩阵
            im_proj = dataset.GetProjection()  # 地图投影信息
            im_data = dataset.ReadAsArray(xoff, yoff, im_width, im_height)  # 将数据写成数组，对应栅格矩阵
            del dataset
            return im_proj, im_geotrans, im_data
        except AttributeError:
            print('gdal打开影像：没有文件！')
        except Exception:
            traceback.print_exc()
            pass

    #写文件，以写成tif为例
    @classmethod
    def write_img(cls, dst_filename, data):
        format = "ENVI"
        driver = gdal.GetDriverByName(format)
        RasterXSize = data.shape[2]  # 遥感影像的sample（列数）
        RasterYSize = data.shape[1]  # 遥感影像的line（行数）
        band = data.shape[0]
        # driver.Create()函数中RasterXSize代表影像的sample（列数），RasterYSize代表影像的line（行数）
        dst_ds = driver.Create(dst_filename, RasterXSize, RasterYSize, band, gdal.GDT_Float32)
        for i in range(band):
            dst_ds.GetRasterBand(i + 1).WriteArray(data[i, :, :])  # gdal的band从1开始，所以dst_ds.GetRasterBand(i+1)
        dst_ds = None