Unified adaptive dataloader for polyp segmentation. Uses Albumentations and strictly handles binary mask conversion.
| 8 | from albumentations.pytorch import ToTensorV2 |
| 9 | |
| 10 | class PolypDataset(data.Dataset): |
| 11 | """ |
| 12 | Unified adaptive dataloader for polyp segmentation. |
| 13 | Uses Albumentations and strictly handles binary mask conversion. |
| 14 | """ |
| 15 | def __init__(self, image_root, gt_root, trainsize, augmentation, split='train', color_image=True): |
| 16 | self.trainsize = trainsize |
| 17 | self.color_image = color_image |
| 18 | self.augmentation = augmentation |
| 19 | self.split = split |
| 20 | |
| 21 | # Load and sort file paths |
| 22 | exts = ('.jpg', '.png', '.jpeg', '.tif') |
| 23 | self.images = sorted([os.path.join(image_root, f) for f in os.listdir(image_root) if f.lower().endswith(exts)]) |
| 24 | self.gts = sorted([os.path.join(gt_root, f) for f in os.listdir(gt_root) if f.lower().endswith(exts)]) |
| 25 | |
| 26 | self.filter_files() |
| 27 | self.size = len(self.images) |
| 28 | |
| 29 | # Transformation Setup |
| 30 | mean = [0.485, 0.456, 0.406] if color_image else [0.5] |
| 31 | std = [0.229, 0.224, 0.225] if color_image else [0.229] |
| 32 | |
| 33 | if self.split == 'train' and self.augmentation: |
| 34 | self.transform = A.Compose([ |
| 35 | A.Rotate(limit=90, p=0.5), |
| 36 | A.VerticalFlip(p=0.5), |
| 37 | A.HorizontalFlip(p=0.5), |
| 38 | A.Resize(height=self.trainsize, width=self.trainsize), |
| 39 | A.Normalize(mean=mean, std=std), |
| 40 | ToTensorV2() |
| 41 | ]) |
| 42 | else: |
| 43 | self.transform = A.Compose([ |
| 44 | A.Resize(height=self.trainsize, width=self.trainsize), |
| 45 | A.Normalize(mean=mean, std=std), |
| 46 | ToTensorV2() |
| 47 | ]) |
| 48 | |
| 49 | def filter_files(self): |
| 50 | valid_images, valid_gts = [], [] |
| 51 | for img_p, gt_p in zip(self.images, self.gts): |
| 52 | if os.path.exists(img_p) and os.path.exists(gt_p): |
| 53 | valid_images.append(img_p) |
| 54 | valid_gts.append(gt_p) |
| 55 | self.images, self.gts = valid_images, valid_gts |
| 56 | |
| 57 | def __getitem__(self, index): |
| 58 | # 1. Load Image |
| 59 | image = cv2.imread(self.images[index]) |
| 60 | image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB if self.color_image else cv2.COLOR_BGR2GRAY) |
| 61 | |
| 62 | # 2. Load Mask |
| 63 | mask_np = cv2.imread(self.gts[index], cv2.IMREAD_GRAYSCALE) |
| 64 | |
| 65 | # 3. Apply Transformations |
| 66 | augmented = self.transform(image=image, mask=mask_np) |
| 67 | image = augmented['image'] |