def __init__(self,
root,
input_transform=None,
target_transform=None,
_train=False):
self.images_root = os.path.join(root, 'Images')
self.labels_root = os.path.join(root, 'Labels')
self._train = _train
self.filenames = [
image_basename(f) for f in os.listdir(self.labels_root)
if is_image(f)
]
random.shuffle(self.filenames)
self.input_transform = input_transform
self.target_transform = target_transform
self.A_transform = A.Compose([
A.Rotate(limit=40, p=0.9),
A.HorizontalFlip(p=0.5),
A.RGBShift(r_shift_limit=25,
g_shift_limit=25,
b_shift_limit=25,
p=0.9),
A.OneOf([A.Blur(blur_limit=2, p=0.5),
A.ColorJitter(p=0.5)], p=1.0)
])
def get_augmentations():
"""Get a list of 'major' and 'minor' augmentation functions for the pipeline in a dictionary."""
return {
"major": {
"shift-scale-rot":
trans.ShiftScaleRotate(
shift_limit=0.05,
rotate_limit=35,
border_mode=cv2.BORDER_REPLICATE,
always_apply=True,
),
"crop":
trans.RandomResizedCrop(100,
100,
scale=(0.8, 0.95),
ratio=(0.8, 1.2),
always_apply=True),
# "elastic": trans.ElasticTransform(
# alpha=0.8,
# alpha_affine=10,
# sigma=40,
# border_mode=cv2.BORDER_REPLICATE,
# always_apply=True,
# ),
"distort":
trans.OpticalDistortion(0.2, always_apply=True),
},
"minor": {
"blur":
trans.GaussianBlur(7, always_apply=True),
"noise":
trans.GaussNoise((20.0, 40.0), always_apply=True),
"bright-contrast":
trans.RandomBrightnessContrast(0.4, 0.4, always_apply=True),
"hsv":
trans.HueSaturationValue(30, 40, 50, always_apply=True),
"rgb":
trans.RGBShift(always_apply=True),
"flip":
trans.HorizontalFlip(always_apply=True),
},
}
def augment(im, params=None):
"""
Perform data augmentation on some image using the albumentations package.
Parameters
----------
im : Numpy array
params : dict or None
Contains the data augmentation parameters
Mandatory keys:
- h_flip ([0,1] float): probability of performing an horizontal left-right mirroring.
- v_flip ([0,1] float): probability of performing an vertical up-down mirroring.
- rot ([0,1] float): probability of performing a rotation to the image.
- rot_lim (int): max degrees of rotation.
- stretch ([0,1] float): probability of randomly stretching an image.
- crop ([0,1] float): randomly take an image crop.
- zoom ([0,1] float): random zoom applied to crop_size.
--> Therefore the effective crop size at each iteration will be a
random number between 1 and crop*(1-zoom). For example:
* crop=1, zoom=0: no crop of the image
* crop=1, zoom=0.1: random crop of random size between 100% image and 90% of the image
* crop=0.9, zoom=0.1: random crop of random size between 90% image and 80% of the image
* crop=0.9, zoom=0: random crop of always 90% of the image
Image size refers to the size of the shortest side.
- blur ([0,1] float): probability of randomly blurring an image.
- pixel_noise ([0,1] float): probability of randomly adding pixel noise to an image.
- pixel_sat ([0,1] float): probability of randomly using HueSaturationValue in the image.
- cutout ([0,1] float): probability of using cutout in the image.
Returns
-------
Numpy array
"""
## 1) Crop the image
effective_zoom = np.random.rand() * params['zoom']
crop = params['crop'] - effective_zoom
ly, lx, channels = im.shape
crop_size = int(crop * min([ly, lx]))
rand_x = np.random.randint(low=0, high=lx - crop_size + 1)
rand_y = np.random.randint(low=0, high=ly - crop_size + 1)
crop = transforms.Crop(x_min=rand_x,
y_min=rand_y,
x_max=rand_x + crop_size,
y_max=rand_y + crop_size)
im = crop(image=im)['image']
## 2) Now add the transformations for augmenting the image pixels
transform_list = []
# Add random stretching
if params['stretch']:
transform_list.append(
imgaug_transforms.IAAPerspective(scale=0.1, p=params['stretch']))
# Add random rotation
if params['rot']:
transform_list.append(
transforms.Rotate(limit=params['rot_lim'], p=params['rot']))
# Add horizontal flip
if params['h_flip']:
transform_list.append(transforms.HorizontalFlip(p=params['h_flip']))
# Add vertical flip
if params['v_flip']:
transform_list.append(transforms.VerticalFlip(p=params['v_flip']))
# Add some blur to the image
if params['blur']:
transform_list.append(
albumentations.OneOf([
transforms.MotionBlur(blur_limit=7, p=1.),
transforms.MedianBlur(blur_limit=7, p=1.),
transforms.Blur(blur_limit=7, p=1.),
],
p=params['blur']))
# Add pixel noise
if params['pixel_noise']:
transform_list.append(
albumentations.OneOf(
[
transforms.CLAHE(clip_limit=2, p=1.),
imgaug_transforms.IAASharpen(p=1.),
imgaug_transforms.IAAEmboss(p=1.),
transforms.RandomBrightnessContrast(contrast_limit=0,
p=1.),
transforms.RandomBrightnessContrast(brightness_limit=0,
p=1.),
transforms.RGBShift(p=1.),
transforms.RandomGamma(p=1.) #,
# transforms.JpegCompression(),
# transforms.ChannelShuffle(),
# transforms.ToGray()
],
p=params['pixel_noise']))
# Add pixel saturation
if params['pixel_sat']:
transform_list.append(
transforms.HueSaturationValue(p=params['pixel_sat']))
# Remove randomly remove some regions from the image
if params['cutout']:
ly, lx, channels = im.shape
scale_low, scale_high = 0.05, 0.25 # min and max size of the squares wrt the full image
scale = np.random.uniform(scale_low, scale_high)
transform_list.append(
transforms.Cutout(num_holes=8,
max_h_size=int(scale * ly),
max_w_size=int(scale * lx),
p=params['cutout']))
# Compose all image transformations and augment the image
augmentation_fn = albumentations.Compose(transform_list)
im = augmentation_fn(image=im)['image']
return im
def augment(im, params=None):
"""
Perform data augmentation on some image using the albumentations package.
Parameters
----------
im : Numpy array
params : dict or None
Contains the data augmentation parameters
Mandatory keys:
- h_flip ([0,1] float): probability of performing an horizontal left-right mirroring.
- v_flip ([0,1] float): probability of performing an vertical up-down mirroring.
- rot ([0,1] float): probability of performing a rotation to the image.
- rot_lim (int): max degrees of rotation.
- stretch ([0,1] float): probability of randomly stretching an image.
- expand ([True, False] bool): whether to pad the image to a square shape with background color canvas.
- crop ([0,1] float): randomly take an image crop.
- invert_col ([0, 1] float): randomly invert the colors of the image. p=1 -> invert colors (VPR)
- zoom ([0,1] float): random zoom applied to crop_size.
--> Therefore the effective crop size at each iteration will be a
random number between 1 and crop*(1-zoom). For example:
* crop=1, zoom=0: no crop of the image
* crop=1, zoom=0.1: random crop of random size between 100% image and 90% of the image
* crop=0.9, zoom=0.1: random crop of random size between 90% image and 80% of the image
* crop=0.9, zoom=0: random crop of always 90% of the image
Image size refers to the size of the shortest side.
- blur ([0,1] float): probability of randomly blurring an image.
- pixel_noise ([0,1] float): probability of randomly adding pixel noise to an image.
- pixel_sat ([0,1] float): probability of randomly using HueSaturationValue in the image.
- cutout ([0,1] float): probability of using cutout in the image.
Returns
-------
Numpy array
"""
## 1) Expand the image by padding it with bg-color canvas
if params["expand"]:
desired_size = max(im.shape)
# check bg
if np.argmax(im.shape) > 0:
bgcol = tuple(np.repeat(int(np.mean(im[[0, -1], :, :])), 3))
else:
bgcol = tuple(np.repeat(int(np.mean(im[:, [0, -1], :])), 3))
im = Image.fromarray(im)
old_size = im.size # old_size[0] is in (width, height) format
ratio = float(desired_size) / max(old_size)
new_size = tuple([int(x * ratio) for x in old_size])
im = im.resize(new_size, Image.ANTIALIAS)
# create a new image and paste the resized on it
new_im = Image.new("RGB", (desired_size, desired_size), color=bgcol)
new_im.paste(im, ((desired_size - new_size[0]) // 2,
(desired_size - new_size[1]) // 2))
im = np.array(new_im)
## 2) Crop the image
if params["crop"] and params["crop"] != 1:
effective_zoom = np.random.rand() * params['zoom']
crop = params['crop'] - effective_zoom
ly, lx, channels = im.shape
crop_size = int(crop * min([ly, lx]))
rand_x = np.random.randint(low=0, high=lx - crop_size + 1)
rand_y = np.random.randint(low=0, high=ly - crop_size + 1)
crop = transforms.Crop(x_min=rand_x,
y_min=rand_y,
x_max=rand_x + crop_size,
y_max=rand_y + crop_size)
im = crop(image=im)['image']
if params["enhance"]:
im = Image.fromarray(im)
enhancer = ImageEnhance.Contrast(im)
im = np.array(enhancer.enhance(params["enhance"]))
## 3) Now add the transformations for augmenting the image pixels
transform_list = []
if params['invert_col']:
transform_list.append(transforms.InvertImg(p=params['invert_col']))
# Add random stretching
if params['stretch']:
transform_list.append(
imgaug_transforms.IAAPerspective(scale=0.1, p=params['stretch']))
# Add random rotation
if params['rot']:
transform_list.append(
transforms.Rotate(limit=params['rot_lim'], p=params['rot']))
# Add horizontal flip
if params['h_flip']:
transform_list.append(transforms.HorizontalFlip(p=params['h_flip']))
# Add vertical flip
if params['v_flip']:
transform_list.append(transforms.VerticalFlip(p=params['v_flip']))
# Add some blur to the image
if params['blur']:
transform_list.append(
albumentations.OneOf([
transforms.MotionBlur(blur_limit=7, p=1.),
transforms.MedianBlur(blur_limit=7, p=1.),
transforms.Blur(blur_limit=7, p=1.),
],
p=params['blur']))
# Add pixel noise
if params['pixel_noise']:
transform_list.append(
albumentations.OneOf(
[
transforms.CLAHE(clip_limit=2, p=1.),
imgaug_transforms.IAASharpen(p=1.),
imgaug_transforms.IAAEmboss(p=1.),
transforms.RandomBrightnessContrast(contrast_limit=0,
p=1.),
transforms.RandomBrightnessContrast(brightness_limit=0,
p=1.),
transforms.RGBShift(p=1.),
transforms.RandomGamma(p=1.) #,
# transforms.JpegCompression(),
# transforms.ChannelShuffle(),
# transforms.ToGray()
],
p=params['pixel_noise']))
# Add pixel saturation
if params['pixel_sat']:
transform_list.append(
transforms.HueSaturationValue(p=params['pixel_sat']))
# Remove randomly remove some regions from the image
if params['cutout']:
ly, lx, channels = im.shape
scale_low, scale_high = 0.05, 0.25 # min and max size of the squares wrt the full image
scale = np.random.uniform(scale_low, scale_high)
transform_list.append(
transforms.Cutout(num_holes=8,
max_h_size=int(scale * ly),
max_w_size=int(scale * lx),
p=params['cutout']))
# Compose all image transformations and augment the image
augmentation_fn = albumentations.Compose(transform_list)
im = augmentation_fn(image=im)['image']
return im
alb_transforms = [
alb.IAAAdditiveGaussianNoise(p=1),
alb.GaussNoise(p=1),
alb.MotionBlur(p=1),
alb.MedianBlur(blur_limit=3, p=1),
alb.Blur(blur_limit=3, p=1),
alb.OpticalDistortion(p=1),
alb.GridDistortion(p=1),
alb.IAAPiecewiseAffine(p=1),
aat.CLAHE(clip_limit=2, p=1),
alb.IAASharpen(p=1),
alb.IAAEmboss(p=1),
aat.HueSaturationValue(p=0.3),
aat.HorizontalFlip(p=1),
aat.RGBShift(),
aat.RandomBrightnessContrast(),
aat.RandomGamma(p=1),
aat.Cutout(2, 10, 10, p=1),
aat.Equalize(mode='cv', p=1),
aat.FancyPCA(p=1),
aat.RandomFog(p=1),
aat.RandomRain(blur_value=3, p=1),
albumentations.IAAAffine(p=1),
albumentations.ShiftScaleRotate(rotate_limit=15, p=1)
]
def one_by_one():
data_dir = 'train_val/pic'
data_anno = pd.read_csv('train_val/keys.csv')