【Python】planetscope UDM文件掩膜代码

用planetscope自带的影像质量文件去除不可用像素的代码https://developers.planet.com/docs/data/udm-2/

自存用

import os
import copy
import numpy as np
import pylab as plt
from osgeo import gdal


rt_file_path = "影像路径"
gl_file_path = "质量文件路径"
out_file_path = "输出路径"


def tif_write(filename, im_data, im_proj, im_geotrans):
    """
    gdal数据类型包括
    gdal.GDT_Byte,
    gdal.GDT_UInt16, gdal.GDT_Int16, gdal.GDT_UInt32, gdal.GDT_Int32,
    gdal.GDT_Float32, gdal.GDT_Float64
    :param filename: 存出文件名
    :param im_data: 输入数据
    :param im_proj: 投影信息
    :param im_geotrans: 放射变换信息
    :return: 0
    """
    datatype = gdal.GDT_UInt16 #数据类型是16 bite unsigned
    # 判读数组维数
    if len(im_data.shape) == 3:
        im_bands, im_height, im_width = im_data.shape
    else:
        im_bands, (im_height, im_width) = 1, im_data.shape  # 多维或1.2维
    # 创建文件
    driver = gdal.GetDriverByName("GTiff")  # 数据类型必须有，因为要计算需要多大内存空间
    dataset = driver.Create(filename, im_width, im_height, im_bands, datatype)
    dataset.SetGeoTransform(im_geotrans)  # 写入仿射变换参数
    dataset.SetProjection(im_proj)  # 写入投影
    if im_bands == 1:
        dataset.GetRasterBand(1).WriteArray(im_data)  # 写入数组数据
    else:
        for i in range(im_bands):
            dataset.GetRasterBand(i + 1).WriteArray(im_data[i])
    del dataset

rt_file_list = os.listdir(rt_file_path)
for rt_file in rt_file_list:
    if rt_file.endswith('SR_8b.tif'): #遥感影像的后缀
        file_name_split = rt_file.split(".")
        rt_hv = file_name_split[0][:23]

        gl_file_list = os.listdir(gl_file_path)
        for gl_file in gl_file_list:
            if gl_file.endswith('udm2.tif'): #掩膜文件的后缀
                if rt_hv in gl_file:
                    rt_file_tif_path = rt_file_path + rt_file
                    gl_file_tif_path = gl_file_path + gl_file

                    #设置新文件名称
                    if not os.path.exists(out_file_path):
                        os.makedirs(out_file_path)
                    out_file_tif_path = os.path.join(out_file_path, rt_file)

                    #读取遥感影像
                    #读取质量文件
                    inDs = gdal.Open(gl_file_path + gl_file)
                    gl_lai_band = inDs.GetRasterBand(1)
                    gl_raster_array = gl_lai_band.ReadAsArray()

                    rt_raster = gdal.Open(rt_file_path + rt_file)
                    geotrans = rt_raster.GetGeoTransform()
                    proj = rt_raster.GetProjection()
                    rt_raster_array = rt_raster.ReadAsArray()
                    rt_lai_array_mask = np.where(gl_raster_array==0, np.nan, rt_raster_array)

                    tif_write(out_file_tif_path, rt_lai_array_mask, proj, geotrans)
                    print(rt_file +' 完成')

                    break