| 17 | matplotlib.use('TkAgg') |
| 18 | |
| 19 | class ChangeSim(data.Dataset): |
| 20 | def __init__(self, crop_size=(320, 240), num_classes=5, set='train'): |
| 21 | """ |
| 22 | ChangeSim Dataloader for Visualization |
| 23 | Please download ChangeSim Dataset in https://github.com/SAMMiCA/ChangeSim |
| 24 | |
| 25 | Args: |
| 26 | crop_size (tuple): Image resize shape (H,W) (default: (320, 240)) |
| 27 | num_classes (int): Number of target change detection class |
| 28 | 5 for multi-class change detection |
| 29 | 2 for binary change detection (default: 5) |
| 30 | set (str): 'train' or 'test' (defalut: 'train') |
| 31 | """ |
| 32 | self.crop_size = crop_size |
| 33 | self.num_classes = num_classes |
| 34 | self.set = set |
| 35 | train_list = ['Warehouse_0', 'Warehouse_1', 'Warehouse_2', 'Warehouse_3', 'Warehouse_4', 'Warehouse_5'] |
| 36 | test_list = ['Warehouse_6', 'Warehouse_7', 'Warehouse_8', 'Warehouse_9'] |
| 37 | self.image_total_files = [] |
| 38 | if set == 'train': |
| 39 | for map in train_list: |
| 40 | self.image_total_files += glob.glob('../Query/Query_Seq_Train/' + map + '/Seq_0/rgb/*.png') |
| 41 | self.image_total_files += glob.glob('../Query/Query_Seq_Train/' + map + '/Seq_1/rgb/*.png') |
| 42 | elif set == 'test': |
| 43 | for map in test_list: |
| 44 | self.image_total_files += glob.glob('../Query/Query_Seq_Test/' + map + '/Seq_0/rgb/*.png') |
| 45 | self.image_total_files += glob.glob('../Query/Query_Seq_Test/' + map + '/Seq_1/rgb/*.png') |
| 46 | |
| 47 | def __len__(self): |
| 48 | return len(self.image_total_files) |
| 49 | |
| 50 | def __getitem__(self, index): |
| 51 | # Get File Paths |
| 52 | test_rgb_path = self.image_total_files[index] |
| 53 | test_depth_path = test_rgb_path.replace('rgb', 'depth') |
| 54 | test_segmentation_path = test_rgb_path.replace('rgb', 'semantic_segmentation') |
| 55 | ref_rgb_path = test_rgb_path.replace('rgb', 't0/rgb') |
| 56 | ref_depth_path = test_rgb_path.replace('rgb', 't0/depth') |
| 57 | change_segmentation_path = test_rgb_path.replace('rgb', 'change_segmentation') |
| 58 | name = '_'.join(test_rgb_path.split('/')[-5:]) |
| 59 | |
| 60 | # RGB |
| 61 | test_rgb = Image.open(test_rgb_path) |
| 62 | ref_rgb = Image.open(ref_rgb_path) |
| 63 | test_rgb = test_rgb.resize(self.crop_size, Image.BICUBIC) |
| 64 | ref_rgb = ref_rgb.resize(self.crop_size, Image.BICUBIC) |
| 65 | test_rgb = ToTensor()(test_rgb) |
| 66 | ref_rgb = ToTensor()(ref_rgb) |
| 67 | |
| 68 | # Depth |
| 69 | test_depth = Image.open(test_depth_path) |
| 70 | test_depth = test_depth.resize(self.crop_size, Image.BICUBIC) |
| 71 | test_depth = np.asarray(test_depth) |
| 72 | test_depth = test_depth.astype('float32') / 255 |
| 73 | test_depth = torch.from_numpy(test_depth) |
| 74 | |
| 75 | ref_depth = Image.open(ref_depth_path) |
| 76 | ref_depth = ref_depth.resize(self.crop_size, Image.BICUBIC) |