def get_data_transforms(self):
cfg = self.cfg
bbox_params = self.get_bbox_params()
transform = Compose([Resize(cfg.data.img_sz, cfg.data.img_sz)],
bbox_params=bbox_params)
augmentors_dict = {
'Blur': Blur(),
'RandomRotate90': RandomRotate90(),
'HorizontalFlip': HorizontalFlip(),
'VerticalFlip': VerticalFlip(),
'GaussianBlur': GaussianBlur(),
'GaussNoise': GaussNoise(),
'RGBShift': RGBShift(),
'ToGray': ToGray()
}
aug_transforms = []
for augmentor in cfg.data.augmentors:
try:
aug_transforms.append(augmentors_dict[augmentor])
except KeyError as e:
log.warning(
'{0} is an unknown augmentor. Continuing without {0}. \
Known augmentors are: {1}'.format(
e, list(augmentors_dict.keys())))
aug_transforms.append(Resize(cfg.data.img_sz, cfg.data.img_sz))
aug_transform = Compose(aug_transforms, bbox_params=bbox_params)
return transform, aug_transform
def get_data_transforms(self) -> Tuple[BasicTransform, BasicTransform]:
"""Get albumentations transform objects for data augmentation.
Returns:
1st tuple arg: a transform that doesn't do any data augmentation
2nd tuple arg: a transform with data augmentation
"""
cfg = self.cfg
bbox_params = self.get_bbox_params()
transform = Compose([Resize(cfg.data.img_sz, cfg.data.img_sz)],
bbox_params=bbox_params)
augmentors_dict = {
'Blur': Blur(),
'RandomRotate90': RandomRotate90(),
'HorizontalFlip': HorizontalFlip(),
'VerticalFlip': VerticalFlip(),
'GaussianBlur': GaussianBlur(),
'GaussNoise': GaussNoise(),
'RGBShift': RGBShift(),
'ToGray': ToGray()
}
aug_transforms = []
for augmentor in cfg.data.augmentors:
try:
aug_transforms.append(augmentors_dict[augmentor])
except KeyError as e:
log.warning(
'{0} is an unknown augmentor. Continuing without {0}. \
Known augmentors are: {1}'.format(
e, list(augmentors_dict.keys())))
aug_transforms.append(Resize(cfg.data.img_sz, cfg.data.img_sz))
aug_transform = Compose(aug_transforms, bbox_params=bbox_params)
return transform, aug_transform
def build_databunch(data_dir, img_sz, batch_sz, class_names, augmentors):
num_workers = 4
train_dir = join(data_dir, 'train')
valid_dir = join(data_dir, 'valid')
augmentors_dict = {
'Blur': Blur(),
'RandomRotate90': RandomRotate90(),
'HorizontalFlip': HorizontalFlip(),
'VerticalFlip': VerticalFlip(),
'GaussianBlur': GaussianBlur(),
'GaussNoise': GaussNoise(),
'RGBShift': RGBShift(),
'ToGray': ToGray()
}
aug_transforms = []
for augmentor in augmentors:
try:
aug_transforms.append(augmentors_dict[augmentor])
except KeyError as e:
log.warning('{0} is an unknown augmentor. Continuing without {0}. \
Known augmentors are: {1}'.format(e,
list(
augmentors_dict.keys())))
aug_transforms = Compose(aug_transforms)
train_ds = AlbumentationDataset(ImageFolder(train_dir,
classes=class_names),
transform=aug_transforms)
valid_ds = AlbumentationDataset(ImageFolder(valid_dir,
classes=class_names))
train_dl = DataLoader(train_ds,
shuffle=True,
batch_size=batch_sz,
num_workers=num_workers,
drop_last=True,
pin_memory=True)
valid_dl = DataLoader(valid_ds,
batch_size=batch_sz,
num_workers=num_workers,
pin_memory=True)
return DataBunch(train_ds, train_dl, valid_ds, valid_dl, class_names)
def build_databunch(data_dir, img_sz, batch_sz, class_names, rare_classes,
desired_prob, augmentors):
num_workers = 4
train_dir = join(data_dir, 'train')
valid_dir = join(data_dir, 'valid')
augmentors_dict = {
'Blur': Blur(),
'RandomRotate90': RandomRotate90(),
'HorizontalFlip': HorizontalFlip(),
'VerticalFlip': VerticalFlip(),
'GaussianBlur': GaussianBlur(),
'GaussNoise': GaussNoise(),
'RGBShift': RGBShift(),
'ToGray': ToGray()
}
aug_transforms = []
for augmentor in augmentors:
try:
aug_transforms.append(augmentors_dict[augmentor])
except KeyError as e:
log.warning('{0} is an unknown augmentor. Continuing without {0}. \
Known augmentors are: {1}'.format(e,
list(
augmentors_dict.keys())))
aug_transforms = Compose(aug_transforms)
train_ds = AlbumentationDataset(ImageFolder(train_dir,
classes=class_names),
transform=aug_transforms)
valid_ds = AlbumentationDataset(ImageFolder(valid_dir,
classes=class_names))
if rare_classes != []:
targets = [target for _, target in train_ds.orig_dataset.imgs]
train_sample_weights = calculate_oversampling_weights(
targets, rare_classes, desired_prob)
num_train_samples = len(train_ds)
train_sampler = WeightedRandomSampler(weights=train_sample_weights,
num_samples=num_train_samples,
replacement=True)
shuffle = False
else:
train_sampler = None
shuffle = True
train_dl = DataLoader(train_ds,
shuffle=shuffle,
batch_size=batch_sz,
num_workers=num_workers,
drop_last=True,
pin_memory=True,
sampler=train_sampler)
valid_dl = DataLoader(valid_ds,
batch_size=batch_sz,
num_workers=num_workers,
pin_memory=True)
return DataBunch(train_ds, train_dl, valid_ds, valid_dl, class_names)