(self, image_root, gt_root, trainsize)
| 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]) |
nothing calls this directly
no test coverage detected