Форма (Нет, Нет) наземных меток при использовании настраиваемого генератора данных Keras и функции потерь

Я пытаюсь реализовать свою собственную модель обнаружения, которая пытается найти объекты в изображении в оттенках серого по его координатам, для этого я создал настраиваемый набор данных, определяемый настраиваемым генератором данных:

class DataGenerator(tf.compat.v2.keras.utils.Sequence):

    def __init__(self, X_data , y_data, batch_size, shuffle = True):
        self.batch_size = batch_size
        self.X_data = X_data
        self.labels = y_data
        self.y_data = y_data
        self.shuffle = shuffle
        self.n = 0
        self.dim = (480, 848)
        self.list_IDs = np.arange(len(self.X_data))
        self.on_epoch_end()

    def __next__(self):
        # Get one batch of data
        data = self.__getitem__(self.n)
        # Batch index
        self.n += 1

        # If we have processed the entire dataset then
        if self.n >= self.__len__():
            self.on_epoch_end
            self.n = 0

        return data
    def __len__(self):
        # Return the number of batches of the dataset
        return math.ceil(len(self.indexes)/self.batch_size)

    def __getitem__(self, index):
        # Generate indexes of the batch
        indexes = self.indexes[index*self.batch_size:
            (index+1)*self.batch_size]
        # Find list of IDs
        list_IDs_temp = [self.list_IDs[k] for k in indexes]

        X = self._generate_x(list_IDs_temp)
        y = self._generate_y(list_IDs_temp)
        return X, y


    def on_epoch_end(self):

        self.indexes = np.arange(len(self.X_data))

        if self.shuffle: 
            np.random.shuffle(self.indexes)

    def _generate_x(self, list_IDs_temp):

        X = np.empty((self.batch_size, *self.dim))
        for i, ID in enumerate(list_IDs_temp):
            X[i,] = cv2.imread(self.X_data[ID],0)
            X = (X/255).astype('float32') # Normalize data  
        return X[:,:,:, np.newaxis]

    def _generate_y(self, list_IDs_temp):

        y = np.empty((self.batch_size, 2))
        for i, ID in enumerate(list_IDs_temp):
            y[i] = self.y_data[ID]
        return y

При вызове он дает следующий результат:

val_generator  = DataGenerator(x_test, y_test, batch_size=4, shuffle=False)
images, labels = next(val_generator)
print(labels.shape)
>>>> (4, 2)

Это то, что вы ожидаете от пакета размером 4 с координатами x и y на изображении.

Модель выглядит следующим образом:

model = Sequential()
model.add(Conv2D(32, kernel_size=(3, 3),
                activation='relu',
                batch_input_shape=(4, 480, 848, 1)))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Dropout(0.25))
model.add(Conv2D(64, (3, 3), activation='relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Dropout(0.25))
model.add(Flatten())
model.add(Dense(64, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(15, activation="relu"))

model.compile(loss=simple_loss, optimizer=keras.optimizers.Adam(lr=0.001))

Вызов модели:

steps_per_epoch = len(train_generator)
validation_steps = len(val_generator)

model.fit(train_generator,
        steps_per_epoch=steps_per_epoch,
        epochs=10,
        validation_data=val_generator,
        validation_steps=validation_steps)

Моя обычная, но простая функция потерь, чтобы проверить, работает ли модель, дала несколько ошибок:

def simple_loss(yTrue, yPred):
    probs_logs, coords = yPred[:,:5], yPred[:,5:]
    coords_2d = tf.reshape(coords, [4, 10]) # 4 batch size and 10 coords flatten out
    _abs = tf.math.abs(yTrue, coords_2d )
    return tf.sqrt(_abs)

Сначала я подумал, что yPred содержит ошибки, но в результате получился следующий тензор: Tensor ("simple_loss / Reshape: 0", shape = (4, 10), dtype = float32)

Затем я посмотрел на yTrue и обнаружил, что форма была (None, None): Tensor ("IteratorGetNext: 1", shape = (None, None), dtype = float32)

Я думаю, что с моим генератором что-то не так, я понятия не имею, что именно, поэтому мне было интересно, может ли кто-нибудь из вас мне помочь?

Спасибо