Retrieve (or add) filenames when loading tensorflow dataset - tensorflow

Question

I am training a 3D CNN using keras and I would like to find a way to preserve the filenames from every input that comes into the dataset. The data consists on a series of .NII images that I load using tools from the Nibabel package. However, when I use the function get_fdata() from nibabel, I lose all the file names from the header and I just get the 3D arrays from the image (which I need to create the tensors that go into the model). I would like to know if there is a way to either preserve those names or add them as an attribute of each 3D array (maybe incorporating them into the data loader). I have followed the example from the keras CT scan classification, with some modifications for these fMRI images. Here I added some code to show what I have so far:
import os
from re import A
import zipfile
import numpy as np
import tensorflow as tf
import random
from tensorflow import keras
from tensorflow.keras import layers
import nibabel as nib
## Loading data and preprocessing
MNAME = "test_A.h5"
def read_nifti_file(filepath):
"""Read and load volume"""
# Read file
scan = nib.load(filepath)
# Get raw data
scan = scan.get_fdata()
return scan
scene_D_paths = np.asarray([
os.path.join(os.getcwd(),'data\scene_D', x)
for x in os.listdir("data\scene_D")
])
scene_H_paths = np.asarray([
os.path.join(os.getcwd(),'data\scene_H', x)
for x in os.listdir("data\scene_H")
])
# Read the scans.
D_scans = np.array([read_nifti_file(path) for path in scene_D_paths])
H_scans = np.array([read_nifti_file(path) for path in scene_H_paths])
# One hot encoding
D_labels = np.array([[1., 0., 0., 0.] for _ in range(len(D_scans))])
H_labels = np.array([[0., 1., 0., 0.] for _ in range(len(H_scans))])
# Divide dataset in 0.7 training and 0.3 testing
# We use a single randomization per scene and we keep track of the indexing
def split_dataset(scene):
n_images = len(scene)
n_train = round((n_images-1)*.7)
train_idx = np.full(n_images, False)
train_idx[random.sample(range(0, n_images), n_train)] = True
test_idx = np.full(len(train_idx), True)
test_idx[train_idx] = False
return(train_idx, test_idx)
D_index = split_dataset(scene_D_paths)
H_index = split_dataset(scene_H_paths)
paths_train = np.concatenate((scene_D_paths[D_index[0]], scene_H_paths[H_index[0]]), axis = 0)
x_train = np.concatenate((D_scans[D_index[0]], H_scans[H_index[0]]), axis=0)
y_train = np.concatenate((D_labels[D_index[0]], H_labels[H_index[0]]), axis=0)
paths_val = np.concatenate((scene_D_paths[D_index[1]], scene_H_paths[H_index[1]]), axis = 0)
x_val = np.concatenate((D_scans[D_index[1]], H_scans[H_index[1]]), axis=0)
y_val = np.concatenate((D_labels[D_index[1]], H_labels[H_index[1]]), axis=0)
# Define data loaders.
train_filenames = tf.constant(paths_train)
train_loader = tf.data.Dataset.from_tensor_slices((#train_filenames,
x_train, y_train))
validation_filenames = tf.constant(paths_val)
validation_loader = tf.data.Dataset.from_tensor_slices((#validation_filenames,
x_val, y_val))
# If increase this value, reduce learning rate
batch_size = 2
# Augment the on the fly during training.
train_dataset = (
train_loader.shuffle(len(x_train))
.batch(batch_size)
.prefetch(2)
)
# Only rescale.
validation_dataset = (
validation_loader.shuffle(len(x_val))
.batch(batch_size)
.prefetch(2)
)
def get_model(width=53, height=63, depth=52):
"""Build a 3D convolutional neural network model."""
inputs = keras.Input((width, height, depth, 1))
x = layers.Conv3D(filters=52, kernel_size=3, activation="relu")(inputs)
x = layers.MaxPool3D(pool_size=2)(x)
x = layers.BatchNormalization()(x)
x = layers.Conv3D(filters=52, kernel_size=3, activation="relu")(x)
x = layers.MaxPool3D(pool_size=2)(x)
x = layers.BatchNormalization()(x)
x = layers.Conv3D(filters=104, kernel_size=3, activation="relu")(x)
x = layers.MaxPool3D(pool_size=2)(x)
x = layers.BatchNormalization()(x)
x = layers.Conv3D(filters=208, kernel_size=3, activation="relu")(x)
x = layers.MaxPool3D(pool_size=2)(x)
x = layers.BatchNormalization()(x)
# Tweak the number of units. More will tend to overfit to the training, less will tend to underfit
# We can increase dropout or reduce the number of parameters
# Reduce number of epoch during the initial training
x = layers.GlobalAveragePooling3D()(x)
x = layers.Dense(units=208, activation="relu")(x)
x = layers.Dropout(0.3)(x)
outputs = layers.Dense(units=4, activation="softmax")(x)
# Define the model.
model = keras.Model(inputs, outputs, name="3d_tinn_fmri")
return model
# Build model.
model = get_model(width=53, height=63, depth=52)
model.summary()
Thanks for your help,
Eduardo