CSharp
CSharp

Reputation: 1476

How to prepare image data stored in a zip file for training in Tensorflow 2?

I have a large set of images I need to prepare for Deep Learning with a Convolutional Neural Network using Tensorflow 2 / Keras. A batch of 61 Images are stored in a zip file with their respective 'masks' (which are simply the segmented version of the image). So for example, zip file Batch-0-of-163.zip contains:

'image-1.png', 'mask-1.png', 'image-2.png', 'mask-2.png' ... 'image-61.png', 'mask-61.png'

Is there a way to create a tensorflow.data.Dataset in Tensorflow 2, that will generate the image and mask data when needed by the GPU for input to my CNN? I want to use a Dataset so I can take advantage of the batching/prefetching functionality provided.

Upvotes: 2

Views: 6614

Answers (1)

Gabriele
Gabriele

Reputation: 949

The way I would solve the problem consists in the following steps:

  • create a Dataset object containing the path to each file
  • map a python function on each element of the Dataset to unzip, load the data and remove the unzipped the folder (that I assume you don't need anymore to be unzipped)
  • go back to tensorflow code to further processing

Here is an example of how the code should look like:

 from scipy import misc
 import os

 # ----------------------------
 # Parsing function with standard python:

 def zip_data_parser(zip_fname):
     os.system('unzip {0}'.format(zip_fname)) # unzip
     folder_name = zip_fname.rsplit('.zip')[0]

     # load data:
     x_stack = []
     y_stack = []
     for i in range(n_images):
         x_stack.append(misc.imread(folder_name + '/image-{0}.png'.format(i)))
         y_stack.append(misc.imread(folder_name + '/mask-{0}.png'.format(i)))
     x = np.array(x_stack)
     y = np.array(y_stack)

     os.system('rm -rf {0}'.format(folder_name)) # remove unzipped folder
     return x, y 

 # ----------------------------
 # Dataset pipeline:

 all_zip_paths = ['file1.zip', 'file2.zip', 'file3.zip'] # list of paths for each zip file
 train_data = tf.constant(all_zip_paths)
 train_data = tf.data.Dataset.from_tensor_slices(train_data)

 train_data = train_data.map(
            lambda filename: tf.py_func(  # Parse the record into tensors
                zip_data_parser,
                [filename],
                [tf.float32, tf.float32]), num_parallel_calls=num_threads)

 # un-batch first, then batch the data again to have dimension [batch_size, N, M, C]
 train_data = train_data.apply(tf.data.experimental.unbatch())
 train_data = train_data.batch(b_size, drop_remainder=True)

Of course, you may need to cast x and y to np.float32 before returning them from zip_data_parser to the Dataset object. I also assumed that the masks are already one-hot encoded in my example.

Upvotes: 2

Related Questions