def transform(self, sample): composed_transforms = transforms.Compose([ tr.RandomHorizontalFlip(), tr.RandomScaleCrop(base_size=self.args.base_size, crop_size=self.args.crop_size), tr.RandomGaussianBlur(), tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), tr.ToTensor()]) return composed_transforms(sample)
def __init__(self, base_dir=settings.root_dir, split='train'): super().__init__() self.split = split self.images = [] self.masks = [] file_list = Path(base_dir) / f'{split}.txt' with open(file_list, 'r') as f: lines = f.read().splitlines() for line in lines: image, mask = line.split(',') image = Path(base_dir) / image mask = Path(base_dir) / mask assert image.exists(), f'File not found: {image}' assert mask.exists(), f'File not found: {mask}' self.images.append(image) self.masks.append(mask) print(f'{split}| images: {len(self.images)}, masks: {len(self.masks)}') if split == 'train': self.composed_transform = transforms.Compose([ tr.FixedResize(settings.resize_height, settings.resize_width), tr.RandomHorizontalFlip(), # tr.RandomRotate(90), # tr.RandomScaleCrop(base_size=settings.base_size, crop_size=settings.crop_size), tr.RandomGaussianBlur(), tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), tr.ToTensor() ]) elif split == 'valid': self.composed_transform = transforms.Compose([ # tr.FixScaleCrop(crop_size=settings.crop_size), tr.FixedResize(settings.resize_height, settings.resize_width), tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), tr.ToTensor() ]) else: raise KeyError("split must be one of 'train' or 'valid'.")