-
Notifications
You must be signed in to change notification settings - Fork 0
/
datasets.py
executable file
·32 lines (26 loc) · 1.09 KB
/
datasets.py
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
import h5py
import numpy as np
from torch.utils.data import Dataset
'''
Source: https://github.com/yjn870/SRCNN-pytorch/blob/master/datasets.py
'''
class TrainDataset(Dataset):
def __init__(self, h5_file):
super(TrainDataset, self).__init__()
self.h5_file = h5_file
def __getitem__(self, idx):
with h5py.File(self.h5_file, 'r') as f:
return np.expand_dims(f['lr'][idx] / 255., 0), np.expand_dims(f['hr'][idx] / 255., 0)
def __len__(self):
with h5py.File(self.h5_file, 'r') as f:
return len(f['lr'])
class EvalDataset(Dataset):
def __init__(self, h5_file):
super(EvalDataset, self).__init__()
self.h5_file = h5_file
def __getitem__(self, idx):
with h5py.File(self.h5_file, 'r') as f:
return np.expand_dims(f['lr'][str(idx)][:, :] / 255., 0), np.expand_dims(f['hr'][str(idx)][:, :] / 255., 0) #difference is we return Cb and Cr channels to perform evaluation. Only Y channel is fed to SR model
def __len__(self):
with h5py.File(self.h5_file, 'r') as f:
return len(f['lr'])