| | 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' |
| |
|
| |
|
