def strong_aug(p=0.9):
return Compose(
[
OneOf([
IAAAdditiveGaussianNoise(scale=(0.01 * 255, 0.05 * 255),
p=1.0),
GaussNoise(var_limit=(20, 120), p=1.0),
RandomGamma(gamma_limit=(80, 120), p=1.0),
],
p=0.9),
RandomBrightnessContrast(p=1.0),
OneOf(
[
# MotionBlur(p=1.0),
# MedianBlur(blur_limit=3, p=1.0),
Blur(blur_limit=5, p=1.0),
IAASharpen(p=1.0),
# IAAEmboss(p=1.0),
# IAASuperpixels(n_segments=10, p_replace=0.05, p=1.0),
],
p=0.9),
OneOf(
[
CLAHE(clip_limit=8, p=1.0),
RGBShift(p=1.0),
ChannelShuffle(p=1.0),
HueSaturationValue(p=1.0),
# ToGray(p=1.0),
],
p=0.9),
# OneOf([
# OpticalDistortion(border_mode=cv2.BORDER_CONSTANT, p=1.0),
# # GridDistortion(border_mode=cv2.BORDER_CONSTANT, p=1.0),
# IAAPiecewiseAffine(nb_rows=4, nb_cols=4, p=1.0),
# IAAPerspective(scale=(0.05, 0.075), p=1.0),
# # IAAAffine(mode='constant', p=1.0),
# ElasticTransform(alpha=alpha, sigma=sigma, alpha_affine=alpha_affine,
# border_mode=cv2.BORDER_CONSTANT,
# p=1.0),
# ], p=0.9),
],
p=p)
def train_transform(p=1.0):
augmentation = Compose([
FlipChannels(),
VerticalFlip(p=p),
HorizontalFlip(p=p),
RandomRotate90(p=p),
RandomGamma(p=p, gamma_limit=(90, 350)),
OpticalDistortion(p=p, border_mode=cv2.BORDER_CONSTANT),
GridDistortion(p=p, border_mode=cv2.BORDER_CONSTANT),
ShiftScaleRotate(p=p, scale_limit=0.2, border_mode=cv2.BORDER_CONSTANT)
],
p=p)
# @contrast_norm
def transform_fun(img):
data = {'image': img}
augmented = augmentation(**data)
return augmented['image']
return transform_fun
def get_augmentations_train():
AUGMENTATIONS_TRAIN = Compose([
HorizontalFlip(p=0.5),
OneOf([
RandomContrast(),
RandomGamma(),
RandomBrightness(),
], p=0.3),
OneOf([
ElasticTransform(
alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03),
GridDistortion(),
OpticalDistortion(distort_limit=2, shift_limit=0.5),
],
p=0.3),
RandomSizedCrop(min_max_height=(176, 256), height=h, width=w, p=0.25),
ToFloat(max_value=1)
],
p=1)
return AUGMENTATIONS_TRAIN
def create_train_transforms(conf):
height = conf['crop_height']
width = conf['crop_width']
return Compose([
RandomSizedCropAroundBbox(min_max_height=(int(
height * 0.8), int(height * 1.2)),
w2h_ratio=1.,
height=height,
width=width,
p=1),
HorizontalFlip(),
VerticalFlip(),
RandomRotate90(),
Transpose(),
Lighting(alphastd=0.3),
RandomBrightnessContrast(p=0.2),
RandomGamma(p=0.2),
RGBShift(p=0.2)
],
additional_targets={'image1': 'image'})
def aug_baseline_CLAHE_Sharpen(image, mask):
aug = Compose([
HorizontalFlip(),
VerticalFlip(),
OneOf([RandomContrast(),
RandomGamma(),
RandomBrightness()], p=0.3),
OneOf([
ElasticTransform(
alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03),
GridDistortion(),
OpticalDistortion(distort_limit=2, shift_limit=0.5)
],
p=0.3),
ShiftScaleRotate(),
CLAHE(),
IAASharpen()
])
aug_img = aug(image=image, mask=mask)
return aug_img['image'], aug_img['mask']
def augment_flips_color(p=.5):
return Compose([
CLAHE(),
RandomRotate90(),
RandomGamma(),
RGBShift(),
Blur(blur_limit=3),
HueSaturationValue(),
OneOf([
MotionBlur(p=.2),
MedianBlur(blur_limit=3, p=0.1),
Blur(blur_limit=3, p=0.1),
],
p=0.2),
OneOf([
IAAAdditiveGaussianNoise(),
GaussNoise(),
], p=0.2),
],
p=p)
def get_augmentations(augmentation, p, image_size = 256):
if augmentation == 'train':
augmentations = Compose([
# RandomScale(scale_limit=0.125),
HorizontalFlip(p=0.5),
RandomBrightnessContrast(p=0.5),
RandomGamma(p=0.3),
ShiftScaleRotate(shift_limit=0.1625, scale_limit=0.6, rotate_limit=15, p=0.6),
# ShiftScaleRotate(rotate_limit=20, p=0.6),
Resize(image_size,image_size)
], p=p)
elif augmentation == 'valid':
augmentations = Compose([
Resize(image_size, image_size)
], p=p)
else:
raise ValueError("Unknown Augmentations")
return augmentations
def aug_with_crop(image_size = 256, crop_prob = 1):
# Monkey-patch lol
albu.augmentations.functional.MAX_VALUES_BY_DTYPE[np.dtype('float64')] = 1.0
return albu.Compose([
RandomCrop(width = image_size, height = image_size, p=crop_prob),
HorizontalFlip(p=0.3),
VerticalFlip(p=0.3),
RandomRotate90(p=0.3),
Transpose(p=0.3),
ShiftScaleRotate(shift_limit=0.01, scale_limit=0.04, rotate_limit=0, p=0.1),
RandomBrightnessContrast(p=0.3),
RandomGamma(p=0.1),
IAAEmboss(p=0.1),
Blur(p=0.001, blur_limit = 3),
OneOf([
ElasticTransform(p=0.1, alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03),
GridDistortion(p=0.1),
OpticalDistortion(p=0.05, distort_limit=2, shift_limit=0.5)
], p=0.5)
], p = 1)
def __getitem__(self, i):
idx = self.ids[i]
mask_file = glob(self.masks_dir + idx + '*')
img_file = glob(self.imgs_dir + idx + '*')
assert len(mask_file) == 1, \
f'Either no mask or multiple masks found for the ID {idx}: {mask_file}'
assert len(img_file) == 1, \
f'Either no image or multiple images found for the ID {idx}: {img_file}'
mask = Image.open(mask_file[0])
img = Image.open(img_file[0])
# random image augmentation
aug = Compose([
OneOf([
ElasticTransform(p=0.5,
alpha=120,
sigma=120 * 0.05,
alpha_affine=120 * 0.03),
GridDistortion(p=0.5),
RandomGamma(gamma_limit=(50, 130), p=0.5),
CLAHE(clip_limit=2.0, p=0.5),
RandomBrightnessContrast(brightness_limit=0.4, p=0.5),
Rotate(limit=20),
RandomScale(scale_limit=0.2)
],
p=0.8)
])
augmented = aug(image=img, mask=mask)
img = augmented['image']
mask = augmented['mask']
assert img.size == mask.size, \
f'Image and mask {idx} should be the same size, but are {img.size} and {mask.size}'
img = self.preprocess(img)
# mask = np.array(mask)
# mask[mask > 1] = 0
mask = self.mask_preprocess(mask)
return {'image': torch.from_numpy(img), 'mask': torch.from_numpy(mask)}
def crazy_custom_aug(targets):
cage_aug = RandomCageOverlay(default_cage_maker, p=0.15)
info_aug = RandomInfoOverlay(image_dict, max_overlay_num=10, p=0.5)
aug = targets_aug(
[
cage_aug,
info_aug,
ShiftScaleRotate(),
RandomCrop(256, 256),
ElasticTransform(p=0.3),
MedianBlur(p=0.3),
RandomBrightness(p=0.3),
HueSaturationValue(p=0.5),
HorizontalFlip(p=0.5),
VerticalFlip(p=0.5),
RandomBrightnessContrast(brightness_limit=0.5, contrast_limit=0.5),
RandomGamma(p=1),
],
targets,
)
return aug
def aug_baseline_randCrop(image, mask):
aug = Compose([
HorizontalFlip(),
VerticalFlip(),
OneOf([RandomContrast(),
RandomGamma(),
RandomBrightness()], p=0.3),
OneOf([
ElasticTransform(
alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03),
GridDistortion(),
OpticalDistortion(distort_limit=2, shift_limit=0.5)
],
p=0.3),
RandomResizedCrop(width=cfg.TRAIN_CROP_SIZE[0],
height=cfg.TRAIN_CROP_SIZE[1],
scale=[0.5, 1.0]),
RandomRotate90()
])
aug_img = aug(image=image, mask=mask)
return aug_img['image'], aug_img['mask']
def get_transforms(stage: str = None, mode: str = None, size: int = 224):
if mode == 'train':
transforms = [
Resize(size, size),
OneOf([
RandomBrightnessContrast(brightness_limit=0.5,
contrast_limit=0.5),
RandomGamma(),
CLAHE(),
],
p=0.5),
ShiftScaleRotate(rotate_limit=45),
OneOf([Blur(), MotionBlur()], p=0.5),
JpegCompression(quality_lower=10),
Normalize(),
ToTensor()
]
else:
transforms = [Resize(size, size), Normalize(), ToTensor()]
return Compose(transforms=transforms)
def train_transform_regr(sz, downscale=1, p=1):
augmentation = Compose([
VerticalFlip(p=0.5),
HorizontalFlip(p=0.5),
RandomRotate90(p=0.5),
RandomGamma(p=0.9, gamma_limit=(80, 150)),
HueSaturationValue(
p=0.9, hue_shift_limit=10, sat_shift_limit=20, val_shift_limit=10),
],
p=p)
def transform_fun(img_path, n=None):
img = read_img(img_path, sz, downscale=downscale)
data = {"image": img}
augmented = augmentation(**data)
img = augmented["image"]
img = norm_fun(img)
return img
return transform_fun
def augment_flips_color(p=.5):
return Compose(
[
CLAHE(),
RandomRotate90(),
ShiftScaleRotate(border_mode=cv2.BORDER_CONSTANT, value=0),
RandomGamma(gamma_limit=(75, 140)),
Blur(blur_limit=3),
# HueSaturationValue(),
OneOf([
MotionBlur(p=.2),
MedianBlur(blur_limit=3, p=0.1),
Blur(blur_limit=3, p=0.1),
],
p=0.2),
OneOf([
IAAAdditiveGaussianNoise(),
GaussNoise(),
], p=0.2),
],
p=p)
def hard_transform():
transforms = [
ShiftScaleRotate(shift_limit=0.1,
scale_limit=0.1,
rotate_limit=15,
border_mode=cv2.BORDER_REFLECT,
p=0.5),
IAAPerspective(scale=(0.02, 0.05), p=0.3),
RandomBrightnessContrast(brightness_limit=0.2,
contrast_limit=0.2,
p=0.3),
RandomGamma(gamma_limit=(85, 115), p=0.3),
HueSaturationValue(p=0.3),
ChannelShuffle(p=0.5),
ToGray(p=0.2),
CLAHE(p=0.3),
RGBShift(p=0.3),
JpegCompression(quality_lower=50),
]
transforms = Compose(transforms)
return transforms
def load_batch(chip_coords, all_density, all_gsd, image, gsd_train_mean,
gsd_train_std, args):
images = np.zeros((len(chip_coords), args.model_image_size[0],
args.model_image_size[1], 3))
densities = np.zeros((len(chip_coords), 1))
gsds = np.zeros((len(chip_coords), 2))
for i in range(len(chip_coords)):
x1, y1, x2, y2 = chip_coords[i]
density = all_density[i]
gsd = all_gsd[i]
sub_image = image[y1:y2, x1:x2, :3]
sub_image = cv2.resize(sub_image, args.model_image_size)
aug = Compose([
VerticalFlip(p=0.5),
RandomRotate90(p=0.5),
HorizontalFlip(p=0.5),
Transpose(p=0.5),
CLAHE(p=0.2),
RandomBrightness(limit=0.2, p=0.2),
RandomGamma(p=0.2),
HueSaturationValue(hue_shift_limit=20,
sat_shift_limit=10,
val_shift_limit=10,
p=0.2),
RandomContrast(limit=0.4, p=0.3)
])
sub_image = aug(image=sub_image)['image']
images[i, :, :, :] = sub_image
densities[i, 0] = density
gsds[i, 0] = (abs(gsd[0]) - gsd_train_mean[0]) / gsd_train_std[0]
gsds[i, 1] = (abs(gsd[1]) - gsd_train_mean[1]) / gsd_train_std[1]
images = imagenet_utils.preprocess_input(images) / 255.0
return images, densities, gsds
def data_augmentation(original_image, original_mask):
"""进行样本和掩膜的随机增强
Args:
original_image: 原始图片
original_mask: 原始掩膜
Return:
image_aug: 增强后的图片
mask_aug: 增强后的掩膜
"""
original_height, original_width = original_image.shape[:2]
augmentations = Compose([
HorizontalFlip(p=0.4),
Rotate(limit=15, p=0.4),
CenterCrop(p=0.3, height=original_height, width=original_width),
# 直方图均衡化
CLAHE(p=0.4),
# 亮度、对比度
RandomGamma(gamma_limit=(80, 120), p=0.1),
RandomBrightnessContrast(p=0.1),
# 模糊
OneOf([
MotionBlur(p=0.1),
MedianBlur(blur_limit=3, p=0.1),
Blur(blur_limit=3, p=0.1),
],
p=0.3),
OneOf([
IAAAdditiveGaussianNoise(),
GaussNoise(),
], p=0.2)
])
augmented = augmentations(image=original_image, mask=original_mask)
image_aug = augmented['image']
mask_aug = augmented['mask']
return image_aug, mask_aug
def initialize(self, config, filename):
self.config = config
self.train = 'train' in filename
self.A_paths = self.get_files(filename)
# self.boxes = self.get_boxes()
self.A_size = self.A_paths.shape[0]
self.batch_size = config['batch_size']
self.size = config['image_size']
self.transform = Compose([
# Resize(self.size[0], self.size[1]),
ShiftScaleRotate(shift_limit=0.0,
scale_limit=(0.0, 0.1),
rotate_limit=10,
p=.3),
OneOf([
IAAAdditiveGaussianNoise(),
GaussNoise(),
], p=0.4),
OneOf([
CLAHE(clip_limit=2),
IAASharpen(),
IAAEmboss(),
RandomContrast(),
RandomBrightness(),
RandomGamma()
],
p=0.6),
OneOf([
RGBShift(),
HueSaturationValue(),
], p=0.6),
JpegCompression(quality_lower=40, quality_upper=100, p=0.3),
Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
),
])
self.used_classes = deque()
def get_transforms(phase):
original_height = 1400
original_width = 2100
list_transforms = []
if phase == "train":
list_transforms.extend([
OneOf([
RandomSizedCrop(min_max_height=(50, 101),
height=original_height,
width=original_width,
p=0.5),
PadIfNeeded(min_height=original_height,
min_width=original_width,
p=0.5)
],
p=1),
VerticalFlip(p=0.5),
# RandomRotate90(p=0.5),
OneOf([
ElasticTransform(p=0.5,
alpha=120,
sigma=120 * 0.05,
alpha_affine=120 * 0.03),
GridDistortion(p=0.5),
OpticalDistortion(p=1, distort_limit=2, shift_limit=0.5)
],
p=0.8),
CLAHE(p=0.8),
RandomBrightnessContrast(p=0.8),
RandomGamma(p=0.8),
])
list_transforms.extend([
Resize(height=original_height,
width=original_width,
interpolation=cv2.INTER_NEAREST),
Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), p=1),
ToTensor(),
])
list_trfms = Compose(list_transforms)
return list_trfms
def augmentation(image, mask, noise=False, transform=False, clahe=True, r_bright=True, r_gamma=True):
aug_list = [
VerticalFlip(p=0.5),
HorizontalFlip(p=0.5),
RandomRotate90(p=0.5),
]
if r_bright:
aug_list += [RandomBrightnessContrast(p=.5)]
if r_gamma:
aug_list += [RandomGamma(p=.5)]
if clahe:
aug_list += [CLAHE(p=1., always_apply=True)]
if noise:
aug_list += [GaussNoise(p=.5, var_limit=1.)]
if transform:
aug_list += [ElasticTransform(p=.5, sigma=1., alpha_affine=20, border_mode=0)]
aug = Compose(aug_list)
augmented = aug(image=image, mask=mask)
image_heavy = augmented['image']
mask_heavy = augmented['mask']
return image_heavy, mask_heavy
def simple_aug(p=0.5):
return Compose(
[
#HorizontalFlip(p=0.5),
ShiftScaleRotate(rotate_limit=10,
scale_limit=0.15,
p=0.5,
border_mode=cv2.BORDER_CONSTANT,
value=[0, 0, 0]),
OneOf([
JpegCompression(quality_lower=80),
Blur(),
], p=0.5),
OneOf([
CLAHE(),
RandomGamma(),
RandomContrast(),
RandomBrightness(),
],
p=0.5)
],
p=p)
def __call__(self, x):
return Compose([
HorizontalFlip(),
VerticalFlip(),
Transpose(),
RandomRotate90(),
ShiftScaleRotate(
shift_limit=0.0625, scale_limit=0.2, rotate_limit=45, p=0.2),
OneOf([
MotionBlur(p=.2),
MedianBlur(blur_limit=3, p=0.1),
Blur(blur_limit=3, p=0.1),
],
p=0.2),
OneOf([
RandomBrightnessContrast(),
RandomGamma(),
CLAHE(),
HueSaturationValue(p=0.3),
])
],
p=self.p)(image=x)['image']
def __init__(self, root_path, file_list, aug=False, is_test=False):
self.is_test = is_test
self.root_path = root_path
self.file_list = file_list
if aug:
self._aug = Compose([
HorizontalFlip(),
# OneOf([
# ShiftScaleRotate(shift_limit=0.2, scale_limit=0.2, rotate_limit=15,
# interpolation=cv2.INTER_NEAREST),
# RandomSizedCrop(min_max_height=(70, 90),
# height=IMAGE_PADDED, width=IMAGE_PADDED,
# interpolation=cv2.INTER_NEAREST)]
# ], p=1),
# ElasticTransform(p=0.5, border_mode=cv2.BORDER_REPLICATE,
# alpha=1, sigma=5, alpha_affine=5, interpolation=cv2.INTER_NEAREST),
# OneOf([GridDistortion(distort_limit=0.1, border_mode=cv2.BORDER_REPLICATE,
# interpolation=cv2.INTER_NEAREST),
# OpticalDistortion(border_mode=cv2.BORDER_REPLICATE, interpolation=cv2.INTER_NEAREST, p=0)], p=1),
# InvertImg(p=0.5),
RandomBrightness(p=0.5, limit=0.5),
RandomContrast(p=0.5, limit=0.5),
RandomGamma(p=0.5, gamma_limit=(70, 130)),
CLAHE(p=0.5),
# OneOf([
# Blur(p=1, blur_limit=4),
# MedianBlur(p=1, blur_limit=4),
# GaussNoise(p=1)
# ], p=1),
Normalize(),
# Resize(IMAGE_PADDED, IMAGE_PADDED),
PadIfNeeded(min_height=IMAGE_TOTAL_SIZE, min_width=IMAGE_TOTAL_SIZE, border_mode=cv2.BORDER_REPLICATE)
])
else:
self._aug = Compose([
Normalize(),
# Resize(IMAGE_PADDED, IMAGE_PADDED),
PadIfNeeded(min_height=IMAGE_TOTAL_SIZE, min_width=IMAGE_TOTAL_SIZE, border_mode=cv2.BORDER_REPLICATE)
])
def augment_hard(image):
aug = OneOf([
Blur(blur_limit=5, p=1.),
RandomGamma(gamma_limit=(50, 150), p=1.),
HueSaturationValue(hue_shift_limit=20,
sat_shift_limit=30,
val_shift_limit=20,
p=1.),
RGBShift(r_shift_limit=15, g_shift_limit=5, b_shift_limit=15,
p=1.),
RandomBrightness(limit=.25, p=1.),
RandomContrast(limit=.25, p=1.),
MedianBlur(blur_limit=5, p=1.),
CLAHE(clip_limit=2.0, tile_grid_size=(8, 8), p=1.)
],
p=1.)
augmented = aug(image=image)
image_augmented = augmented['image']
return image_augmented
def aug_image(image):
from albumentations import (
HorizontalFlip, IAAPerspective, ShiftScaleRotate, CLAHE,
RandomRotate90, RandomGamma, VerticalFlip, Transpose, ShiftScaleRotate,
Blur, OpticalDistortion, GridDistortion, HueSaturationValue,
IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur,
RandomBrightnessContrast, IAAPiecewiseAffine, IAASharpen, IAAEmboss,
Flip, OneOf, Compose, Rotate, RandomContrast, RandomBrightness,
RandomCrop, Resize, OpticalDistortion)
transforms = Compose(
[
#Rotate(limit=30, p=0.5),
#Rotate(limit=180, p=0.5),
#RandomRotate90(p=1.0)
#Transpose(p=1.0)
# Resize(248,248, p=1), # resize 후 크롭
# RandomCrop(224,224, p=1), # 위에꺼랑 세트
OneOf(
[
RandomContrast(
p=1, limit=(-0.5, 2)
), # -0.5 ~ 2 까지가 현장과 가장 비슷함 -- RandomBrightnessContrast
RandomBrightness(p=1, limit=(-0.2, 0.4)),
RandomGamma(p=1, gamma_limit=(80, 200)),
],
p=1),
# OneOf([
# Rotate(limit=30, p=0.3),
# RandomRotate90(p=0.3),
# VerticalFlip(p=0.3)
# ], p=0.3),
MotionBlur(p=0.2), # 움직일때 흔들리는 것 같은 이미지
#ShiftScaleRotate(shift_limit=0.001, scale_limit=0.1, rotate_limit=30, p=0.3, border_mode=1),
#Resize(224,224, p=1),
],
p=1)
return transforms(image=image)['image']
def get_transforms(phase_config):
list_transforms = []
if phase_config.Noise:
list_transforms.append(
OneOf([
GaussNoise(),
IAAAdditiveGaussianNoise(),
], p=0.5), )
if phase_config.Contrast:
list_transforms.append(
OneOf([
RandomContrast(0.5),
RandomGamma(),
RandomBrightness(),
],
p=0.5), )
if phase_config.Blur:
list_transforms.append(
OneOf([
MotionBlur(p=.2),
MedianBlur(blur_limit=3, p=0.1),
Blur(blur_limit=3, p=0.1),
],
p=0.5))
if phase_config.Distort:
list_transforms.append(
OneOf([
OpticalDistortion(p=0.3),
GridDistortion(p=.1),
IAAPiecewiseAffine(p=0.3),
],
p=0.5))
list_transforms.extend([
Normalize(mean=phase_config.mean, std=phase_config.std, p=1),
ToTensor(),
])
return Compose(list_transforms)
def data_augmentation(original_image,
original_mask,
crop=False,
height=None,
width=None):
augmentations = Compose([
HorizontalFlip(p=0.4),
VerticalFlip(p=0.4),
ShiftScaleRotate(shift_limit=0.07, rotate_limit=0, p=0.4),
CLAHE(p=0.3),
RandomGamma(gamma_limit=(80, 120), p=0.1),
RandomBrightnessContrast(p=0.1),
OneOf([
MotionBlur(p=0.1),
MedianBlur(blur_limit=3, p=0.1),
Blur(blur_limit=3, p=0.1),
],
p=0.3),
OneOf([
IAAAdditiveGaussianNoise(),
GaussNoise(),
], p=0.2)
])
if crop:
assert height and width
crop_aug = RandomCrop(height=height, width=width, always_apply=True)
crop_sample = crop_aug(image=original_image, mask=original_mask)
original_image = crop_sample['image']
original_mask = crop_sample['mask']
augmented = augmentations(image=original_image, mask=original_mask)
image_aug = augmented['image']
mask_aug = augmented['mask']
return image_aug, mask_aug
def __init__(self, dataset_dir, cfg, index_list=None, preprocessing=None):
super(TrainSetLoader, self).__init__()
self.preprocessing = preprocessing
self.dataset_dir = os.path.join(dataset_dir)
self.all_file_list = [
os.path.splitext(file)[0] for file in sorted(
glob(os.path.join(self.dataset_dir, 'imgs', '*.png')))
]
if index_list is None:
self.file_list = self.all_file_list
else:
self.file_list = []
for index in index_list:
self.file_list.append(self.all_file_list[index])
assert len(self.all_file_list) // 5 == len(
self.file_list) // 4, 'the length of train file_list is error'
self.aug = Compose(
[
RandomSizedCrop(p=1,
min_max_height=(int(
cfg.size * 1), int(cfg.size * 1)),
height=cfg.size,
width=cfg.size),
RandomBrightnessContrast(p=0.8),
RandomGamma(p=0.8),
HorizontalFlip(p=0.5),
VerticalFlip(p=0.5),
RandomRotate90(p=0.5),
Transpose(p=0.5),
# OneOf([ElasticTransform(p=1, alpha=120, sigma=120*0.05, alpha_affine=120*0.03),
# GridDistortion(p=1),
# OpticalDistortion(p=1, distort_limit=2, shift_limit=0.5),
# ], p=0.8)
# ElasticTransform(p=0.9, alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03),
# GridDistortion(p=0.8),
# OpticalDistortion(p=0.8, distort_limit=0.5, shift_limit=0.5),
],
p=1)
def __init__(self, image_ids, transform=True, preprocessing_fn=None):
"""
Dataset class for segmentation problem
:param image_ids: ids of the images, list
:param transform: True/False, no transform in validation
:param preprocessing_fn: a function for preprocessing image
"""
# Create empty dictionary to store image and mask paths
self.data = defaultdict(dict)
# for augmentations
self.transform = transform
# preprocessing function to normalize images
self.preprocessing_fn = preprocessing_fn
# albumentations augmentations
self.aug = Compose([
ShiftScaleRotate(shift_limit=0.0625,
scale_limit=0.1,
rotate_limit=10,
p=0.8),
OneOf(
[
RandomGamma(gamma_limit=(90, 110)),
RandomBrightnessContrast(brightness_limit=0.1,
contrast_limit=0.1),
],
p=0.5,
),
])
# going over all image_ids to store image and mask paths
counter = 0
for imgid in image_ids:
files = glob.glob(os.path.join(TRAIN_PATH, imgid, "*.png"))
self.data[counter] = {
"img_path": os.path.join(TRAIN_PATH, imgid + ".png"),
"mask_path": os.path.join(MASK_PATH, imgid + ".png")
}
counter += 1
def get_augmentations(params):
augmentations = {
"light": Compose([
HorizontalFlip(p = 0.5),
Rotate(p = 0.5)])
,
"medium": Compose([
OneOf([RandomSizedCrop(min_max_height = (params.img_shape - 40, params.img_shape - 40),
height = params.img_shape,
width = params.img_shape,
p = 0.1),
PadIfNeeded(min_height = params.img_shape, min_width = params.img_shape, p = 1)], p = 1),
VerticalFlip(p = 0.2),
RandomRotate90(p = 0.2),
Rotate(p = 0.5),
RandomBrightnessContrast(brightness_limit = 0.6, contrast_limit = 0.6, p = 0.1),
GaussNoise(p = 0.1, var_limit = (10.0, 25.0)),
ISONoise(color_shift = (0.01, 0.5), intensity = (0.1, 0.9), p = 0.1),
RandomGamma(gamma_limit = (50, 150), p = 0.1)],
)}
return augmentations