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 transform_val(self, sample):
composed_transforms = transforms.Compose([
tr.FixScaleCrop(crop_size=self.args.crop_size),
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'.")
