How to pass y_true as a dict to a custom loss function unchanged?

I need to implement simple OCR model using tf.keras.*. But:

• Blank class is not zero, like tf expects, but (num_classes - 1) instead.
• Width of input images is not known beforehand (it is different for different batches).

I want to utilize tf.nn.ctc_loss which has some nice argument: blank_index. So I made a simple wrapper to compute CTC loss:

    class CTCLossWrapper(tf.keras.losses.Loss):
        def __init__(self, blank_class: int, reduction: str = tf.keras.losses.Reduction.AUTO, name: str = 'ctc_loss'):
            super().__init__(reduction=reduction, name=name)
            self.blank_class = blank_class
        
        def call(self, y_true, y_pred):
            output = y_true['output']
            targets, target_lenghts = output['targets'], output['target_lengths']
            y_pred = tf.math.log(tf.transpose(y_pred, perm=[1, 0, 2]) + K.epsilon())
            max_input_len = K.cast(K.shape(y_pred)[1], dtype='int32')
            input_lengths = tf.ones((K.shape(y_pred)[0]), dtype='int32') * max_input_len
            return tf.nn.ctc_loss(
                labels=targets,
                logits=y_pred,
                label_length=target_lenghts,
                logit_length=input_lengths,
                blank_index=self.blank_class
            )

I also wrote a simple generator function which yields training samples:

def generator(dataset, batch_size: int, shuffle=False):
    indexes = np.arange(len(dataset))
    while True:
        if shuffle:
            indexes = np.random.permutation(indexes)
        for i in range(0, len(dataset), batch_size):
            # Get next batch
            batch = dataset[indexes[i:i+batch_size]]
            images, image_widths = batch['images'], batch['image_widths']
            targets, target_lengths = batch['targets'], batch['target_lengths']
            # Re-arrange dimensions (B, H, W, C) -> (B, W, H, C)
            # Important Note: width=W and height=H are swapped from typical Keras convention
            # because width is the time dimension when it gets fed into the RNN
            images = np.transpose(images, axes=(0, 2, 1, 3)).astype(np.float32) / 255.0
            # Change zero target length to 1 due to invalid implementation of ctc_batch_cost in keras
            target_lengths[target_lengths == 0] = 1
            # Add singleton dimension
            # image_widths = image_widths[:, np.newaxis]
            # target_lengths = target_lengths[:, np.newaxis]
            # Construct output value
            outputs = {
                'images': images,  # (batch_size, max_image_width, 32, 1)
                'image_widths': image_widths,  # (batch_size,)
                'targets': targets,  # (batch_size, max_target_len)
                'target_lengths': target_lengths,  # (batch_size,)
            }
            yield images, dict(output=outputs)

As you may see, generator outputs not just (x, y_true) but 4 values:

• input images,
• input image widths,
• target sequences,
• lengths for each target sequence.

This is so because tf.nn.ctc_loss also requires at least 4 arguments to work.

My plan was to pass input images as x and a dictionary of all 4 values as y_true.

Then of course I compile the model using my CTCLossWrapper and blank_class:

    model.compile(
        optimizer=Adam(),
        loss=CTCLossWrapper(blank_class=blank_class),
    )

After that I can start training by:

    model.fit(
        x=generator(train_dataset, batch_size=batch_size, shuffle=True),
        steps_per_epoch=int(len(train_dataset) // batch_size),
        epochs=200
    )

The problem is that when my CTCLossWrapper is invoked it does not get dict() as y_true. It gets only one of tensors from it.

How is it possible to avoid or turn off tensorflow preprocessing and get y_true values in the same form as they were supplied from dataset?