Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import os | |
| from torch.utils.data import Dataset | |
| from PIL import Image | |
| from scipy.ndimage import gaussian_filter | |
| class RGBDD_Dataset(Dataset): | |
| """RGB-D-D Dataset.""" | |
| def __init__(self, root_dir="./dataset/RGB-D-D/", scale=4, downsample='real', train=True, | |
| 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 | |
| types = ['models', 'plants', 'portraits'] | |
| if train: | |
| if self.downsample == 'real': | |
| self.GTs = [] | |
| self.LRs = [] | |
| self.RGBs = [] | |
| for type in types: | |
| list_dir = os.listdir('%s/%s/%s_train'% (root_dir, type, type)) | |
| for n in list_dir: | |
| self.RGBs.append('%s/%s/%s_train/%s/%s_RGB.jpg' % (root_dir, type, type, n, n)) | |
| self.GTs.append('%s/%s/%s_train/%s/%s_HR_gt.png' % (root_dir, type, type, n, n)) | |
| self.LRs.append('%s/%s/%s_train/%s/%s_LR_fill_depth.png' % (root_dir, type, type, n, n)) | |
| else: | |
| self.GTs = [] | |
| self.RGBs = [] | |
| for type in types: | |
| list_dir = os.listdir('%s/%s/%s_train'% (root_dir, type, type)) | |
| for n in list_dir: | |
| self.RGBs.append('%s/%s/%s_train/%s/%s_RGB.jpg' % (root_dir, type, type, n, n)) | |
| self.GTs.append('%s/%s/%s_train/%s/%s_HR_gt.png' % (root_dir, type, type, n, n)) | |
| else: | |
| if self.downsample == 'real': | |
| self.GTs = [] | |
| self.LRs = [] | |
| self.RGBs = [] | |
| for type in types: | |
| list_dir = os.listdir('%s/%s/%s_test'% (root_dir, type, type)) | |
| for n in list_dir: | |
| self.RGBs.append('%s/%s/%s_test/%s/%s_RGB.jpg' % (root_dir, type, type, n, n)) | |
| self.GTs.append('%s/%s/%s_test/%s/%s_HR_gt.png' % (root_dir, type, type, n, n)) | |
| self.LRs.append('%s/%s/%s_test/%s/%s_LR_fill_depth.png' % (root_dir, type, type, n, n)) | |
| else: | |
| self.GTs = [] | |
| self.RGBs = [] | |
| for type in types: | |
| list_dir = os.listdir('%s/%s/%s_test'% (root_dir, type, type)) | |
| for n in list_dir: | |
| self.RGBs.append('%s/%s/%s_test/%s/%s_RGB.jpg' % (root_dir, type, type, n, n)) | |
| self.GTs.append('%s/%s/%s_test/%s/%s_HR_gt.png' % (root_dir, type, type, n, n)) | |
| def __len__(self): | |
| return len(self.GTs) | |
| def __getitem__(self, idx): | |
| if self.downsample == 'real': | |
| image = np.array(Image.open(self.RGBs[idx]).convert("RGB")).astype(np.float32) | |
| name = self.RGBs[idx][-22:-8] | |
| gt = np.array(Image.open(self.GTs[idx])).astype(np.float32) | |
| h, w = gt.shape | |
| lr = np.array(Image.open(self.LRs[idx]).resize((w, h), Image.BICUBIC)).astype(np.float32) | |
| else: | |
| image = Image.open(self.RGBs[idx]).convert("RGB") | |
| name = self.RGBs[idx][-22:-8] | |
| image = np.array(image).astype(np.float32) | |
| gt = Image.open(self.GTs[idx]) | |
| w, h = gt.size | |
| s = self.scale | |
| lr = np.array(gt.resize((w // s, h // s), Image.BICUBIC).resize((w, h), Image.BICUBIC)).astype(np.float32) | |
| gt = np.array(gt).astype(np.float32) | |
| # 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) | |
| maxx = np.max(image) | |
| minn = np.min(image) | |
| image = (image - minn) / (maxx - minn) | |
| 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) | |
| 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 | |