卫星遥感图像处理——保存tiff文件以及直方图均衡化

对图像处理后出现图像太暗可以采用直方图均衡化：

import numpy as np
from PIL import Image
from skimage import exposure
import tifffile as tf
import cv2
import imageio
from osgeo import gdal
import matplotlib.pyplot as plt

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

# 直方图均衡化
work_dir = 'G:\SPEI_RF-Dataset_test\spei12'      # 自己手动改spei n 对每个时间尺度的图片均衡化
path = f'{work_dir}\\spei_test_result.tif'
before = 'G:\SPEI_RF-Dataset_test\SPEI_1_200101.tif'        # 用于作为模板的样本
after = f'{work_dir}\\spei_light_12.tif'          # 自己手动改spei n 对每个时间尺度的图片均衡化
image = tf.imread(path)
equalized_image = exposure.equalize_hist(image)
# spei = tf.imread(before)       # 用于作为模板的样本读为用于测试的数组
save_tiff(after, equalized_image, before)
equalized_image[np.isnan(equalized_image)] = 0      # 只要有一个nan，计算出来的就都是nan,此处将nan转为0
im = Image.fromarray(equalized_image)
if im.mode == "F":
    im = im.convert('RGB')
np.save(f'{work_dir}\\spei_light_12.npy', arr = equalized_image)      # 存入
cv_show('spei_light', equalized_image)