Bug? Properly using images and masks with tensorflow datasets - cv2.error in function 'warpAffine' when applying Rotate, CLAHE, OpticalDistortion..

Hi , I am having difficulties to run a segmentation task, where I use images and masks. Things that I tried:

def transform(image, mask):
   
    augmentations = A.Compose(
    [
        A.HorizontalFlip(p = 0.5), # Apply horizontal flip to 50% of images
        A.OneOf(
            [
                A.RandomContrast(), 
                A.RandomGamma(), 
                A.RandomBrightness(), 
            ],
            p = 0.5 
        ),
        A.Rotate(p=0.3),
        A.CLAHE (p=0.5),
        
    ]
    )   
    return augmentations(image=image)['image'], augmentations(mask=mask)['mask']
    
    
def prepare(ds, augment=True):
    
    ds = ds.map(resize_and_rescale, num_parallel_calls=opt.AUTOTUNE)
   
    
    # Batch dataset
    ds = ds.batch(opt.BATCH_SIZE)
    
    if augment:
      ds = ds.map(transform, num_parallel_calls=opt.AUTOTUNE)
    
    return ds.prefetch(buffer_size=opt.AUTOTUNE)


train_ds = prepare(train_ds)

This results in : TypeError: image must be numpy array type and I can’t find a way to overcome this.

Another approach that I used:

def aug_fn_img(image):

    data = {"image": image}
    aug_data = transform(**data)
    aug_img = aug_data["image"]

    return aug_img

def aug_fn_mask(mask):

    data = {'mask': mask}
    aug_data = transform(**data)
    aug_mask = aug_data["mask"]

    return aug_mask

def process_aug(image, mask):

    aug_img = tf.numpy_function(func=aug_fn_img, inp=[image],
                                          Tout=[tf.float32])
    
    aug_mask = tf.numpy_function(func=aug_fn_mask, inp=[mask],
                                          Tout=[tf.float32])
    return aug_img, aug_mask

and then call if augment: ds = ds.map(process_aug, num_parallel_calls=opt.AUTOTUNE)

it gives me : ValueError: as_list() is not defined on an unknown TensorShape.

How can I deal with images and masks with tensorflow?

[Andredance]

I think that the main problem is that you are trying to augment images after you batch them (Albumentations process one image/mask at a time). The proper way would be:

def prepare(ds, shuffle=False, augment=False):
    
    ds = ds.map(resize_and_rescale, num_parallel_calls=AUTOTUNE)
    
    if augment:
      ds = ds.map(process_data, num_parallel_calls=AUTOTUNE)

    if shuffle:
        ds = ds.shuffle(buffer_size=1000)
    
    ds = ds.batch(BATCH_SIZE)
    return ds.prefetch(buffer_size=AUTOTUNE)

Also, you may find a bit better way to resize your images. Maybe you can try something from Albumentations like: LongestMaxSize + PadIfNeeded. But this all is your choice.

[ggous]

Thanks @Andredance ! It works fine now!