| 137 | |
| 138 | |
| 139 | class test_dataset: |
| 140 | def __init__(self, image_root, gt_root, testsize): |
| 141 | self.testsize = testsize |
| 142 | self.images = [image_root + f for f in os.listdir(image_root) if f.endswith('.jpg') or f.endswith('.png')] |
| 143 | self.gts = [gt_root + f for f in os.listdir(gt_root) if f.endswith('.tif') or f.endswith('.png' or f.endswith('.jpg'))] |
| 144 | self.images = sorted(self.images) |
| 145 | self.gts = sorted(self.gts) |
| 146 | self.transform = transforms.Compose([ |
| 147 | transforms.Resize((self.testsize, self.testsize)), |
| 148 | transforms.ToTensor(), |
| 149 | transforms.Normalize([0.485, 0.456, 0.406], |
| 150 | [0.229, 0.224, 0.225])]) |
| 151 | self.gt_transform = transforms.ToTensor() |
| 152 | self.size = len(self.images) |
| 153 | self.index = 0 |
| 154 | |
| 155 | def load_data(self): |
| 156 | image = self.rgb_loader(self.images[self.index]) |
| 157 | image = self.transform(image).unsqueeze(0) |
| 158 | gt = self.binary_loader(self.gts[self.index]) |
| 159 | name = self.images[self.index].split('/')[-1] |
| 160 | if name.endswith('.jpg'): |
| 161 | name = name.split('.jpg')[0] + '.png' |
| 162 | self.index += 1 |
| 163 | return image, gt, name |
| 164 | |
| 165 | def rgb_loader(self, path): |
| 166 | with open(path, 'rb') as f: |
| 167 | img = Image.open(f) |
| 168 | return img.convert('RGB') |
| 169 | |
| 170 | def binary_loader(self, path): |
| 171 | with open(path, 'rb') as f: |
| 172 | img = Image.open(f) |
| 173 | return img.convert('L') |
nothing calls this directly
no outgoing calls
no test coverage detected