| import torch.utils.data as data | |
| from PIL import Image | |
| from random import randrange | |
| from torchvision.transforms import Compose, ToTensor, Normalize | |
| import re | |
| from PIL import ImageFile | |
| from os import path | |
| import numpy as np | |
| import torch | |
| ImageFile.LOAD_TRUNCATED_IMAGES = True | |
| import os | |
| # --- Training dataset --- # | |
| import torch as th | |
| import cv2 | |
| import math | |
| import random | |
| seed = np.random.RandomState(112311) | |
| class ValData(data.Dataset): | |
| def __init__(self, dataset_path, crop_size=[256,256]): | |
| super().__init__() | |
| # train_list = train_data_dir + train_filename | |
| # with open(train_list) as f: | |
| # contents = f.readlines() | |
| # input_names = [i.strip() for i in contents] | |
| # gt_names = [i.strip().replace('input','gt') for i in input_names] | |
| # self.train_data_dir = '/media/labuser/cb8bb1ad-451a-4aa4-870c-2d3eeafe2525/ECCV_2022/diffusion_ema_rain_imagenet/rain_sub1/' | |
| # self.train_data_dir = '/media/labuser/cb8bb1ad-451a-4aa4-870c-2d3eeafe2525/ICIP_Turbulence_files/Tubfaces89/300M/tubimages/' | |
| # self.train_data_dir = "/media/malsha/47a8802b-e0b7-47a8-8a4d-1649cc3ad408/sar_optical/optical/" | |
| self.noisy_path = os.path.join(dataset_path, 'noisy') | |
| # self.noisy_path = dataset_path | |
| # self.clean_path = dataset_path | |
| self.clean_path = os.path.join(dataset_path, 'clean') | |
| self.images_list = os.listdir(self.noisy_path) | |
| self.crop_size = crop_size | |
| def __len__(self): | |
| return len(os.listdir(self.noisy_path)) | |
| def __getitem__(self, idx): | |
| image_filename = self.images_list[idx] | |
| noisy_im = cv2.imread(os.path.join(self.noisy_path, image_filename)) | |
| clean_im = cv2.imread(os.path.join(self.clean_path, image_filename)) | |
| arr1=np.array(clean_im) | |
| arr2=np.array(noisy_im) | |
| arr3 = arr1+ 1e-9 | |
| arr3 = np.divide(arr2,arr3) | |
| arr1 = cv2.resize(arr1, (256,256), interpolation=cv2.INTER_LINEAR) | |
| arr2= cv2.resize(arr2, (256,256), interpolation=cv2.INTER_LINEAR) | |
| arr3= cv2.resize(arr3, (256,256), interpolation=cv2.INTER_LINEAR) | |
| ## for grayscale images | |
| # arr1 = arr1[..., np.newaxis] | |
| # arr2 = arr2[..., np.newaxis] | |
| # arr3 = arr3[..., np.newaxis] | |
| # arr3 = np.square(arr3) | |
| # # for log data | |
| # arr1 = (arr1.astype(np.float32) + 1 )/256.0 | |
| # arr2 = (arr2.astype(np.float32) + 1 )/256.0 | |
| # arr1 = np.log(np.absolute(arr1)) | |
| # arr2 = np.log(np.absolute(arr2)) | |
| # # arr1 = arr1.astype(np.float32) / (0.5*np.log(256.0)) - 1 | |
| # # arr2 = arr2.astype(np.float32) / (0.5*np.log(256.0)) - 1 | |
| # arr1 = 2*(arr1.astype(np.float32) + np.log(256.0))/ np.log(256.0) - 1 | |
| # arr2 = 2*(arr2.astype(np.float32) + np.log(256.0))/ np.log(256.0) - 1 | |
| # ## correct normalization for log | |
| # arr1 = (arr1.astype(np.float32))/255.0 | |
| # arr2 = (arr2.astype(np.float32))/255.0 | |
| # arr1 = arr1*(math.exp(1)-math.exp(-1)) + math.exp(-1) | |
| # arr2 = arr2*(math.exp(1)-math.exp(-1)) + math.exp(-1) | |
| # arr1 = np.log(arr1) | |
| # arr2 = np.log(arr2) | |
| # arr1 = arr1.astype(np.float32) | |
| # arr2 = arr2.astype(np.float32) | |
| arr1 = arr1.astype(np.float32) / 127.5 - 1 | |
| arr2 = arr2.astype(np.float32) / 127.5 - 1 | |
| # arr3 = arr3.astype(np.float32) / 127.5 - 1 | |
| # arr3 = arr3.astype(np.float32) | |
| arr2 = np.transpose(arr2, [2, 0, 1]) | |
| arr1 = np.transpose(arr1, [2, 0, 1]) | |
| arr3 = np.transpose(arr3, [2, 0, 1]) | |
| # return arr3, {'SR': arr2, 'HR': arr1 , 'Index': image_filename} | |
| return arr1, {'SR': arr2, 'HR': arr1 , 'Index': image_filename} | |
| # return arr2, {'SR': arr2, 'HR': arr2 , 'Index': image_filename} | |
| # return arr1, {'noise': arr2, 'Index': image_filename} | |
| class ValDataNew(data.Dataset): | |
| def __init__(self, dataset_path, crop_size=[256,256]): | |
| super().__init__() | |
| # train_list = train_data_dir + train_filename | |
| # with open(train_list) as f: | |
| # contents = f.readlines() | |
| # input_names = [i.strip() for i in contents] | |
| # gt_names = [i.strip().replace('input','gt') for i in input_names] | |
| # self.train_data_dir = '/media/labuser/cb8bb1ad-451a-4aa4-870c-2d3eeafe2525/ECCV_2022/diffusion_ema_rain_imagenet/rain_sub1/' | |
| # self.train_data_dir = '/media/labuser/cb8bb1ad-451a-4aa4-870c-2d3eeafe2525/ICIP_Turbulence_files/Tubfaces89/300M/tubimages/' | |
| # self.train_data_dir = "/media/malsha/47a8802b-e0b7-47a8-8a4d-1649cc3ad408/sar_optical/optical/" | |
| # self.noisy_path = os.path.join(dataset_path, 'noisy') | |
| self.noisy_path = dataset_path | |
| self.clean_path = dataset_path | |
| # self.clean_path = os.path.join(dataset_path, 'clean') | |
| self.images_list = os.listdir(self.noisy_path) | |
| self.crop_size = crop_size | |
| def __len__(self): | |
| return len(os.listdir(self.noisy_path)) | |
| def __getitem__(self, idx): | |
| image_filename = self.images_list[idx] | |
| pil_image = cv2.imread(os.path.join(self.noisy_path, image_filename)) ## Clean image | |
| pil_image = cv2.cvtColor(pil_image, cv2.COLOR_BGR2GRAY) | |
| pil_image = np.repeat(pil_image[:,:,np.newaxis],3, axis=2) | |
| # print(pil_image.shape) | |
| im1 = ((np.float32(pil_image)+1.0)/256.0)**2 | |
| gamma_noise = seed.gamma(size=im1.shape, shape=1.0, scale=1.0).astype(im1.dtype) | |
| syn_sar = np.sqrt(im1 * gamma_noise) | |
| pil_image1 = syn_sar * 256-1 ## Noisy image | |
| arr1=np.array(pil_image) | |
| arr2=np.array(pil_image1) | |
| arr1 = cv2.resize(arr1, (256,256), interpolation=cv2.INTER_LINEAR) | |
| arr2= cv2.resize(arr2, (256,256), interpolation=cv2.INTER_LINEAR) | |
| arr1 = arr1.astype(np.float32) / 127.5 - 1 | |
| arr2 = arr2.astype(np.float32) / 127.5 - 1 | |
| arr2 = np.transpose(arr2, [2, 0, 1]) | |
| arr1 = np.transpose(arr1, [2, 0, 1]) | |
| return arr1, {'SR': arr2, 'HR': arr1 , 'Index': image_filename} | |
| class ValDataNewReal(data.Dataset): | |
| def __init__(self, dataset_path, crop_size=[256,256]): | |
| super().__init__() | |
| # train_list = train_data_dir + train_filename | |
| # with open(train_list) as f: | |
| # contents = f.readlines() | |
| # input_names = [i.strip() for i in contents] | |
| # gt_names = [i.strip().replace('input','gt') for i in input_names] | |
| # self.train_data_dir = '/media/labuser/cb8bb1ad-451a-4aa4-870c-2d3eeafe2525/ECCV_2022/diffusion_ema_rain_imagenet/rain_sub1/' | |
| # self.train_data_dir = '/media/labuser/cb8bb1ad-451a-4aa4-870c-2d3eeafe2525/ICIP_Turbulence_files/Tubfaces89/300M/tubimages/' | |
| # self.train_data_dir = "/media/malsha/47a8802b-e0b7-47a8-8a4d-1649cc3ad408/sar_optical/optical/" | |
| # self.noisy_path = os.path.join(dataset_path, 'noisy') | |
| self.noisy_path = dataset_path | |
| self.clean_path = dataset_path | |
| # self.clean_path = os.path.join(dataset_path, 'clean') | |
| self.images_list = os.listdir(self.noisy_path) | |
| self.crop_size = crop_size | |
| def __len__(self): | |
| return len(os.listdir(self.noisy_path)) | |
| def __getitem__(self, idx): | |
| image_filename = self.images_list[idx] | |
| pil_image = cv2.imread(os.path.join(self.noisy_path, image_filename),0) ## SAR image | |
| # pil_image = cv2.cvtColor(pil_image, cv2.COLOR_BGR2GRAY) | |
| pil_image = np.repeat(pil_image[:,:,np.newaxis],3, axis=2) | |
| # print(pil_image.shape) | |
| # im1 = ((np.float32(pil_image)+1.0)/256.0)**2 | |
| # gamma_noise = seed.gamma(size=im1.shape, shape=1.0, scale=1.0).astype(im1.dtype) | |
| # syn_sar = np.sqrt(im1 * gamma_noise) | |
| # pil_image1 = syn_sar * 256-1 ## Noisy image | |
| # pil_image = np.repeat(pil_image[:,:,np.newaxis],3, axis=2) | |
| # pil_image1 = np.repeat(pil_image1[:,:,np.newaxis],3, axis=2) | |
| arr1=np.array(pil_image) | |
| arr2=np.array(pil_image) | |
| arr3 = arr1 + 1e-9 | |
| # print(arr3.dtype) | |
| arr3 = np.divide(arr2,arr3) | |
| arr1 = cv2.resize(arr1, (256,256), interpolation=cv2.INTER_LINEAR) | |
| arr2= cv2.resize(arr2, (256,256), interpolation=cv2.INTER_LINEAR) | |
| arr3= cv2.resize(arr3, (256,256), interpolation=cv2.INTER_LINEAR) | |
| ## for grayscale images | |
| # arr1 = arr1[..., np.newaxis] | |
| # arr2 = arr2[..., np.newaxis] | |
| # arr3 = arr3[..., np.newaxis] | |
| # arr3 = np.square(arr3) | |
| # # for log data | |
| # arr1 = (arr1.astype(np.float32) + 1 )/256.0 | |
| # arr2 = (arr2.astype(np.float32) + 1 )/256.0 | |
| # arr1 = np.log(np.absolute(arr1)) | |
| # arr2 = np.log(np.absolute(arr2)) | |
| # # arr1 = arr1.astype(np.float32) / (0.5*np.log(256.0)) - 1 | |
| # # arr2 = arr2.astype(np.float32) / (0.5*np.log(256.0)) - 1 | |
| # arr1 = 2*(arr1.astype(np.float32) + np.log(256.0))/ np.log(256.0) - 1 | |
| # arr2 = 2*(arr2.astype(np.float32) + np.log(256.0))/ np.log(256.0) - 1 | |
| # ## correct normalization for log | |
| # arr1 = (arr1.astype(np.float32))/255.0 | |
| # arr2 = (arr2.astype(np.float32))/255.0 | |
| # arr1 = arr1*(math.exp(1)-math.exp(-1)) + math.exp(-1) | |
| # arr2 = arr2*(math.exp(1)-math.exp(-1)) + math.exp(-1) | |
| # arr1 = np.log(arr1) | |
| # arr2 = np.log(arr2) | |
| # arr1 = arr1.astype(np.float32) | |
| # arr2 = arr2.astype(np.float32) | |
| arr1 = arr1.astype(np.float32) / 127.5 - 1 | |
| arr2 = arr2.astype(np.float32) / 127.5 - 1 | |
| # arr3 = arr3.astype(np.float32) / 127.5 - 1 | |
| # arr3 = arr3.astype(np.float32) | |
| arr2 = np.transpose(arr2, [2, 0, 1]) | |
| arr1 = np.transpose(arr1, [2, 0, 1]) | |
| arr3 = np.transpose(arr3, [2, 0, 1]) | |
| # return arr3, {'SR': arr2, 'HR': arr1 , 'Index': image_filename} | |
| return arr1, {'SR': arr2, 'HR': arr1 , 'Index': image_filename} | |