| 35 | |
| 36 | |
| 37 | class SrDataset(Dataset): |
| 38 | def __init__(self, hr_dir: str, lr_dir: str): |
| 39 | self.input_size = np.asanyarray([224, 224]) |
| 40 | self.hr = [] |
| 41 | self.lr = [] |
| 42 | |
| 43 | for file in sorted(os.listdir(hr_dir)): |
| 44 | self.hr.append(self._resize_img(os.path.join(hr_dir, file), 2)) |
| 45 | |
| 46 | for file in sorted(os.listdir(lr_dir)): |
| 47 | self.lr.append(self._resize_img(os.path.join(lr_dir, file), 1)) |
| 48 | |
| 49 | if len(self.hr) != len(self.lr): |
| 50 | raise AssertionError( |
| 51 | "The number of high resolution pics is not equal to low " |
| 52 | "resolution pics" |
| 53 | ) |
| 54 | |
| 55 | def __getitem__(self, idx: int): |
| 56 | return self.hr[idx], self.lr[idx] |
| 57 | |
| 58 | def __len__(self): |
| 59 | return len(self.lr) |
| 60 | |
| 61 | def _resize_img(self, file: str, scale: int): |
| 62 | with Image.open(file) as img: |
| 63 | return to_tensor(img.resize(tuple(self.input_size * scale))).unsqueeze(0) |
| 64 | |
| 65 | |
| 66 | def get_b100( |
no outgoing calls
no test coverage detected