制作遥感图像超分辨率数据集（下采样4倍）

# FileName:gen_lr.py
# Author:Li Rumei
# Time :2023/2/20 11:01

import torch
import os
import argparse
import cv2
from osgeo import gdal
import random
import numpy as np


# 保存为TIF格式
def arr2img(self, save_path, arr):
    # 保存为jpg格式
    # plt.imsave(save_path, arr)

    driver = gdal.GetDriverByName("GTiff")
    datasetnew = driver.Create(save_path, int(self.RasterXSize/4), int(self.RasterYSize/4), 1, gdal.GDT_UInt16)
    geotransform = list(self.GetGeoTransform())
    geotransform[1] *= 4
    geotransform[5] *= 4
    datasetnew.SetGeoTransform(geotransform)
    datasetnew.SetProjection(self.GetProjection())
    band = datasetnew.GetRasterBand(1)
    band.WriteArray(arr)
    datasetnew.FlushCache()  # Write to disk.必须有清除缓存





# 添加椒盐噪声
def sp_noise(noise_img, proportion):
    '''
    proportion的值表示加入噪声的量，可根据需要自行调整
    return: img_noise
    '''
    height, width = noise_img.shape[0], noise_img.shape[1]#获取高度宽度像素值
    num = int(height * width * proportion) #一个准备加入多少噪声小点
    for i in range(num):
        w = random.randint(0, width - 1)
        h = random.randint(0, height - 1)
        if random.randint(0, 1) == 0:
            noise_img[h, w] = 0
        else:
            noise_img[h, w] = 65535
            #16 bit对应65535
    return noise_img


# add gaussian noise
def gaussian_noise(img, mean, sigma):
    '''
    此函数用将产生的高斯噪声加到图片上
    传入:
        img   :  原图
        mean  :  均值
        sigma :  标准差
    返回:
        gaussian_out : 噪声处理后的图片
    '''
    # 将图片灰度标准化
    img = img / 65535
    # 产生高斯 noise
    noise = np.random.normal(mean, sigma, img.shape)
    # 将噪声和图片叠加
    gaussian_out = img + noise
    # 将超过 1 的置 1，低于 0 的置 0
    gaussian_out = np.clip(gaussian_out, 0, 1)
    # 将图片灰度范围的恢复为 0-65535
    gaussian_out = np.uint16(gaussian_out*65535)
    return gaussian_out# 这里也会返回噪声，注意返回值




def poisson_noise(lr_img):
    img = lr_img
    vals = len(np.unique(img))
    vals = 2 ** np.ceil(np.log2(vals))
    # 给图片添加泊松噪声
    noisy_img = np.random.poisson(img * vals)/ int(vals)
    # 保存图片
    return noisy_img


# Downsample HR images
def downsample_HRNTL(hr_image_dir,lr_image_dir):
    supported_img_formats = (".TIF", ".tiff", ".tif", ".TIFF")
    for filename in os.listdir(hr_image_dir):
        if not filename.endswith(supported_img_formats):
            continue
        name, ext = os.path.splitext(filename)
        save_path = os.path.join(lr_image_dir + "LR" + name[2:] + ext)
        # Read HR image
        self_hr_img = gdal.Open(os.path.join(hr_image_dir, filename))
        # hr_img_dims = (self_hr_img.RasterXSize, self_hr_img.RasterYSize)#读入图片的（宽度，高度）
        hr_img = self_hr_img.GetRasterBand(1).ReadAsArray()

        # Blur with Gaussian kernel
        lr_img = cv2.GaussianBlur(hr_img, (7, 7),0)
        # Downsample image 2x
        lr_img = cv2.resize(lr_img, (0, 0), fx=0.5, fy=0.5,
                               interpolation=cv2.INTER_NEAREST)
        # add noise
        # lr_img_4x = sp_noise(lr_img_4x,0.025)
        lr_img = gaussian_noise(lr_img,0,0.012)
        # lr_img = poisson_noise(lr_img)


        #once more
        # Blur with Gaussian kernel
        lr_img = cv2.GaussianBlur(lr_img, (7, 7), 0)
        # Downsample image 2x
        lr_img = cv2.resize(lr_img, (0, 0), fx=0.5, fy=0.5,
                            interpolation=cv2.INTER_NEAREST)
        # add noise
        # lr_img_4x = sp_noise(lr_img_4x,0.025)
        lr_img = gaussian_noise(lr_img, 0, 0.012)
        # lr_img = poisson_noise(lr_img)

        #save
        arr2img(self_hr_img, save_path, lr_img)



if __name__ == '__main__':
    hr_img_dir ="E:\yan1\夜光\datasets\\NTLSR DATASET\HR"    # HR dir
    lr_image_dir ="E:\yan1\夜光\datasets\\NTLSR DATASET\LR\\" # lr dir
    downsample_HRNTL(hr_img_dir, lr_image_dir)