def test_random_distributions_2D(self): ### test whether all 4 possible mirrorings occur in approximately equal frquencies in 2D batch_gen = BasicDataLoader((self.x_2D, self.y_2D), self.batch_size, number_of_threads_in_multithreaded=None) batch_gen = SingleThreadedAugmenter(batch_gen, MirrorTransform((0, 1))) counts = np.zeros(shape=(4,)) for b in range(self.num_batches): batch = next(batch_gen) for ix in range(self.batch_size): if (batch['data'][ix, :, :, :] == self.cam_left).all(): counts[0] = counts[0] + 1 elif (batch['data'][ix, :, :, :] == self.cam_updown).all(): counts[1] = counts[1] + 1 elif (batch['data'][ix, :, :, :] == self.cam_updown_left).all(): counts[2] = counts[2] + 1 elif (batch['data'][ix, :, :, :] == self.cam).all(): counts[3] = counts[3] + 1 self.assertTrue([1 if (2200 < c < 2800) else 0 for c in counts] == [1]*4, "2D Images were not mirrored along " "all axes with equal probability. " "This may also indicate that " "mirroring is not working")
def main(): args = get_arguments() utils.reproducibility(args, seed) # utils.make_dirs(args.save) if not os.path.exists(args.save): os.makedirs(args.save) # training_generator, val_generator, full_volume, affine = medical_loaders.generate_datasets(args, training_generator, val_generator, full_volume, affine, dataset = medical_loaders.generate_datasets(args, path='/data/hejy/MedicalZooPytorch_2cls/datasets') model, optimizer = medzoo.create_model(args) criterion = DiceLoss(classes=2, skip_index_after=args.classes, weight = torch.tensor([1, 1]).cuda(), sigmoid_normalization=True) # criterion = WeightedCrossEntropyLoss() if args.cuda: model = model.cuda() # model.restore_checkpoint(args.pretrained) dataloader_train = MICCAI2020_RIBFRAC_DataLoader3D(dataset, args.batchSz, args.dim, num_threads_in_multithreaded=2) tr_transforms = get_train_transform(args.dim) training_generator_aug = SingleThreadedAugmenter(dataloader_train, tr_transforms,) trainer = train.Trainer(args, model, criterion, optimizer, train_data_loader=training_generator, valid_data_loader=val_generator, lr_scheduler=None, dataset = dataset, train_data_loader_aug=training_generator_aug) trainer.training()
def create_data_gen_pipeline(patient_data, cf, is_training=True): """ create mutli-threaded train/val/test batch generation and augmentation pipeline. :param patient_data: dictionary containing one dictionary per patient in the train/test subset. :param is_training: (optional) whether to perform data augmentation (training) or not (validation/testing) :return: multithreaded_generator """ # create instance of batch generator as first element in pipeline. data_gen = BatchGenerator(patient_data, batch_size=cf.batch_size, cf=cf) # add transformations to pipeline. my_transforms = [] if is_training: mirror_transform = Mirror(axes=cf.mirror_axes) #np.arange(cf.dim) my_transforms.append(mirror_transform) spatial_transform = SpatialTransform( patch_size=cf.patch_size[:cf.dim], patch_center_dist_from_border=cf.da_kwargs['rand_crop_dist'], do_elastic_deform=cf.da_kwargs['do_elastic_deform'], alpha=cf.da_kwargs['alpha'], sigma=cf.da_kwargs['sigma'], do_rotation=cf.da_kwargs['do_rotation'], angle_x=cf.da_kwargs['angle_x'], angle_y=cf.da_kwargs['angle_y'], angle_z=cf.da_kwargs['angle_z'], do_scale=cf.da_kwargs['do_scale'], scale=cf.da_kwargs['scale'], random_crop=cf.da_kwargs['random_crop']) my_transforms.append(spatial_transform) my_transforms.append( RandomChannelDeleteTransform(cf.droppable_channels, cf.channel_drop_p)) else: my_transforms.append( CenterCropTransform(crop_size=cf.patch_size[:cf.dim])) my_transforms.append( ConvertSegToBoundingBoxCoordinates( cf.dim, get_rois_from_seg_flag=False, class_specific_seg_flag=cf.class_specific_seg_flag)) all_transforms = Compose(my_transforms) if cf.debugging: multithreaded_generator = SingleThreadedAugmenter( data_gen, all_transforms) else: multithreaded_generator = MultiThreadedAugmenter( data_gen, all_transforms, num_processes=cf.n_workers, seeds=range(cf.n_workers)) return multithreaded_generator
def test_segmentations_2D(self): ### test whether segmentations are mirrored coherently with images batch_gen = BasicDataLoader((self.x_2D, self.y_2D), self.batch_size, number_of_threads_in_multithreaded=None) batch_gen = SingleThreadedAugmenter(batch_gen, MirrorTransform((0, 1))) equivalent = True for b in range(self.num_batches): batch = next(batch_gen) for ix in range(self.batch_size): if (batch['data'][ix] != batch['seg'][ix]).all(): equivalent = False self.assertTrue(equivalent, "2D images and seg were not mirrored in the same way (they should though because " "seg needs to match the corresponding data")
def main(): # --------- Parse arguments --------------------------------------------------------------- parser = argparse.ArgumentParser() parser.add_argument('--config', type=str, default='./config_0.yaml', help='Path to the configuration file.') parser.add_argument('--dice', action='store_true') # be aware of this argument!!! parser.add_argument('--start-epoch', default=0, type=int, metavar='N', help='manual epoch number (useful on restarts)') parser.add_argument('--resume', default='', type=str, metavar='PATH', help='path to latest checkpoint (default: none)') parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true', help='evaluate model on validation set') parser.add_argument('-i', '--inference', default='', type=str, metavar='PATH', help='run inference on data set and save results') # 1e-8 works well for lung masks but seems to prevent # rapid learning for nodule masks parser.add_argument('--no-cuda', action='store_true') parser.add_argument('--save') parser.add_argument('--seed', type=int, default=1) args = parser.parse_args() # ---------- get the config file(config.yaml) -------------------------------------------- config = get_config(args.config) args.cuda = not args.no_cuda and torch.cuda.is_available() args.save = os.path.join('./work', (datestr() + ' ' + config['filename'])) nll = True if config['dice']: nll = False weight_decay = config['weight_decay'] num_threads_for_kits19 = config['num_of_threads'] patch_size = (160, 160, 128) num_batch_per_epoch = config['num_batch_per_epoch'] setproctitle.setproctitle(args.save) start_epoch = 1 # -------- Record best kidney segmentation dice ------------------------------------------- best_tk = 0.0 torch.manual_seed(args.seed) if args.cuda: torch.cuda.manual_seed(args.seed) print("build vnet") # Embed attention module model = vnet.VNet(elu=False, nll=nll, attention=config['attention'], nclass=3) # mark batch_size = config['ngpu'] * config['batchSz'] save_iter = config['model_save_iter'] # batch_size = args.ngpu*args.batchSz gpu_ids = range(config['ngpu']) # print(gpu_ids) model.apply(weights_init) # ------- Resume training from saved model ----------------------------------------------- if config['resume']: if os.path.isfile(config['resume']): print("=> loading checkpoint '{}'".format(config['resume'])) checkpoint = torch.load(config['resume']) # .tar files if config['resume'].endswith('.tar'): # print(checkpoint, "tar") start_epoch = checkpoint['epoch'] best_tk = checkpoint['best_tk'] checkpoint_model = checkpoint['model_state_dict'] model.load_state_dict( {k.replace('module.', ''): v for k, v in checkpoint_model.items()}) # .pkl files for the whole model else: # print(checkpoint, "pkl") model.load_state_dict(checkpoint.state_dict()) print("=> loaded checkpoint (epoch {})".format( checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(config['resume'])) exit(-1) else: pass # ------- Which loss function to use ------------------------------------------------------ if nll: training = train_bg validate = test_bg # class_balance = True else: training = train_bg_dice validate = test_bg_dice # class_balance = False # ----------------------------------------------------------------------------------------- print(' + Number of params: {}'.format( sum([p.data.nelement() for p in model.parameters()]))) # -------- Set on GPU --------------------------------------------------------------------- if args.cuda: model = model.cuda() if os.path.exists(args.save): shutil.rmtree(args.save) # create the output directory os.makedirs(args.save) # save the config file to the output folder shutil.copy(args.config, os.path.join(args.save, 'config.yaml')) # kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {} # ------ Load Training and Validation set -------------------------------------------- preprocessed_folders = "/home/data_share/npy_data/" patients = get_list_of_patients( preprocessed_data_folder=preprocessed_folders) # split num_split cross-validation sets # train, val = get_split_deterministic( # patients, fold=0, num_splits=5, random_state=12345) train, val = patients[0:147], patients[147:189] # VALIDATION DATA CANNOT BE LOADED IN CASE DUE TO THE LARGE SHAPE... # PRINT VALIDATION CASES FOR LATER TEST USE!! print("Validation cases:\n", val) # set max shape for validation set shapes = [Kits2019DataLoader3D.load_patient( i)[0].shape[1:] for i in val] max_shape = np.max(shapes, 0) max_shape = np.max((max_shape, patch_size), 0) # data loading + augmentation dataloader_train = Kits2019DataLoader3D( train, batch_size, patch_size, num_threads_for_kits19) dataloader_validation = Kits2019DataLoader3D( val, batch_size * 2, patch_size, num_threads_for_kits19) tr_transforms = get_train_transform(patch_size, prob=config['prob']) # whether to use single/multiThreadedAugmenter ------------------------------------------ if num_threads_for_kits19 > 1: tr_gen = MultiThreadedAugmenter(dataloader_train, tr_transforms, num_processes=num_threads_for_kits19, num_cached_per_queue=3,seeds=None, pin_memory=True) val_gen = MultiThreadedAugmenter(dataloader_validation, None, num_processes=max(1, num_threads_for_kits19//2), num_cached_per_queue=1, seeds=None, pin_memory=False) tr_gen.restart() val_gen.restart() else: tr_gen = SingleThreadedAugmenter(dataloader_train, transform=tr_transforms) val_gen = SingleThreadedAugmenter(dataloader_validation, transform=None) # ------- Set learning rate scheduler ---------------------------------------------------- lr_schdl = lr_scheduler.LR_Scheduler(mode=config['lr_policy'], base_lr=config['lr'], num_epochs=config['nEpochs'], iters_per_epoch=num_batch_per_epoch, lr_step=config['step_size'], warmup_epochs=config['warmup_epochs']) # ------ Choose Optimizer ---------------------------------------------------------------- if config['opt'] == 'sgd': optimizer = optim.SGD(model.parameters(), lr=config['lr'], momentum=0.99, weight_decay=weight_decay) elif config['opt'] == 'adam': optimizer = optim.Adam( model.parameters(), lr=config['lr'], weight_decay=weight_decay) elif config['opt'] == 'rmsprop': optimizer = optim.RMSprop( model.parameters(), lr=config['lr'], weight_decay=weight_decay) lr_plateu = optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.5, verbose=True, threshold=1e-3, patience=5) # ------- Apex Mixed Precision Acceleration ---------------------------------------------- model, optimizer = amp.initialize(model, optimizer, opt_level="O1") model = nn.parallel.DataParallel(model, device_ids=gpu_ids) # ------- Save training data ------------------------------------------------------------- trainF = open(os.path.join(args.save, 'train.csv'), 'w') trainF.write('Epoch,Loss,Kidney_Dice,Tumor_Dice\n') testF = open(os.path.join(args.save, 'test.csv'), 'w') testF.write('Epoch,Loss,Kidney_Dice,Tumor_Dice\n ') # ------- Training Pipeline -------------------------------------------------------------- for epoch in range(start_epoch, config['nEpochs'] + start_epoch): torch.cuda.empty_cache() training(args, epoch, model, tr_gen, optimizer, trainF, config, lr_schdl) torch.cuda.empty_cache() print('==>lr decay to:', optimizer.param_groups[0]['lr']) print('testing validation set...') composite_dice = validate(args, epoch, model, val_gen, optimizer, testF, config, lr_plateu) torch.cuda.empty_cache() # save model with best result and routinely if composite_dice > best_tk or epoch % config['model_save_iter'] == 0: # model_name = 'vnet_epoch_step1_' + str(epoch) + '.pkl' model_name = 'vnet_step1_' + str(epoch) + '.tar' # torch.save(model, os.path.join(args.save, model_name)) torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'best_tk': best_tk }, os.path.join(args.save, model_name)) best_tk = composite_dice # ---------------------------------------------------------------------------------------- trainF.close() testF.close()
def run(self, img_data, seg_data): # Create a parser for the batchgenerators module data_generator = DataParser(img_data, seg_data) # Initialize empty transform list transforms = [] # Add mirror augmentation if self.mirror: aug_mirror = MirrorTransform(axes=self.config_mirror_axes) transforms.append(aug_mirror) # Add contrast augmentation if self.contrast: aug_contrast = ContrastAugmentationTransform( self.config_contrast_range, preserve_range=True, per_channel=True, p_per_sample=self.config_p_per_sample) transforms.append(aug_contrast) # Add brightness augmentation if self.brightness: aug_brightness = BrightnessMultiplicativeTransform( self.config_brightness_range, per_channel=True, p_per_sample=self.config_p_per_sample) transforms.append(aug_brightness) # Add gamma augmentation if self.gamma: aug_gamma = GammaTransform(self.config_gamma_range, invert_image=False, per_channel=True, retain_stats=True, p_per_sample=self.config_p_per_sample) transforms.append(aug_gamma) # Add gaussian noise augmentation if self.gaussian_noise: aug_gaussian_noise = GaussianNoiseTransform( self.config_gaussian_noise_range, p_per_sample=self.config_p_per_sample) transforms.append(aug_gaussian_noise) # Add spatial transformations as augmentation # (rotation, scaling, elastic deformation) if self.rotations or self.scaling or self.elastic_deform or \ self.cropping: # Identify patch shape (full image or cropping) if self.cropping: patch_shape = self.cropping_patch_shape else: patch_shape = img_data[0].shape[0:-1] # Assembling the spatial transformation aug_spatial_transform = SpatialTransform( patch_shape, [i // 2 for i in patch_shape], do_elastic_deform=self.elastic_deform, alpha=self.config_elastic_deform_alpha, sigma=self.config_elastic_deform_sigma, do_rotation=self.rotations, angle_x=self.config_rotations_angleX, angle_y=self.config_rotations_angleY, angle_z=self.config_rotations_angleZ, do_scale=self.scaling, scale=self.config_scaling_range, border_mode_data='constant', border_cval_data=0, border_mode_seg='constant', border_cval_seg=0, order_data=3, order_seg=0, p_el_per_sample=self.config_p_per_sample, p_rot_per_sample=self.config_p_per_sample, p_scale_per_sample=self.config_p_per_sample, random_crop=self.cropping) # Append spatial transformation to transformation list transforms.append(aug_spatial_transform) # Compose the batchgenerators transforms all_transforms = Compose(transforms) # Assemble transforms into a augmentation generator augmentation_generator = SingleThreadedAugmenter( data_generator, all_transforms) # Perform the data augmentation x times (x = cycles) aug_img_data = None aug_seg_data = None for i in range(0, self.cycles): # Run the computation process for the data augmentations augmentation = next(augmentation_generator) # Access augmentated data from the batchgenerators data structure if aug_img_data is None and aug_seg_data is None: aug_img_data = augmentation["data"] aug_seg_data = augmentation["seg"] # Concatenate the new data augmentated data with the cached data else: aug_img_data = np.concatenate( (augmentation["data"], aug_img_data), axis=0) aug_seg_data = np.concatenate( (augmentation["seg"], aug_seg_data), axis=0) # Transform data from channel-first back to channel-last structure # Data structure channel-first 3D: (batch, channel, x, y, z) # Data structure channel-last 3D: (batch, x, y, z, channel) aug_img_data = np.moveaxis(aug_img_data, 1, -1) aug_seg_data = np.moveaxis(aug_seg_data, 1, -1) # Return augmentated image and segmentation data return aug_img_data, aug_seg_data
def get_moreDA_augmentation_ae(dataloader_train, dataloader_val, patch_size, params=default_3D_augmentation_params, border_val_seg=-1, seeds_train=None, seeds_val=None, order_seg=1, order_data=3, deep_supervision_scales=None, soft_ds=False, classes=None, pin_memory=True, anisotropy=False, extra_label_keys=None, val_mode=False, use_conf=False, global_params=None): ''' Work as Dataloader with augmentation :return: train_loader, val_loader for each iterator, return {'data': (B, D, H, W), 'target': (B, D, H, W)} ''' if global_params: for key, value in global_params.items(): params[key] = value print(global_params) if not val_mode: assert params.get( 'mirror' ) is None, "old version of params, use new keyword do_mirror" tr_transforms = [] if params.get("selected_data_channels") is not None: tr_transforms.append( DataChannelSelectionTransform( params.get("selected_data_channels"))) if params.get("selected_seg_channels") is not None: tr_transforms.append( SegChannelSelectionTransform( params.get("selected_seg_channels"))) # anistropic setting if anisotropy or params.get("dummy_2D"): ignore_axes = (0, ) tr_transforms.append( Convert3DTo2DTransform(extra_label_keys=extra_label_keys)) patch_size = patch_size[1:] # 2D patch size print('Using dummy2d data augmentation') params["elastic_deform_alpha"] = (0., 200.) params["elastic_deform_sigma"] = (9., 13.) params["rotation_x"] = (-180. / 360 * 2. * np.pi, 180. / 360 * 2. * np.pi) params["rotation_y"] = (-0. / 360 * 2. * np.pi, 0. / 360 * 2. * np.pi) params["rotation_z"] = (-0. / 360 * 2. * np.pi, 0. / 360 * 2. * np.pi) else: ignore_axes = None # 1. Spatial Transform: rotation, scaling tr_transforms.append( SpatialTransform( patch_size, patch_center_dist_from_border=None, do_elastic_deform=params.get("do_elastic"), alpha=params.get("elastic_deform_alpha"), sigma=params.get("elastic_deform_sigma"), do_rotation=params.get("do_rotation"), angle_x=params.get("rotation_x"), angle_y=params.get("rotation_y"), angle_z=params.get("rotation_z"), p_rot_per_axis=params.get("rotation_p_per_axis"), do_scale=params.get("do_scaling"), scale=params.get("scale_range"), border_mode_data=params.get("border_mode_data"), border_cval_data=0, order_data=order_data, border_mode_seg="constant", border_cval_seg=border_val_seg, order_seg=order_seg, random_crop=params.get("random_crop"), p_el_per_sample=params.get("p_eldef"), p_scale_per_sample=params.get("p_scale"), p_rot_per_sample=params.get("p_rot"), independent_scale_for_each_axis=params.get( "independent_scale_factor_for_each_axis"), extra_label_keys=extra_label_keys, do_translate=params.get("do_translate"), p_trans=params.get("p_trans"), trans_max_shifts=params.get("trans_max_shifts"), trans_const_channel=params.get("trans_const_channel"), )) if anisotropy or params.get("dummy_2D"): tr_transforms.append( Convert2DTo3DTransform(extra_label_keys=extra_label_keys)) ''' # 2. Noise Augmentation: gaussian noise, gaussian blur tr_transforms.append(GaussianNoiseTransform(p_per_sample=0.1)) tr_transforms.append(GaussianBlurTransform((0.5, 1.), different_sigma_per_channel=True, p_per_sample=0.2, p_per_channel=0.5)) # 3. Color Augmentation: brightness, constrast, low resolution, gamma_transform tr_transforms.append(BrightnessMultiplicativeTransform(multiplier_range=(0.75, 1.25), p_per_sample=0.15)) if params.get("do_additive_brightness"): tr_transforms.append(BrightnessTransform(params.get("additive_brightness_mu"), params.get("additive_brightness_sigma"), True, p_per_sample=params.get("additive_brightness_p_per_sample"), p_per_channel=params.get("additive_brightness_p_per_channel"))) tr_transforms.append(ContrastAugmentationTransform(p_per_sample=0.15)) tr_transforms.append(SimulateLowResolutionTransform(zoom_range=(0.5, 1), per_channel=True, p_per_channel=0.5, order_downsample=0, order_upsample=3, p_per_sample=0.25, ignore_axes=ignore_axes)) tr_transforms.append( GammaTransform(params.get("gamma_range"), True, True, retain_stats=params.get("gamma_retain_stats"), p_per_sample=0.1)) # inverted gamma if params.get("do_gamma"): tr_transforms.append( GammaTransform(params.get("gamma_range"), False, True, retain_stats=params.get("gamma_retain_stats"), p_per_sample=params["p_gamma"])) ''' # 4. Mirror Transform if params.get("do_mirror") or params.get("mirror"): tr_transforms.append( MirrorTransform(params.get("mirror_axes"), extra_label_keys=extra_label_keys)) # if params.get("mask_was_used_for_normalization") is not None: # mask_was_used_for_normalization = params.get("mask_was_used_for_normalization") # tr_transforms.append(MaskTransform(mask_was_used_for_normalization, mask_idx_in_seg=0, set_outside_to=0)) tr_transforms.append( RemoveLabelTransform(-1, 0, extra_label_keys=extra_label_keys)) tr_transforms.append(RenameTransform('data', 'image', True)) tr_transforms.append(RenameTransform('seg', 'gt', True)) if deep_supervision_scales is not None: if soft_ds: assert classes is not None tr_transforms.append( DownsampleSegForDSTransform3(deep_supervision_scales, 'gt', 'gt', classes)) else: tr_transforms.append( DownsampleSegForDSTransform2( deep_supervision_scales, 0, 0, input_key='gt', output_key='gt', extra_label_keys=extra_label_keys)) toTensorKeys = [ 'image', 'gt' ] + extra_label_keys if extra_label_keys is not None else [ 'image', 'gt' ] tr_transforms.append(NumpyToTensor(toTensorKeys, 'float')) tr_transforms = Compose(tr_transforms) if seeds_train is not None: seeds_train = [seeds_train] * params.get('num_threads') if use_conf: num_threads = 1 num_cached_per_thread = 1 else: num_threads, num_cached_per_thread = params.get( 'num_threads'), params.get("num_cached_per_thread") batchgenerator_train = MultiThreadedAugmenter(dataloader_train, tr_transforms, num_threads, num_cached_per_thread, seeds=seeds_train, pin_memory=pin_memory) val_transforms = [] val_transforms.append( RemoveLabelTransform(-1, 0, extra_label_keys=extra_label_keys)) if params.get("selected_data_channels") is not None: val_transforms.append( DataChannelSelectionTransform( params.get("selected_data_channels"))) if params.get("selected_seg_channels") is not None: val_transforms.append( SegChannelSelectionTransform( params.get("selected_seg_channels"))) val_transforms.append(RenameTransform('data', 'image', True)) val_transforms.append(RenameTransform('seg', 'gt', True)) if deep_supervision_scales is not None: if soft_ds: assert classes is not None val_transforms.append( DownsampleSegForDSTransform3(deep_supervision_scales, 'gt', 'gt', classes)) else: val_transforms.append( DownsampleSegForDSTransform2( deep_supervision_scales, 0, 0, input_key='gt', output_key='gt', extra_label_keys=extra_label_keys)) val_transforms.append(NumpyToTensor(toTensorKeys, 'float')) val_transforms = Compose(val_transforms) # batchgenerator_val = MultiThreadedAugmenter(dataloader_val, val_transforms, max(params.get('num_threads') // 2, 1), # params.get("num_cached_per_thread"), # seeds=seeds_val, pin_memory=pin_memory) if seeds_val is not None: seeds_val = [seeds_val] * 1 # batchgenerator_val = MultiThreadedAugmenter(dataloader_val, val_transforms, 1, # params.get("num_cached_per_thread"), # seeds=seeds_val, pin_memory=False) batchgenerator_val = SingleThreadedAugmenter(dataloader_val, val_transforms) else: val_transforms = [] val_transforms.append( RemoveLabelTransform(-1, 0, extra_label_keys=extra_label_keys)) if params.get("selected_data_channels") is not None: val_transforms.append( DataChannelSelectionTransform( params.get("selected_data_channels"))) if params.get("selected_seg_channels") is not None: val_transforms.append( SegChannelSelectionTransform( params.get("selected_seg_channels"))) val_transforms.append(RenameTransform('data', 'image', True)) val_transforms.append(RenameTransform('seg', 'gt', True)) if deep_supervision_scales is not None: if soft_ds: assert classes is not None val_transforms.append( DownsampleSegForDSTransform3(deep_supervision_scales, 'gt', 'gt', classes)) else: val_transforms.append( DownsampleSegForDSTransform2( deep_supervision_scales, 0, 0, input_key='gt', output_key='gt', extra_label_keys=extra_label_keys)) toTensorKeys = [ 'image', 'gt' ] + extra_label_keys if extra_label_keys is not None else [ 'image', 'gt' ] val_transforms.append(NumpyToTensor(toTensorKeys, 'float')) val_transforms = Compose(val_transforms) batchgenerator_val = SingleThreadedAugmenter(dataloader_val, val_transforms) batchgenerator_train = SingleThreadedAugmenter(dataloader_train, val_transforms) return batchgenerator_train, batchgenerator_val
def __init__(self, data_loader: BaseDataLoader, transforms, n_process_augmentation, sampler, sampler_queues: list, num_cached_per_queue=2, seeds=None, **kwargs): """ Parameters ---------- data_loader : :class:`BaseDataLoader` the dataloader providing the actual data transforms : Callable or None the transforms to use. Can be single callable or None n_process_augmentation : int the number of processes to use for augmentation (only necessary if not in debug mode) sampler : :class:`AbstractSampler` the sampler to use; must be used here instead of inside the dataloader to avoid duplications and oversampling due to multiprocessing sampler_queues : list of :class:`multiprocessing.Queue` queues to pass the sample indices to the actual dataloader num_cached_per_queue : int the number of samples to cache per queue (only necessary if not in debug mode) seeds : int or list the seeds for each process (only necessary if not in debug mode) **kwargs : additional keyword arguments """ self._batchsize = data_loader.batch_size # don't use multiprocessing in debug mode if get_current_debug_mode(): augmenter = SingleThreadedAugmenter(data_loader, transforms) else: assert isinstance(n_process_augmentation, int) # no seeds are given -> use default seed of 1 if seeds is None: seeds = 1 # only an int is gien as seed -> replicate it for each process if isinstance(seeds, int): seeds = [seeds] * n_process_augmentation # avoid same seeds for all processes if any([seeds[0] == _seed for _seed in seeds[1:]]): for idx in range(len(seeds)): seeds[idx] = seeds[idx] + idx augmenter = MultiThreadedAugmenter( data_loader, transforms, num_processes=n_process_augmentation, num_cached_per_queue=num_cached_per_queue, seeds=seeds, **kwargs) self._augmenter = augmenter self._sampler = sampler self._sampler_queues = sampler_queues self._queue_id = 0
def get_moreDA_augmentation(dataloader_train, dataloader_val, patch_size, params=default_3D_augmentation_params, border_val_seg=-1, seeds_train=None, seeds_val=None, order_seg=1, order_data=3, deep_supervision_scales=None, soft_ds=False, classes=None, pin_memory=True): assert params.get( 'mirror') is None, "old version of params, use new keyword do_mirror" tr_transforms = [] if params.get("selected_data_channels") is not None: tr_transforms.append( DataChannelSelectionTransform( params.get("selected_data_channels"))) if params.get("selected_seg_channels") is not None: tr_transforms.append( SegChannelSelectionTransform(params.get("selected_seg_channels"))) # don't do color augmentations while in 2d mode with 3d data because the color channel is overloaded!! if params.get("dummy_2D") is not None and params.get("dummy_2D"): ignore_axes = (0, ) tr_transforms.append(Convert3DTo2DTransform()) else: ignore_axes = None tr_transforms.append( SpatialTransform(patch_size, patch_center_dist_from_border=None, do_elastic_deform=params.get("do_elastic"), alpha=params.get("elastic_deform_alpha"), sigma=params.get("elastic_deform_sigma"), do_rotation=params.get("do_rotation"), angle_x=params.get("rotation_x"), angle_y=params.get("rotation_y"), angle_z=params.get("rotation_z"), p_rot_per_axis=params.get("rotation_p_per_axis"), do_scale=params.get("do_scaling"), scale=params.get("scale_range"), border_mode_data=params.get("border_mode_data"), border_cval_data=0, order_data=order_data, border_mode_seg="constant", border_cval_seg=border_val_seg, order_seg=order_seg, random_crop=params.get("random_crop"), p_el_per_sample=params.get("p_eldef"), p_scale_per_sample=params.get("p_scale"), p_rot_per_sample=params.get("p_rot"), independent_scale_for_each_axis=params.get( "independent_scale_factor_for_each_axis"))) if params.get("dummy_2D"): tr_transforms.append(Convert2DTo3DTransform()) # we need to put the color augmentations after the dummy 2d part (if applicable). Otherwise the overloaded color # channel gets in the way tr_transforms.append(GaussianNoiseTransform(p_per_sample=0.1)) tr_transforms.append( GaussianBlurTransform((0.5, 1.), different_sigma_per_channel=True, p_per_sample=0.2, p_per_channel=0.5)) tr_transforms.append( BrightnessMultiplicativeTransform(multiplier_range=(0.75, 1.25), p_per_sample=0.15)) if params.get("do_additive_brightness"): tr_transforms.append( BrightnessTransform( params.get("additive_brightness_mu"), params.get("additive_brightness_sigma"), True, p_per_sample=params.get("additive_brightness_p_per_sample"), p_per_channel=params.get("additive_brightness_p_per_channel"))) tr_transforms.append(ContrastAugmentationTransform(p_per_sample=0.15)) tr_transforms.append( SimulateLowResolutionTransform(zoom_range=(0.5, 1), per_channel=True, p_per_channel=0.5, order_downsample=0, order_upsample=3, p_per_sample=0.25, ignore_axes=ignore_axes)) tr_transforms.append( GammaTransform(params.get("gamma_range"), True, True, retain_stats=params.get("gamma_retain_stats"), p_per_sample=0.1)) # inverted gamma if params.get("do_gamma"): tr_transforms.append( GammaTransform(params.get("gamma_range"), False, True, retain_stats=params.get("gamma_retain_stats"), p_per_sample=params["p_gamma"])) if params.get("do_mirror") or params.get("mirror"): tr_transforms.append(MirrorTransform(params.get("mirror_axes"))) if params.get("mask_was_used_for_normalization") is not None: mask_was_used_for_normalization = params.get( "mask_was_used_for_normalization") tr_transforms.append( MaskTransform(mask_was_used_for_normalization, mask_idx_in_seg=0, set_outside_to=0)) tr_transforms.append(RemoveLabelTransform(-1, 0)) if params.get("move_last_seg_chanel_to_data") is not None and params.get( "move_last_seg_chanel_to_data"): tr_transforms.append( MoveSegAsOneHotToData(1, params.get("all_segmentation_labels"), 'seg', 'data')) if params.get( "cascade_do_cascade_augmentations") is not None and params.get( "cascade_do_cascade_augmentations"): if params.get("cascade_random_binary_transform_p") > 0: tr_transforms.append( ApplyRandomBinaryOperatorTransform( channel_idx=list( range(-len(params.get("all_segmentation_labels")), 0)), p_per_sample=params.get( "cascade_random_binary_transform_p"), key="data", strel_size=params.get( "cascade_random_binary_transform_size"), p_per_label=params.get( "cascade_random_binary_transform_p_per_label"))) if params.get("cascade_remove_conn_comp_p") > 0: tr_transforms.append( RemoveRandomConnectedComponentFromOneHotEncodingTransform( channel_idx=list( range(-len(params.get("all_segmentation_labels")), 0)), key="data", p_per_sample=params.get("cascade_remove_conn_comp_p"), fill_with_other_class_p=params.get( "cascade_remove_conn_comp_max_size_percent_threshold" ), dont_do_if_covers_more_than_X_percent=params.get( "cascade_remove_conn_comp_fill_with_other_class_p") )) tr_transforms.append(RenameTransform('seg', 'target', True)) if deep_supervision_scales is not None: if soft_ds: assert classes is not None tr_transforms.append( DownsampleSegForDSTransform3(deep_supervision_scales, 'target', 'target', classes)) else: tr_transforms.append( DownsampleSegForDSTransform2(deep_supervision_scales, 0, 0, input_key='target', output_key='target')) tr_transforms.append(NumpyToTensor(['data', 'target'], 'float')) tr_transforms = Compose(tr_transforms) if config.DEBUG_MODE: # zhuc batchgenerator_train = SingleThreadedAugmenter(dataloader_train, tr_transforms) else: batchgenerator_train = MultiThreadedAugmenter( dataloader_train, tr_transforms, params.get('num_threads'), params.get("num_cached_per_thread"), seeds=seeds_train, pin_memory=pin_memory) val_transforms = [] val_transforms.append(RemoveLabelTransform(-1, 0)) if params.get("selected_data_channels") is not None: val_transforms.append( DataChannelSelectionTransform( params.get("selected_data_channels"))) if params.get("selected_seg_channels") is not None: val_transforms.append( SegChannelSelectionTransform(params.get("selected_seg_channels"))) if params.get("move_last_seg_chanel_to_data") is not None and params.get( "move_last_seg_chanel_to_data"): val_transforms.append( MoveSegAsOneHotToData(1, params.get("all_segmentation_labels"), 'seg', 'data')) val_transforms.append(RenameTransform('seg', 'target', True)) if deep_supervision_scales is not None: if soft_ds: assert classes is not None val_transforms.append( DownsampleSegForDSTransform3(deep_supervision_scales, 'target', 'target', classes)) else: val_transforms.append( DownsampleSegForDSTransform2(deep_supervision_scales, 0, 0, input_key='target', output_key='target')) val_transforms.append(NumpyToTensor(['data', 'target'], 'float')) val_transforms = Compose(val_transforms) if config.DEBUG_MODE: # zhuc batchgenerator_val = SingleThreadedAugmenter(dataloader_val, val_transforms) else: batchgenerator_val = MultiThreadedAugmenter( dataloader_val, val_transforms, max(params.get('num_threads') // 2, 1), params.get("num_cached_per_thread"), seeds=seeds_val, pin_memory=pin_memory) return batchgenerator_train, batchgenerator_val