def img_augment(p=1.): return Compose([ RandomSizedCrop((280, 345), 350, 525, p=0.9, w2h_ratio=1.5), HorizontalFlip(.5), VerticalFlip(.5), OneOf([ CLAHE(clip_limit=2), IAASharpen(), IAAEmboss(), RandomContrast(), RandomBrightness(), ], p=0.3), # ShiftScaleRotate(shift_limit=0.15, scale_limit=0.15, rotate_limit=20, p=.75, border_mode=cv2.BORDER_REFLECT), Blur(blur_limit=3, p=.33), GaussNoise(p=0.8), OpticalDistortion(p=.33), GridDistortion(p=.33), HueSaturationValue(p=.33) ], 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 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 strong_aug(p=.5): return Compose([ RandomRotate90(p=0.2), Transpose(p=0.2), HorizontalFlip(p=0.5), VerticalFlip(p=0.5), OneOf([ GaussNoise(), ], p=0.2), ShiftScaleRotate(p=0.2), OneOf([ CLAHE(clip_limit=2), Sharpen(), Emboss(), RandomBrightnessContrast(), ], p=0.2), HueSaturationValue(p=0.2), ], p=p)
def generate_transforms(image_size): train_transform = Compose([ Resize(height=image_size[0], width=image_size[1]), OneOf( [RandomBrightness(limit=0.1, p=1), RandomContrast(limit=0.1, p=1)]), OneOf([ MotionBlur(blur_limit=3), MedianBlur(blur_limit=3), GaussianBlur(blur_limit=3) ], p=0.5), VerticalFlip(p=0.5), HorizontalFlip(p=0.5), ShiftScaleRotate( shift_limit=0.2, scale_limit=0.2, rotate_limit=20, interpolation=cv2.INTER_LINEAR, border_mode=cv2.BORDER_REFLECT_101, p=1, ), Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0, p=1.0), ]) val_transform = Compose([ Resize(height=image_size[0], width=image_size[1]), Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0, p=1.0), ]) return { "train_transforms": train_transform, "val_transforms": val_transform }
def main(): # 输入 train_path = base_train_path + 'image/*.tif' mask_path = base_train_path + 'label/*.png' # 增强结果输出目录 augtrain_path = base_train_path + 'image_aug/' augmask_path = base_train_path + 'label_aug/' train_img, masks = data_num(train_path, mask_path) for data in range(len(train_img)): file_name = train_img[data].split('\\')[1].split('.')[0] image = cv2.imread(train_img[data]) mask = np.array(Image.open(masks[data])) # 水平翻转 augmented_1 = HorizontalFlip(p=1)(image=image, mask=mask) aug_image_1 = augmented_1['image'] aug_mask_1 = Image.fromarray(augmented_1['mask']) cv2.imwrite(augtrain_path + "/{}_{}.tif".format(file_name, 1), aug_image_1) aug_mask_1.save(augmask_path + "/{}_{}.png".format(file_name, 1)) # 垂直翻转 augmented_2 = VerticalFlip(p=1)(image=image, mask=mask) aug_image_2 = augmented_2['image'] aug_mask_2 = Image.fromarray(augmented_2['mask']) cv2.imwrite(augtrain_path + "/{}_{}.tif".format(file_name, 2), aug_image_2) aug_mask_2.save(augmask_path + "/{}_{}.png".format(file_name, 2)) # 水平 + 垂直 翻转 augmented_3 = Transpose(p=1)(image=image, mask=mask) aug_image_3 = augmented_3['image'] aug_mask_3 = Image.fromarray(augmented_3['mask']) cv2.imwrite(augtrain_path + "/{}_{}.tif".format(file_name, 3), aug_image_3) aug_mask_3.save(augmask_path + "/{}_{}.png".format(file_name, 3)) if data % 1000 == 0: print(data)
def produce_dataset(data_dirs, save_dir, times=1): image_dir = data_dirs + '\image_crop' label_dir = data_dirs + '\label_crop' aug = Compose([ PadIfNeeded(min_height=img_width, min_width=img_height, 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=0.5, distort_limit=1, shift_limit=0.5) ], p=0.8) ]) i = 0 for epoch in range(times): for file in os.listdir(image_dir): ori_image = np.load(os.path.join(image_dir, file)) ori_label = np.load(os.path.join(label_dir, file)) assert ori_image.dtype == 'uint8' assert ori_label.dtype == 'int' augmented = aug(image=ori_image, mask=ori_label) image = augmented['image'] label = augmented['mask'] if label.max() == 0: a = np.random.rand() if a < 0.5: continue np.save(os.path.join(save_dir + '\image', str(i) + '.npy'), image) np.save(os.path.join(save_dir + '\label', str(i) + '.npy'), label) i += 1 print('epoch = ' + str(epoch)) print(' Load over!\n i = ' + str(i))
def train(cls_model='b2', shape=(320,320)): kfold = StratifiedKFold(n_splits=4, random_state=133, shuffle=True) train_df, img_2_vector = preprocess() albumentations_train = Compose([ VerticalFlip(), HorizontalFlip(), Rotate(limit=20), GridDistortion() ], p=1) for n_fold, (train_indices, val_indices) in enumerate(kfold.split(train_df['Image'].values, train_df['Class'].map(lambda x: str(sorted(list(x)))))): train_imgs = train_df['Image'].values[train_indices] val_imgs = train_df['Image'].values[val_indices] data_generator_train = DataGenenerator(train_imgs, augmentation=albumentations_train, resized_height=shape[0], resized_width=shape[1], img_2_ohe_vector=img_2_vector) data_generator_train_eval = DataGenenerator(train_imgs, shuffle=False, resized_height=shape[0], resized_width=shape[1], img_2_ohe_vector=img_2_vector) data_generator_val = DataGenenerator(val_imgs, shuffle=False, resized_height=shape[0], resized_width=shape[1], img_2_ohe_vector=img_2_vector) model = get_model(cls_model, shape=shape) model.compile(optimizer=RAdam(), loss='binary_crossentropy', metrics=['accuracy']) train_metric_callback = PrAucCallback(data_generator_train_eval) checkpoint_name = cls_model + '_' + str(n_fold) val_callback = PrAucCallback(data_generator_val, stage='val', checkpoint_name=checkpoint_name) history_0 = model.fit_generator(generator=data_generator_train, validation_data=data_generator_val, epochs=20, callbacks=[train_metric_callback, val_callback], workers=42, verbose=1 )
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 get_transforms(phase, mean, std): list_transforms = [] if phase == "train": list_transforms.extend( [ Compose([ HorizontalFlip(p=0.5), # only horizontal flip as of now VerticalFlip(p=0.5)], p=1.0), ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.1, rotate_limit=20, p=0.5), # GaussNoise() # ] ) list_transforms.extend( [ Normalize(mean=mean, std=std, p=1), ToTensor(), ] ) list_trfms = Compose(list_transforms) return list_trfms
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_train_transforms(config): return Compose([ RandomResizedCrop(config.image_size, config.image_size), HorizontalFlip(p=0.4), VerticalFlip(p=0.4), Transpose(p=0.3), ShiftScaleRotate(p=0.3), HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, p=0.3), RandomBrightnessContrast( brightness_limit=(-0.1, 0.1), contrast_limit=(-0.1, 0.1), p=0.3), Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0), Cutout(p=0.3), ToTensorV2(p=1.0), ], p=1.)
def get_train_dataloader(file_list, opt): data_transform = Compose([ PadIfNeeded(min_height=opt.train_crop_height, min_width=opt.train_crop_width, p=1), RandomCrop( height=opt.train_crop_height, width=opt.train_crop_width, p=1), VerticalFlip(p=0.5), HorizontalFlip(p=0.5), Normalize(p=1) ], p=1) train_dataset = RoboticsDataset(file_names=file_list, transform=data_transform, problem_type=opt.problem_type) train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.workers, pin_memory=True) return train_dataloader
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 generate_transforms_resnext(img_size): train_transform = Compose([ RandomResizedCrop(img_size, img_size), Transpose(p=0.5), HorizontalFlip(p=0.5), VerticalFlip(p=0.5), ShiftScaleRotate(p=0.5), Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], ), ToTensorV2(), ]) val_transform = Compose([ Resize(height=img_size, width=img_size), Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ToTensorV2(), ]) return {"train": train_transform, "val": val_transform}
def argument(x_train, msk_train, y_train): aug_img = [] aug_msk = [] aug_y = [] augments = [ (1, HorizontalFlip(p=1)), (0.25, VerticalFlip(p=1)), (0.25, RandomRotate90(p=1)), (0.25, Transpose(p=1)), (0.25, ElasticTransform(p=1, alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03)), (0.25, GridDistortion(p=1)), (0.25, OpticalDistortion(p=1, distort_limit=2, shift_limit=0.5)), # (0.5, RandomSizedCrop(p=1, min_max_height=(int(img_size_ori / 2), img_size_ori), height=img_size_ori, width=img_size_ori)), ] for ratio, aug in tqdm(augments): selidx = np.random.choice(x_train.shape[0], int(x_train.shape[0] * ratio), replace=False) for idx in tqdm(selidx): augmented = aug(image=x_train[idx], mask=msk_train[idx]) aimg = augmented['image'] amsk = augmented['mask'] if len(aimg.shape) < 3: aimg = aimg[...,np.newaxis] if len(amsk.shape) < 3: amsk = amsk[...,np.newaxis] aug_img.append(aimg) aug_msk.append(amsk) aug_y.append(y_train[idx]) aug_img = np.asarray(aug_img) aug_msk = np.asarray(aug_msk) aug_y = np.asarray(aug_y) x_train = np.append(x_train, aug_img, axis=0) msk_train = np.append(msk_train, aug_msk, axis=0) y_train = np.append(y_train, aug_y, axis=0) print(x_train.shape) print(msk_train.shape) print(y_train.shape) return x_train, msk_train, y_train
def aug_train_heavy(resolution, p=1.0): return Compose([ Resize(resolution, resolution), OneOf([ RandomRotate90(), Flip(), Transpose(), HorizontalFlip(), VerticalFlip() ], p=1.0), OneOf([ IAAAdditiveGaussianNoise(), GaussNoise(), ], p=0.5), OneOf([ MotionBlur(p=.2), MedianBlur(blur_limit=3, p=0.1), Blur(blur_limit=3, p=0.1), ], p=0.2), ShiftScaleRotate( shift_limit=0.0625, scale_limit=0.2, rotate_limit=45, p=0.2), OneOf([ OpticalDistortion(p=0.3), GridDistortion(p=.1), IAAPiecewiseAffine(p=0.3), ], p=0.1), OneOf([ CLAHE(clip_limit=2), IAASharpen(), IAAEmboss(), RandomBrightnessContrast(), ], p=0.3), HueSaturationValue(p=0.3), Normalize() ], p=p)
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 get_transform(model_name): if 'TAPNet' in model_name: # transform for sequences of images is very tricky # TODO: more transforms should be adopted for better results train_transform_ops = [ PadIfNeeded(min_height=args.input_height, min_width=args.input_width, p=1), Normalize(p=1), # optional transforms Resize(height=args.input_height, width=args.input_width, p=1), # CenterCrop(height=args.input_height, width=args.input_width, p=1) ] else: train_transform_ops = [ VerticalFlip(p=0.5), HorizontalFlip(p=0.5), PadIfNeeded(min_height=args.input_height, min_width=args.input_width, p=1), Normalize(p=1), # optional transforms # Resize(height=args.input_height, width=args.input_width, p=1), # CenterCrop(height=args.input_height, width=args.input_width, p=1) RandomCrop(height=args.input_height, width=args.input_width, p=1), ] valid_transform_ops = [ Normalize(p=1), PadIfNeeded(min_height=args.input_height, min_width=args.input_width, p=1), # optional transforms Resize(height=args.input_height, width=args.input_width, p=1), # CenterCrop(height=args.input_height, width=args.input_width, p=1) ] return Compose( train_transform_ops, p=1, ), Compose(valid_transform_ops, p=1)
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 __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 get_transforms(*, data): assert data in ('train', 'valid') if data == 'train': return Compose([ HorizontalFlip(p=0.5), VerticalFlip(p=0.5), Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], ), ToTensorV2(), ]) elif data == 'valid': return Compose([ Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], ), ToTensorV2(), ])
def strong_aug(p=0.75): return Compose( [ ShiftScaleRotate(scale_limit=0.1, rotate_limit=90), Transpose(), #IAAAffine(shear=0.1), #IAAPerspective(), Cutout(num_holes=20, max_h_size=8, max_w_size=8), HorizontalFlip(), VerticalFlip(), GaussNoise(), JpegCompression(), #RandomShadow(shadow_roi=(0, 0, 1, 1), p=0.75), OneOf([MotionBlur(), GaussianBlur()]), OneOf([ToGray(), ToSepia()]), RandomBrightnessContrast(brightness_limit=0.75, p=0.75) ], bbox_params=BboxParams("pascal_voc", label_fields=["category_id"], min_area=0.0, min_visibility=0.5), p=p)
def __init__(self, dataframe, vertical_flip, horizontal_flip, is_train=True): self.dataframe, self.is_train = dataframe, is_train self.vertical_flip, self.horizontal_flip = vertical_flip, horizontal_flip # Data Augmentation (custom for each dataset type) if is_train: self.transform = Compose([RandomResizedCrop(height=224, width=224, scale=(0.7, 1.0)), ShiftScaleRotate(rotate_limit=90, scale_limit = [0.7, 1]), HorizontalFlip(p = self.horizontal_flip), VerticalFlip(p = self.vertical_flip), HueSaturationValue(sat_shift_limit=[0.7, 1.3], hue_shift_limit=[-0.1, 0.1]), RandomBrightnessContrast(brightness_limit=[0.01, 0.1], contrast_limit= [0.01, 0.1]), #Normalize(), ToTensor()]) else: self.transform = Compose([ # Normalize(), Resize(height=224, width=224), ToTensor()])
def make_seq_transf(seq_len, NUM_CHAN, image_preproc=None): additional_targets = OrderedDict() for i in range(seq_len): for j in range(NUM_CHAN): if i == 0 and j == 0: continue additional_targets['image' + '_' + str(i) + '_' + str(j)] = 'image' aug = Compose([ VerticalFlip(p=0.4), RandomRotate90(p=0.4), Transpose(p=0.4), RandomBrightnessContrast( p=0.5, brightness_limit=0.3, contrast_limit=0.3, ), ], additional_targets=additional_targets) return lambda imgs: seq_transformer(aug, imgs, additional_targets, seq_len, image_preproc)
def __call__(self, original_image): self.augmentation_pipeline = Compose( [ Resize(650, 650, always_apply=True), HorizontalFlip(p=0.5), VerticalFlip(p=0.5), ShiftScaleRotate(rotate_limit=25.0, p=0.7), Resize(self.height, self.width, always_apply=True), Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], always_apply=True ), ToTensor() ] ) augmented = self.augmentation_pipeline( image=original_image ) image = augmented["image"] return image
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 image_loader(path, batch_size, augmentation=None, samples_per_class=64, val_samples_per_class=64, random_state=24): if augmentation is None: transform = Compose([ Normalize(mean=config.img_means, std=config.img_stds), ToTensor() ]) else: transform = Compose([ HorizontalFlip(p=AUG[augmentation]['p_fliph']), VerticalFlip(p=AUG[augmentation]['p_flipv']), OneOf([RandomSizedCrop((64, 90), 96, 96, p=1.0, interpolation=cv2.INTER_LANCZOS4)], p=AUG[augmentation]['p_crop']), ShiftScaleRotate(p=AUG[augmentation]['p_ssr'], interpolation=cv2.INTER_LANCZOS4), OneOf([RGBShift(p=1.0), RandomBrightness(p=1.0, limit=0.35), Blur(p=1.0, blur_limit=4), OpticalDistortion(p=1.0), GridDistortion(p=1.0)], p=AUG[augmentation]['p_aug']), Cutout(num_holes=2, max_h_size=24, max_w_size=24, p=AUG[augmentation]['p_cut']), Normalize(mean=config.img_means, std=config.img_stds), ToTensor() ]) transform_val = Compose([ Normalize(mean=config.img_means, std=config.img_stds), ToTensor() ]) sup_train_data = AlbumentationsDataset( f'{path}/supervised/train', transform=transform, samples_per_class=samples_per_class, random_state=random_state) sup_val_data = AlbumentationsDataset( f'{path}/supervised/val', transform=transform_val, samples_per_class=val_samples_per_class, random_state=random_state) data_loader_sup_train = torch.utils.data.DataLoader( sup_train_data, batch_size=batch_size, shuffle=True, num_workers=0) data_loader_sup_val = torch.utils.data.DataLoader( sup_val_data, batch_size=batch_size, shuffle=False, num_workers=0) return data_loader_sup_train, data_loader_sup_val
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