File size: 900 Bytes
8ec10cf |
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 |
from data.dataset import Dataset
from utils import interact
class REDS(Dataset):
"""REDS train, val, test subset class
"""
def __init__(self, args, mode='train'):
super(REDS, self).__init__(args, mode)
def set_modes(self):
self.modes = ('train', 'val', 'test')
def set_keys(self):
super(REDS, self).set_keys()
# self.blur_key = 'blur'
# self.sharp_key = 'sharp'
self.non_blur_keys = ['blur', 'blur_comp', 'blur_bicubic']
self.non_blur_keys.remove(self.blur_key)
self.non_sharp_keys = ['sharp_bicubic', 'sharp']
self.non_sharp_keys.remove(self.sharp_key)
def __getitem__(self, idx):
blur, sharp, pad_width, idx, relpath = super(REDS, self).__getitem__(idx)
relpath = relpath.replace('{}/{}/'.format(self.mode, self.blur_key), '')
return blur, sharp, pad_width, idx, relpath
|