Keras - understanding ImageDataGenerator dimensions

Question

I'm trying to implement custom object detection by taking a trained YOLOv2 model in Keras, removing the last layer and retraining it with new data. I'm confused about how to feed the data to Keras, though. I have annotated a bunch of pictures with bounding boxes using the YOLO annotation, and put them in two separate folders (images where the .jpgs reside and annots where the .txt annotations are).

I also removed the last layer from the model and added a custom one (I'm trying to predict bounding boxes for 2 classes).

I'm trying to pass my data with an ImageDataGenerator, as my dataset is quite small.

I have the following input objects:

np.shape(train_images) # this contains RGB data from 79 pictures 
(79, 1, 608, 608, 3)
np.shape(train_y)
(79,)

I'm trying to pass these to the ImageDataGenerator, but I get an error:

train_datagen = ImageDataGenerator(
      rotation_range=20,
      width_shift_range=0.2,
      height_shift_range=0.2,
      horizontal_flip=True,
      fill_mode='nearest')

train_generator = train_datagen.flow(
        train_images,
        train_y)

ValueError: `x` (images tensor) and `y` (labels) should have the same length. Found: x.shape = (1, 608, 608, 3), y.shape = (79,)

I don't understand what the problem is. Somehow the first dimension of my images data is completely gone and thus does not match... What's wrong with it?

Answer 1

I think your train_images array should have shape (79, 698, 608, 3). The generator works through each of the first dimensions of those arrays, so is passing a batch of 4d numpy arrays, instead of a batch of 3d numpy arrays.

You can try seeing if that helps, using numpy.squeeze(), like this:

In [1]: import numpy as np                                                      

In [2]: a = np.random.randint(0, 10, (2, 1, 10, 10, 3))                         

In [3]: np.squeeze(a, axis=1).shape                                             
Out[3]: (2, 10, 10, 3)

So be sure to set the axis=1 argument. Then just pass the updated array as you were doing.