Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import os | |
| from torch.utils.data import Dataset, DataLoader | |
| from PIL import Image | |
| from scipy.ndimage import gaussian_filter | |
| class TOFDSR_Dataset(Dataset): | |
| def __init__(self, root_dir="./dataset/", scale=4, downsample='real', train=True, txt_file='./TOFDSR_Train.txt' , | |
| transform=None, isNoisy=False, blur_sigma=1.2): | |
| self.root_dir = root_dir | |
| self.transform = transform | |
| self.scale = scale | |
| self.downsample = downsample | |
| self.train = train | |
| self.isNoisy = isNoisy | |
| self.blur_sigma = blur_sigma | |
| self.image_list = txt_file | |
| with open(self.image_list, 'r') as f: | |
| self.filename = f.readlines() | |
| def __len__(self): | |
| return len(self.filename) | |
| def __getitem__(self, idx): | |
| sample_path = self.filename[idx].strip('\n') | |
| sample_path_ = sample_path.split(',') | |
| rgb_path = sample_path_[0] | |
| gt_path = sample_path_[1] | |
| lr_path = sample_path_[2] | |
| name = gt_path[20:-4] | |
| rgb_path = os.path.join(self.root_dir, rgb_path) | |
| gt_path = os.path.join(self.root_dir, gt_path) | |
| lr_path = os.path.join(self.root_dir, lr_path) | |
| if self.downsample == 'real': | |
| image = np.array(Image.open(rgb_path).convert("RGB")).astype(np.float32) | |
| gt = np.array(Image.open(gt_path)).astype(np.float32) | |
| h, w = gt.shape | |
| lr = np.array(Image.open(lr_path).resize((w, h), Image.BICUBIC)).astype(np.float32) | |
| else: | |
| image = np.array(Image.open(rgb_path).convert("RGB")).astype(np.float32) | |
| gt = Image.open(gt_path) | |
| w, h = gt.size | |
| lr = np.array(gt.resize((w, h), Image.BICUBIC)).astype(np.float32) | |
| gt = np.array(gt).astype(np.float32) | |
| image_max = np.max(image) | |
| image_min = np.min(image) | |
| image = (image - image_min) / (image_max - image_min) | |
| # normalization | |
| if self.train: | |
| max_out = 5000.0 | |
| min_out = 0.0 | |
| lr = (lr - min_out) / (max_out - min_out) | |
| gt = (gt-min_out)/(max_out-min_out) | |
| else: | |
| max_out = 5000.0 | |
| min_out = 0.0 | |
| lr = (lr - min_out) / (max_out - min_out) | |
| lr_minn = np.min(lr) | |
| lr_maxx = np.max(lr) | |
| if not self.train: | |
| np.random.seed(42) | |
| if self.isNoisy: | |
| lr = gaussian_filter(lr, sigma=self.blur_sigma) | |
| gaussian_noise = np.random.normal(0, 0.07, lr.shape) | |
| lr = lr + gaussian_noise | |
| lr = np.clip(lr, lr_minn, lr_maxx) | |
| if self.transform: | |
| image = self.transform(image).float() | |
| gt = self.transform(np.expand_dims(gt, 2)).float() | |
| lr = self.transform(np.expand_dims(lr, 2)).float() | |
| sample = {'guidance': image, 'lr': lr, 'gt': gt, 'max': max_out, 'min': min_out,'name': name} | |
| return sample | |