How to apply 3D data augmentation to my dataset?

General Discussion
, , ,
#1

I have the augmentations in functions, this is one example:

def random_rotatation(img_numpy, min_angle, max_angle):
   all_axes = [(1, 0), (1, 2), (0, 2)]
   angle = np.random.randint(low=min_angle, high=max_angle+1)
   axes_random_id = np.random.randint(low=0, high=len(all_axes))
   axes = all_axes[axes_random_id]
   return rotate(img_numpy, angle, axes=axes)

And I want to implement them in my dataset, this is my pipeline:

def data_generator(nifti_files, mask_files):
    for nifti_file, mask_file in zip(nifti_files, mask_files):
        nifti_image = np.load(nifti_file)
        nifti_mask = np.load(mask_file)

        yield (nifti_image, nifti_mask)

# Create datasets
dataset = tf.data.Dataset.from_generator(
    lambda: data_generator(train_volumes, train_masks),
    output_signature=(
        tf.TensorSpec(shape=(128, 128, 128, 1), dtype=tf.float32),
        tf.TensorSpec(shape=(128, 128, 128, 4), dtype=tf.float32)
    )
)

dataset_val = tf.data.Dataset.from_generator(
    lambda: data_generator(val_volumes, val_masks),
    output_signature=(
        tf.TensorSpec(shape=(128, 128, 128, 1), dtype=tf.float32),
        tf.TensorSpec(shape=(128, 128, 128, 4), dtype=tf.float32)
    )
)

dataset = dataset.batch(4).prefetch(tf.data.experimental.AUTOTUNE)
dataset_val = dataset_val.batch(4).prefetch(tf.data.experimental.AUTOTUNE)

What should I do?

#3

Hi @matca, Once you have created a dataset, could you please try by using the map method like

augmented_dataset=dataset.map(random_rotatation)

Thank You.