custom augmentation in keras_segmentation gives me error

[webcluster4u]

Hi, I am uisng a very similar code to the helping documents for doing a sample custom augmentation in keras_segmentation

from keras_segmentation.models.unet import vgg_unet
from imgaug import augmenters as iaa

def custom_augmentation():
    return  iaa.Sequential(
        [
            # apply the following augmenters to most images
            iaa.Fliplr(0.5),  # horizontally flip 50% of all images
            iaa.Flipud(0.5), # horizontally flip 50% of all images
        ])

model = vgg_unet(n_classes=51 ,  input_height=416, input_width=608)

model.train(
    train_images =  "dataset1/images_prepped_train/",
    train_annotations = "dataset1/annotations_prepped_train/",
    checkpoints_path = "/tmp/vgg_unet_1" , epochs=5, 
    do_augment=True, # enable augmentation 
    custom_augmentation=custom_augmentation # sets the augmention function to use
)

However, I am getting below error.

Traceback (most recent call last):
  File "<stdin>", line 7, in <module>
TypeError: train() got an unexpected keyword argument 'custom_augmentation'

I am using Keras == 2.4.3 Tensorflow== 2.4.1

Can anyone advise me how to resolve above error?

[webcluster4u]

Can anyone help me on this issue?

[momoai96]

Hey :),

customizing is not really working in this version. You can change it in the scripts inside the package. Have a look in the issues. You can change to a newer branch

[mixxen]

I get the same error. Which commit do I need to roll back to for custom augmentation to work again?

[mixxen]

Use this repo instead:

https://github.com/sushantag9/image-segmentation-keras