Spaces:
Runtime error
Runtime error
File size: 4,879 Bytes
68c537d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
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
|