I am trying to do image colorization. I have 5000 images (256x256x3) and would like not to load all data in my program (for memory reason). I have found that it is possible to use ImageDataGenerator.flow_from_directory()
but I use LAB images and I would like to feed my model with a numpy array of the L component (256, 256, 1). My targets are A and B components (256, 256, 2). To have my image I then merge the input and output to have a LAB image (256, 256, 3). The problem i that ImageDataGenerator.flow_from_directory()
only works with image type files (so a 256x256x3 image) and I would like to know if there is a way to do the same thing with numpy arrays.
I tried using tf.data.Dataset.list_files()
, I had all my files but I did not found how to load my numpy array to feed my model. I guess I need to use some sort of generator but I do not really understand how to use it. This is what I have for now :
HEIGHT = 256
WIDTH = HEIGHT
Batch_size = 50
dir_X_train = 'data/X_train_np/train_black_resized/*.npy'
dir_X_test = 'data/X_test/test_black_resized/*.npy'
dir_y_train = 'data/y_train_np/train_color_resized/*.npy'
dir_y_test = 'data/y_test/test_color_resized/*.npy'
X_train_dataset = tf.data.Dataset.list_files(dir_X_train, shuffle=False).batch(Batch_size)
y_train_dataset = tf.data.Dataset.list_files(dir_y_train, shuffle=False).batch(Batch_size)
def process_path(file_path):
return tf.io.read_file(file_path[0])
X_train_dataset = X_train_dataset.map(process_path)
y_train_dataset = y_train_dataset.map(process_path)
train_dataset = tf.data.Dataset.zip((X_train_dataset, y_train_dataset))
for image_black, image_color in train_dataset.take(1):
print(image_black.numpy()[:100])
print(type(image_black))
print(image_color.numpy()[:100])
print(type(image_color))
Output :
b"\x93NUMPY\x01\x00v\x00{'descr': '<f4', 'fortran_order': False, 'shape': (256, 256), } "
<class 'tensorflow.python.framework.ops.EagerTensor'>
b"\x93NUMPY\x01\x00v\x00{'descr': '<f4', 'fortran_order': False, 'shape': (256, 256, 2), } "
<class 'tensorflow.python.framework.ops.EagerTensor'>
The shape seems to be correct but I don't know how to have the numpy.array