How to apply volumentations to dataset for semantic segmentation

[kenbridge6]

Hi, thanks for sharing this variable tool!
I'm tring to use volumentations to augment 3d images for semantic segmentation, but I've got an error.
The code is like this:

#21 images and masks with the size of 64*64*64 voxels
#image_list.shape: (21, 64, 64, 64, 1) 
#mask_list.shape: (21, 64, 64, 64, 1) 

def get_augmentation(patch_size):
    return Compose([
        Rotate((-5, 5), (-5, 5), (-5, 5), p=0.5),
        Flip(0, p=0.5),
    ], p=1.0)

aug = get_augmentation((64, 64, 64))

def aug_image_mask(image, mask):    
    aug_data = aug(**{'image':image, 'mask':mask})
    img, msk = aug_data['image'], aug_data['mask']
    return img, msk

@tf.function
def aug_dataset(image, mask):
    aug_image, aug_mask = tf.numpy_function(func=aug_image_mask, inp=[image, mask], Tout=tf.float32)
    return aug_image, aug_mask

train_dataset = tf.data.Dataset.from_tensor_slices((image_list, mask_list))
aug_train_dataset = train_dataset.map(aug_dataset)

And I've got a following error message:
OperatorNotAllowedInGraphError: Iterating over a symbolic tf.Tensor is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

I'd like to know how to make an augmented dataset.

Thanks for your help in advance!

[ZFTurbo]

From error message I can see you try to apply function to tf.Tensor. Please check if you give on input of function numpy arrays.

[kenbridge6]

Thank you for your comment!
Finally I could augment images offline by applying the function to original numpy arrays:)