| 13 | |
| 14 | |
| 15 | class MyImageFolder(Dataset): |
| 16 | def __init__(self, root_dir): |
| 17 | super(MyImageFolder, self).__init__() |
| 18 | self.data = [] |
| 19 | self.root_dir = root_dir |
| 20 | self.class_names = os.listdir(root_dir) |
| 21 | |
| 22 | for index, name in enumerate(self.class_names): |
| 23 | files = os.listdir(os.path.join(root_dir, name)) |
| 24 | self.data += list(zip(files, [index] * len(files))) |
| 25 | # print('self.data: {}'.format(self.data)) |
| 26 | # break |
| 27 | |
| 28 | def __len__(self): |
| 29 | return len(self.data) |
| 30 | |
| 31 | def __getitem__(self, index): |
| 32 | img_file, label = self.data[index] |
| 33 | #得到图像的完整路径 |
| 34 | root_and_dir = os.path.join(self.root_dir, self.class_names[label]) |
| 35 | |
| 36 | #将读取的图像转换为numpy形式 |
| 37 | image = np.array(Image.open(os.path.join(root_and_dir, img_file))) |
| 38 | #对图像进行预处理 |
| 39 | image = config.both_transforms(image=image)["image"] |
| 40 | #对高分辨图像进行处理 |
| 41 | high_res = config.highres_transform(image=image)["image"] |
| 42 | #将[96 x 96] => [24 x 24] |
| 43 | low_res = config.lowres_transform(image=image)["image"] |
| 44 | return low_res, high_res |
| 45 | |
| 46 | |
| 47 | def test(): |