Apply preprocessing only to images , not to masks-CodePudding

I have a dataset which has images and masks.

images_p = tf.keras.utils.image_dataset_from_directory(
        path_imgs, batch_size=None, shuffle=False, label_mode=None)
    
masks_p = tf.keras.utils.image_dataset_from_directory(
        path_masks, batch_size=None, shuffle=False, label_mode=None,
        color_mode='grayscale')
    
dataset = tf.data.Dataset.zip((images_p, masks_p))

Now, I want to apply some preprocessing. But, it will be different for images and masks. For example, masks should not be scaled.

I have something like that:

def resize_and_rescale(image, mask=False):
    image = tf.image.resize(image, (IMG_HEIGHT, IMG_WIDTH))
    if mask is False:
        image = image / 255.0
    return image
    
def prepare(ds, shuffle=False, augment=False):
    
    # Resize and rescale 
    ds = ds.map(lambda x, y: (resize_and_rescale(x), y), 
                num_parallel_calls=AUTOTUNE)
    
    
    if shuffle:
        ds = ds.shuffle(buffer_size=1000)
    
    # Batch dataset
    ds = ds.batch(BATCH_SIZE)
    
    # Use data augmentation only on the training set
    if augment:
      ds = ds.map(lambda x, y: (transform(x), y), 
                  num_parallel_calls=AUTOTUNE)
    
    return ds.prefetch(buffer_size=AUTOTUNE)

How can I apply resize on both images and masks and rescale only to images?

I am splitting first the dataset.

train_ds, val_ds, test_ds = split_train_test_val(dataset,
                                                 SEED,
                                                 train_split=0.8, 
                                                 val_split=0.1, 
                                                 test_split=0.1,
                                                 shuffle=True, 
                                                 shuffle_size=1000)

And then I have to apply the prepare function.

Something like that:

train_ds = prepare(train_ds, shuffle=True, augment=True)
val_ds = prepare(val_ds)
test_ds = prepare(test_ds)

CodePudding user response：

Maybe just explicitly only resize the masks like this:

import tensorflow as tf

path_imgs = ('/content/images/*.jpg')
path_masks = ('/content/masks/*.jpg')

images_p = tf.keras.utils.image_dataset_from_directory(
        '/content/images/', batch_size=None, shuffle=False, label_mode=None)
    
masks_p = tf.keras.utils.image_dataset_from_directory(
        '/content/masks/', batch_size=None, shuffle=False, label_mode=None,
        color_mode='grayscale')
dataset = tf.data.Dataset.zip((images_p, masks_p))


IMG_HEIGHT = 64
IMG_WIDTH = 64
BATCH_SIZE = 2

def resize_and_rescale(image, mask):
    image = tf.image.resize(image, (IMG_HEIGHT, IMG_WIDTH))
    mask = tf.image.resize(mask, (IMG_HEIGHT, IMG_WIDTH))
    image = image / 255.0
    return image, mask
    
def prepare(ds, shuffle=False, augment=False):
    
    # Resize and rescale 
    ds = ds.map(resize_and_rescale, num_parallel_calls=tf.data.AUTOTUNE)
    
    
    if shuffle:
        ds = ds.shuffle(buffer_size=1000)
    
    # Batch dataset
    ds = ds.batch(BATCH_SIZE)

    return ds.prefetch(buffer_size=tf.data.AUTOTUNE)

ds = prepare(dataset)