Search code examples
pythontensorflowopencvimage-augmentationalbumentations

I got error trying to use albumentations on tensorflow data pipeline

Im pretty new at deep learning and tensorflow, then when i try to use albumentations on tensorflow data pipeline, this error occurs (i use google colabs):

error: OpenCV(4.1.2) /io/opencv/modules/core/src/matrix.cpp:757: error: (-215:Assertion failed) dims <= 2 && step[0] > 0 in function 'locateROI'

Here is my code:

  • First i load the dataset with keras.utils.image_dataset_from_directory and do some normalization
images= keras.utils.image_dataset_from_directory(
    directory= IMAGE_DIR, 
    batch_size= BATCH_SIZE,
    image_size= RESIZED_IMAGE_SIZE,
    shuffle= True,
    labels='inferred',
    class_names= CLASS_LABELS
)

def normalize(image, label):
    return image/255, label

images= images.map(normalize)
  • Then i try to plot with this functions
def view_image(ds, class_labels):
    image, label = next(iter(ds))
    image = image.numpy()
    label = label.numpy()
    
    fig = plt.figure(figsize=(22, 22))
    for i in range(5):
        try:
            ax = fig.add_subplot(1, 5, i+1, xticks=[], yticks=[])
            ax.imshow(image[i])
            ax.set_title(f"{class_labels[label[i]]}")
        except:
            pass

view_image(images, CLASS_LABELS)

Here is the output

Kinda messed up but at least i know that the data loading was correct

from albumentations.core.composition import OneOf
transforms = A.Compose([
            # A.RandomResizedCrop(height= RESIZED_IMAGE_SIZE[0], width= RESIZED_IMAGE_SIZE[1], scale= (0.75, 1), p=0.8),
            # A.Rotate(limit=50, p=1),
            # A.ColorJitter(brightness=0.75, contrast=0.4, saturation=0.5, hue= 0, p=1),
            A.GaussianBlur(p=1),
            # A.Flip(p=0.25),
            # A.OneOf([
            #     A.GridDistortion(p=0.5),
            #     A.OpticalDistortion(p=0.5)
            # ], p=1),
        ])

def aug_fn(image, img_size):
    data = {"image":image}
    aug_data = transforms(**data)
    aug_img = aug_data["image"]
    aug_img = tf.cast(aug_img, tf.float32)
    
    return aug_img

def process_data(image, label, img_size):
    aug_img = tf.numpy_function(func=aug_fn, inp=[image, img_size], Tout=tf.float32)
    return aug_img, label

aug_images= images.map(partial(process_data, img_size=[416, 312]), num_parallel_calls=AUTOTUNE).prefetch(AUTOTUNE)
  • Then when i called view_image() function or simply the next(iter(aug_images)) it then raised the error i specified above. Here is the full error msg:
UnknownError: error: OpenCV(4.1.2) /io/opencv/modules/core/src/matrix.cpp:757: error: (-215:Assertion failed) dims <= 2 && step[0] > 0 in function 'locateROI'

Traceback (most recent call last):

  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/script_ops.py", line 271, in __call__
    ret = func(*args)

  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/impl/api.py", line 642, in wrapper
    return func(*args, **kwargs)

  File "<ipython-input-49-d101fe365707>", line 3, in aug_fn
    aug_data = transforms(**data)

  File "/usr/local/lib/python3.7/dist-packages/albumentations/core/composition.py", line 210, in __call__
    data = t(force_apply=force_apply, **data)

  File "/usr/local/lib/python3.7/dist-packages/albumentations/core/transforms_interface.py", line 97, in __call__
    return self.apply_with_params(params, **kwargs)

  File "/usr/local/lib/python3.7/dist-packages/albumentations/core/transforms_interface.py", line 112, in apply_with_params
    res[key] = target_function(arg, **dict(params, **target_dependencies))

  File "/usr/local/lib/python3.7/dist-packages/albumentations/augmentations/transforms.py", line 1963, in apply
    return F.gaussian_blur(image, ksize, sigma=sigma)

  File "/usr/local/lib/python3.7/dist-packages/albumentations/augmentations/functional.py", line 54, in wrapped_function
    result = func(img, *args, **kwargs)

  File "/usr/local/lib/python3.7/dist-packages/albumentations/augmentations/functional.py", line 615, in gaussian_blur
    return blur_fn(img)

  File "/usr/local/lib/python3.7/dist-packages/albumentations/augmentations/functional.py", line 189, in __process_fn
    img = process_fn(img, **kwargs)

cv2.error: OpenCV(4.1.2) /io/opencv/modules/core/src/matrix.cpp:757: error: (-215:Assertion failed) dims <= 2 && step[0] > 0 in function 'locateROI'



     [[{{node PyFunc}}]] [Op:IteratorGetNext]

What i have done

I tried to reread the tensorflow code sample several times and i didnt think I was doing anything wrong (idk i might be wrong tho), then i decided to create the map function below, it still doesnt work.

def augmentate_image(image, label):
    aug_img = transforms(image=image)['image']
    aug_img = tf.cast(aug_img, tf.float32)
    return aug_img, label

I would be very grateful if someone answered it :')

Solution

  • In the albumentations implementations with tensorflow docs it states that the dataset losses their shape after applying tf.numpy_function. You need to reset the shape of the data. You can try this bit of code.

    def set_shapes(img, label, img_shape=<you desired shape in 3d>):
    img.set_shape(img_shape)
    label.set_shape([])
    return img, label

    Then map it to aug_images; aug_images = aug_images.map(set_shapes,num_parallel_calls=AUTOTUNE).prefetch(AUTOTUNE)