| import os |
| from data import srdata |
|
|
| class DIV2K(srdata.SRData): |
| def __init__(self, args, name='DIV2K', train=True, benchmark=False): |
| data_range = [r.split('-') for r in args.data_range.split('/')] |
| if train: |
| data_range = data_range[0] |
| else: |
| if args.test_only and len(data_range) == 1: |
| data_range = data_range[0] |
| else: |
| data_range = data_range[1] |
|
|
| self.begin, self.end = list(map(lambda x: int(x), data_range)) |
| super(DIV2K, self).__init__( |
| args, name=name, train=train, benchmark=benchmark |
| ) |
|
|
| def _scan(self): |
| names_hr, names_edge, names_lr = super(DIV2K, self)._scan() |
| names_hr = names_hr[self.begin - 1:self.end] |
| names_edge = names_edge[self.begin - 1:self.end] |
| names_lr = [n[self.begin - 1:self.end] for n in names_lr] |
|
|
| return names_hr, names_edge, names_lr |
|
|
| def _set_filesystem(self, dir_data): |
| super(DIV2K, self)._set_filesystem(dir_data) |
| self.dir_hr = os.path.join(self.apath, 'DIV2K_train_HR') |
| self.dir_edge = os.path.join(self.apath, 'DIV2K_train_EDGE_disturbed') |
| self.dir_lr = os.path.join(self.apath, 'DIV2K_train_LR_bicubic') |
| if self.input_large: self.dir_lr += 'L' |
|
|
|
|
