Landsat 7 去条带

# landsat 7 条带修复

import gdal
import os
import numpy as np

# 条带修复


def gdal_repair(tif_name, out_name, bands):
    """
      tif_name(string): 源影像名
      out_name(string): 输出影像名 
      bands(integer): 影像波段数
    """
    # 打开影像文件
    tif = gdal.Open(tif_name)

    # 根据文件类型获取对应的驱动程序
    driver = gdal.GetDriverByName('GTiff')

    # 根据指定文件的驱动程序，使用现有数据集创建新的可写数据集
    # 所有支持创建新文件的驱动程序都支持该`CreateCopy()`方法，   # 但仅`Create()`部分支持该方法
    # CreateCopy的第一个参数为目标文件名，第二个参数为源数据集
    # 第三个参数的值是`0`或`1`，值是`0`。即使无法将原始数据准确地转换为目标数据，程序仍将执行
    new_img = driver.CreateCopy(out_name, tif, 0)

    for i in range(1, bands+1):
        # 分别对每个波段处理
        
        band = new_img.GetRasterBand(i)

        # 使用FillNodata对条带部分进行插值
        gdal.FillNodata(targetBand=band, maskBand=band,
                        maxSearchDist=15, smoothingIterations=0)

        # 将修复好的波段写入新数据集中
        new_img.GetRasterBand(i).WriteArray(band.ReadAsArray())
    print('repair done!')
    del new_img


# 打开landsat 7  MTL文件
# 仅打开1、2、3、4、5、7波段
def openMTL(path_):
    if not (os.path.exists(path_)):
        print('文件不存在')
    name_ = {}
    for i in range(1, 6):
        name_['FILE_NAME_BAND_' + str(i)] = ''
    name_['FILE_NAME_BAND_7'] = ''

    # 提取各波段文件名
    with open(path_) as MTLfile:
        for line in MTLfile:
            for key_ in name_.keys():
                if key_ in line:
                    name_[key_] = line.split("\"")[-2]
    # test
    # for key in name_.keys():
    #     print(key, '\t', name_[key])

    root_ = os.path.dirname(path_)

    data, proj, trans = None, None, None
    for key in name_.keys():
        # 输出文件名
        # if os.path.exists(os.path.join(root_, name_[key])):
        #     print(os.path.join(root_, name_[key]))
        if proj == None:
            band_ = gdal.Open(os.path.join(root_, name_[key]))
            data = np.array([band_.ReadAsArray()])
            proj = band_.GetProjection()
            trans = band_.GetGeoTransform()
        else:
            band_ = gdal.Open(os.path.join(root_, name_[key]))
            data = np.append(data, np.array([band_.ReadAsArray()]), axis=0)

    return data, trans, proj


def WriteTiff(data_, proj, trans, path):
    data_[np.isinf(data_)] = 0
    data_[np.isneginf(data_)] = 0
    data_[np.isnan(data_)] = 0
    if 'int8' in data_.dtype.name:
        datatype = gdal.GDT_Byte
    elif 'int16' in data_.dtype.name:
        datatype = gdal.GDT_UInt16
    else:
        datatype = gdal.GDT_Float32
    if len(data_.shape) == 3:
        band, height, width = data_.shape
    elif len(data_.shape) == 2:
        data_ = np.array([data_])
        band, height, width = data_.shape
    else:
        print('???')

    driver = gdal.GetDriverByName('GTiff')
    dataset = driver.Create(path, width, height, band, datatype)
    if dataset != None:
        dataset.SetGeoTransform(trans)
        dataset.SetProjection(proj)
    for i in range(band):
        dataset.GetRasterBand(i+1).WriteArray(data_[i])
    print('write ' + path + ' successful!')
    del dataset


def main():
    date_list = ['2000', '2010', '2015', '2020']
    for date in date_list:
        print('================================' + date + '================================')
        for root_, dirs, names in os.walk(os.path.join('../1-data_orign', date)):
            for name_ in names:
                if name_[-7:] == 'MTL.txt':
                    print('\nnew process...')
                    print(os.path.abspath(os.path.join(root_, name_)))
                    # 通过MTL打开数据
                    data, trans, proj = openMTL(os.path.join(root_, name_))
                    # 保存为Tiff数据
                    tiff_path = '../temp/temp.tif'
                    if not os.path.exists(os.path.dirname('../temp/temp.tif')):
                        os.makedirs(os.path.dirname('../temp/temp.tif'))
                    WriteTiff(data, proj, trans, tiff_path)
                    # 条带处理
                    out_root = '../2-data_gapfill/' + date
                    if not os.path.exists(out_root):
                        os.makedirs(out_root)
                    out_path = os.path.join(out_root, (name_[:-7] + '.tif'))
                    gdal_repair(tiff_path, out_path, 6)
                    os.remove(tiff_path)


if __name__ == '__main__':
    # tif_path = r'D:\A_zhang\data\1-data_orign\temp\test.tif'
    # repair_path = r'D:\A_zhang\data\1-data_orign\temp\repair.tif'
    # data, trans, proj = openMTL(r'D:\A_zhang\data\1-data_orign\2010\LE07_L1TP_119034_20101210_20200910_02_T1\LE07_L1TP_119034_20101210_20200910_02_T1_MTL.txt')
    # WriteTiff(data, proj, trans, tif_path)
    # gdal_repair(tif_path, repair_path, 6)
    main()