How to mantain the nested structure of a tf.dataset after applying map?

Question

I'm creating a tf.dataset object containing 2 images as inputs and a mask as target. All of them are 3D in grayscale. After applying a custom map, the shape of the object changes from ((TensorSpec(shape=(), dtype=tf.string, name=None), TensorSpec(shape=(), dtype=tf.string, name=None)), TensorSpec(shape=(), dtype=tf.string, name=None)) to (TensorSpec(shape=<unknown>, dtype=tf.float32, name=None), TensorSpec(shape=<unknown>, dtype=tf.float32, name=None), TensorSpec(shape=<unknown>, dtype=tf.int32, name=None)), losing the nested structure. When I fit the data, my model throws an error because it only detects one input instead of 2. Here is what I'm doing:

x, y = get_filenames(train_data_path, img_type='FLAIR')
z = get_filenames(train_data_path, img_type='mask')

path_dataset = tf.data.Dataset.from_tensor_slices((x, y))
mask_dataset = tf.data.Dataset.from_tensor_slices(z)
dataset = tf.data.Dataset.zip((path_dataset, mask_dataset)).shuffle(50).repeat(10)

ds = dataset. \
    map(lambda xx, zz: ((tf.py_function(load, [xx], [tf.float32, tf.float32])),
                        tf.py_function(load_mask, [zz], [tf.int32])),
        num_parallel_calls=tf.data.AUTOTUNE)

ds = ds.map(lambda xx, zz: (tf.py_function(random_crop_flip, [xx, zz],
                                           [tf.float32, tf.float32, tf.int32])),
            num_parallel_calls=tf.data.AUTOTUNE)
ds = ds.batch(2)
ds = ds.prefetch(tf.data.AUTOTUNE)

I can't map separately the images and the masks because they need the same seed for the random cropping and flipping. Is it possible to change the shape after the map so that I can feed it to my 2 input model?

My random_crop_flip function is as follows:

def random_crop_flip(images, mask, width=128, height=128, depth=128):
    img_bl, img_fu = images
    img_bl = img_bl.numpy()
    img_fu = img_fu.numpy()
    mask = mask.numpy()
    x_rand = random.randint(0, img_bl.shape[2] - width)
    y_rand = random.randint(0, img_bl.shape[1] - height)
    z_rand = random.randint(0, img_bl.shape[3] - depth)
    img_bl_f = img_bl[:, y_rand:y_rand + height, x_rand:x_rand + width, z_rand:z_rand + depth, :]
    img_fu_f = img_fu[:, y_rand:y_rand + height, x_rand:x_rand + width, z_rand:z_rand + depth, :]
    mask_f = mask[:, y_rand:y_rand + height, x_rand:x_rand + width, z_rand:z_rand + depth, :]
    flip_x = random.choice([True, False])
    flip_y = random.choice([True, False])
    flip_z = random.choice([True, False])

    if flip_x:
        img_bl_f = np.flip(img_bl_f, axis=2)
        img_fu_f = np.flip(img_fu_f, axis=2)
        mask_f = np.flip(mask_f, axis=2)

    if flip_y:
        img_bl_f = np.flip(img_bl_f, axis=1)
        img_fu_f = np.flip(img_fu_f, axis=1)
        mask_f = np.flip(mask_f, axis=1)

    if flip_z:
        img_bl_f = np.flip(img_bl_f, axis=3)
        img_fu_f = np.flip(img_fu_f, axis=3)
        mask_f = np.flip(mask_f, axis=3)

    return (img_bl_f, img_fu_f), mask_f

The tuple in the output isn't solving my problem. Is it possible to modify the return to get my desired output?

Answer 1

I have managed to solve this problem by "flattening" (eliminating the parenthesis) the return of the random_crop_flip, and applying another map on top of them, where I specified the shapes and returned my desired structure (x ,y), z:

def _set_shapes(img_bl, img_fu, mask):
    img_bl.set_shape([128, 128, 128, 1])
    img_fu.set_shape([128, 128, 128, 1])
    mask.set_shape([128, 128, 128, 1])

    return (img_bl, img_fu), mask 

Then my code looks like this:

x, y = get_filenames(train_data_path, img_type='FLAIR')
z = get_filenames(train_data_path, img_type='mask')

path_dataset = tf.data.Dataset.from_tensor_slices((x, y))
mask_dataset = tf.data.Dataset.from_tensor_slices(z)
dataset = tf.data.Dataset.zip((path_dataset, mask_dataset)).shuffle(50).repeat(10)

ds = dataset. \
    map(lambda xx, zz: ((tf.py_function(load, [xx], [tf.float32, tf.float32])),
                        tf.py_function(load_mask, [zz], [tf.int32])),
        num_parallel_calls=tf.data.AUTOTUNE)

ds = ds.map(lambda xx, zz: (tf.py_function(random_crop_flip, [xx, zz],
                                           [tf.float32, tf.float32, tf.int32])),
            num_parallel_calls=tf.data.AUTOTUNE)
ds = ds.map(_set_shapes)
ds = ds.batch(2)
ds = ds.prefetch(tf.data.AUTOTUNE)