(self, image_root, gt_root, trainsize, augmentation, split='train', color_image=True)
| 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 = [], [] |
nothing calls this directly
no test coverage detected