卫星遥感图像处理——SPEI线性回归

数据来源：科学数据库 (scidb.cn)

import cv2       # opencv读取格式是BGR
import os
import numpy as np
import tifffile as tf         # tifffile是tiff文件的读取库
import imageio
from osgeo import gdal
from skimage import exposure
import sys
import matplotlib.pyplot as plt
import pylab
from PIL import Image

def cv_show(name, img):    # 定义展示图片的函数
    cv2.namedWindow(name, 0)
    cv2.imshow(name, img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

"""
保存tiff文件
:param out_path: 保存tif文件路径
:param in_data: 原始数据路径 或 数据array(h,w), array(c,h,w) 或 array(h,w,c) 要求c<h,w
:param tiff_sample: tif样本路径, 用于提供投影和变换
:param dtype: gdal数据类型, 默认自动根据输入识别
"""
def save_tiff(out_path, in_data, tiff_sample = None, dtype = None):
    # 数据
    if isinstance(in_data, str):
        in_data = imageio.imread(in_data)
    # 调整维度
    if len(in_data.shape) == 2:
        in_data = in_data.reshape(1, *in_data.shape)
        dim = 1
    elif len(in_data.shape) == 3:
        if in_data.shape[0] > in_data.shape[-1]:
            in_data = in_data.transpose(2, 0, 1)
        dim = in_data.shape[0]
    else:
        raise NotImplementedError
    # 检测类型
    if dtype is None:
        if in_data.dtype == np.dtype('uint8'):
            dtype = gdal.GDT_Byte
        else:
            dtype = gdal.GDT_Float32
    # 参照文件
    if isinstance(tiff_sample, str):
        tiff_sample = gdal.Open(tiff_sample)
    # 创建TIFF文件
    tiff_driver = gdal.GetDriverByName('GTiff').Create(out_path, *[i for i in in_data.shape[::-1]], dtype)
    if tiff_sample:
        tiff_driver.SetProjection(tiff_sample.GetProjection())  # 投影
        tiff_driver.SetGeoTransform(tiff_sample.GetGeoTransform())  # 变换
    # 写入数据
    for i in range(dim):
        tiff_driver.GetRasterBand(i + 1).WriteArray(in_data[i])
    tiff_driver.FlushCache()
    return

def slop(spei_scal, n):
    img_list = os.listdir(spei_scal)        # img_list：一个时间尺度中所有的图像文件名
    spei_collection = []        # 存放一个时间尺度的spei矩阵集合
    for img in img_list:      # img:一个时间尺度中的一个图像文件名
        # print(img)
        spei_tf = tf.imread(f'{spei_scal}//{img}')      # img:一个时间尺度中的一个图像文件的spei矩阵
        spei_tf[np.isnan(spei_tf)] = 0      # 只要有一个nan，计算出来的就都是nan,此处将nan转为0
        # for i in spei_tf:     # 将SPEI矩阵打印查看值
        #     print(i)
        # print("************** SPEI ***************")
        # print(spei_tf)
        spei_collection.append(spei_tf)         # 将一个时间尺度所有的SPEI矩阵加入
        print(len(spei_collection))     # 加入一张SPEI的图像矩阵打印一次集合长度，最后为一个时间尺度内所有月数
    fenzi_a = np.zeros((4063, 6947))
    fenzi_b1 = 0
    fenzi_b2 = np.zeros((4063, 6947))
    fenmu_1 = 0
    for i in range(n):      # 遍历一个SPEI时间尺度内的240个月对应的图像矩阵
        spei = spei_collection[i]       # 一个时间尺度中的一个SPEI矩阵
        # 迭代输出行
        for x in range(len(spei)):
            # 迭代输出列
            for y in range(len(spei[0])):
                fenzi_a[x][y] += spei[x][y] * (i + 1)
                fenzi_b2[x][y] += spei[x][y]
        fenzi_b1 += (i + 1)
        fenmu_1 += (i + 1) * (i + 1)
        print("img" + str(i + 1))
    fenzi_a = n * fenzi_a
    print(fenzi_a)
    fenzi_b = fenzi_b1 * fenzi_b2
    print(fenzi_b)
    fenmu_1 = n * fenmu_1
    fenmu_2 = fenzi_b1 * fenzi_b1
    fenmu = fenmu_1 - fenmu_2
    print(fenmu)
    result = (fenzi_a - fenzi_b) / fenmu
    print(result)
    np.save('H:\SPEI_RF-Dataset_test\spei_test_result.npy', arr=result)  # 存入
    out_file = 'H:\SPEI_RF-Dataset_test\spei_test_result.tif'
    in_file = 'H:\SPEI_RF-Dataset_test\SPEI24\SPEI_24_200101.tif'       # 模板数据
    save_tiff(out_file, result, in_file)
    cv_show('result', result)

work_dir = 'H:\SPEI_RF-Dataset_test'
n = 240     # 总共240个月：2001.01-2010.12
# n = 60
img_scale_dir = os.listdir(work_dir)      # img_scale_dir：指定所有的SPEI时间尺度文件夹名称['SPEI1', 'SPEI12', 'SPEI24', 'SPEI3', 'SPEI6', 'SPEI9']
for img_scale in img_scale_dir:      # img_name：一个时间尺度的文件夹名称，遍历
    slop(f'{work_dir}\\{img_scale}', n)
    print('***************************** SCALE ************************************')

# 直方图均衡化
path = 'H:\SPEI_RF-Dataset_test\spei_test_result.tif'
image = tf.imread(path)
equalized_image = exposure.equalize_hist(image)
np.save('H:\SPEI_RF-Dataset_test\spei_light.npy', arr=equalized_image)  # 存入
out_file = 'H:\SPEI_RF-Dataset_test\spei_light.tif'
in_file = path       # 模板数据
save_tiff(out_file, equalized_image, in_file)
cv_show('spei', equalized_image)

 结果：