| 13 | from torch.utils.data import Dataset |
| 14 | |
| 15 | class MapDataset(Dataset): |
| 16 | def __init__(self,root_dir): |
| 17 | super(MapDataset, self).__init__() |
| 18 | |
| 19 | self.root_dir = root_dir |
| 20 | self.list_files = os.listdir(self.root_dir) |
| 21 | |
| 22 | def __len__(self): |
| 23 | return len(self.list_files) |
| 24 | |
| 25 | def __getitem__(self, index): |
| 26 | img_file = self.list_files[index] |
| 27 | img_path = os.path.join(self.root_dir,img_file) |
| 28 | image = np.array(Image.open(img_path)) |
| 29 | #图片由两张图片拼接的大小为1200 x 1200大小， |
| 30 | input_image = image[:,:600,:] |
| 31 | target_image = image[:,600:,:] |
| 32 | #图像增强 |
| 33 | augmentations = config.both_transform( |
| 34 | image = input_image, |
| 35 | image0 = target_image |
| 36 | ) |
| 37 | input_image,target_image = augmentations["image"],augmentations["image0"] |
| 38 | |
| 39 | input_image = config.transform_only_input(image = input_image)["image"] |
| 40 | target_image = config.transform_only_mask(image = target_image)["image"] |
| 41 | |
| 42 | return input_image,target_image |