I'm creating a tf.dataset object containing 2 images as inputs and a mask as target. All of them are 3D. After applying a custom map, the shape of the object changes from <RepeatDataset shapes: (((), ()), ()), types: ((tf.string, tf.string), tf.string)> to <PrefetchDataset shapes: (<unknown>, <unknown>, <unknown>), types: (tf.float32, tf.float32, tf.int32)>, and 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?
EDIT :
My random_crop_flip functions 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)
images = zip(img_bl_f, img_fu_f)
return images, mask_f
The zip isn't solving my problem. Is it possible to modify the return to get my desired output?
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)