Tensorflow data pipeline for 3D volume data

Question

I am currently developing a 3D segmentation model however I am a little stumped when it comes to creating a data pipeline in TensorFlow for this.

I have previously developed a 2D segmentation model, with a data pipeline similar in structure to

@tf_function
def parse_function(path_pair):
    
    img_path = path_pair[0]
    mask_path = path_pair[1]

    img_raw = tf.io.read_file(img_path)
    img = tf.io.decode_image(
        image_raw, channels = 1, dtype = th.uint8
    )

    #repeat for mask

    return img, mask

dataset = tf.data.Dataset.from_tensor_slices(path_pairs)
#path pairs is a list of tuples of strings of filepaths for image/mask pairs
dataset = dataset.shuffle(len(path_pairs))
dataset = dataset.map(parse_function)
dataset = dataset.batch(BATCH_SIZE)
dataset = dataset.prefetch(tf.data.AUTOTUNE)

My question is essentially, can I use a similar process for the 3D image data to output batches of

(BATCH_SIZE, IMG_HEIGHT, IMG_WIDTH, IMG_DEPTH, CHANNELS)

Within the parse_function(), tf.io.decode_image will not work (as far as I am aware) for my data as it anticipates 2D images + channel. My file format of the 3D images (single channel) can be whatever I need it to be. All of my data cannot be loaded into memory at once.

Tensorflow data pipeline for 3D volume data

0 Answers