dataloader for polyp segmentation tasks
| 5 | |
| 6 | |
| 7 | class PolypDataset(data.Dataset): |
| 8 | """ |
| 9 | dataloader for polyp segmentation tasks |
| 10 | """ |
| 11 | def __init__(self, image_root, gt_root, trainsize): |
| 12 | self.trainsize = trainsize |
| 13 | self.images = [image_root + f for f in os.listdir(image_root) if f.endswith('.jpg') or f.endswith('.png')] |
| 14 | self.gts = [gt_root + f for f in os.listdir(gt_root) if f.endswith('.png')] |
| 15 | self.images = sorted(self.images) |
| 16 | self.gts = sorted(self.gts) |
| 17 | self.filter_files() |
| 18 | self.size = len(self.images) |
| 19 | self.img_transform = transforms.Compose([ |
| 20 | transforms.Resize((self.trainsize, self.trainsize)), |
| 21 | transforms.ToTensor(), |
| 22 | transforms.Normalize([0.485, 0.456, 0.406], |
| 23 | [0.229, 0.224, 0.225])]) |
| 24 | self.gt_transform = transforms.Compose([ |
| 25 | transforms.Resize((self.trainsize, self.trainsize)), |
| 26 | transforms.ToTensor()]) |
| 27 | |
| 28 | def __getitem__(self, index): |
| 29 | image = self.rgb_loader(self.images[index]) |
| 30 | gt = self.binary_loader(self.gts[index]) |
| 31 | image = self.img_transform(image) |
| 32 | gt = self.gt_transform(gt) |
| 33 | return image, gt |
| 34 | |
| 35 | def filter_files(self): |
| 36 | assert len(self.images) == len(self.gts) |
| 37 | images = [] |
| 38 | gts = [] |
| 39 | for img_path, gt_path in zip(self.images, self.gts): |
| 40 | img = Image.open(img_path) |
| 41 | gt = Image.open(gt_path) |
| 42 | if img.size == gt.size: |
| 43 | images.append(img_path) |
| 44 | gts.append(gt_path) |
| 45 | self.images = images |
| 46 | self.gts = gts |
| 47 | |
| 48 | def rgb_loader(self, path): |
| 49 | with open(path, 'rb') as f: |
| 50 | img = Image.open(f) |
| 51 | return img.convert('RGB') |
| 52 | |
| 53 | def binary_loader(self, path): |
| 54 | with open(path, 'rb') as f: |
| 55 | img = Image.open(f) |
| 56 | # return img.convert('1') |
| 57 | return img.convert('L') |
| 58 | |
| 59 | def resize(self, img, gt): |
| 60 | assert img.size == gt.size |
| 61 | w, h = img.size |
| 62 | if h < self.trainsize or w < self.trainsize: |
| 63 | h = max(h, self.trainsize) |
| 64 | w = max(w, self.trainsize) |