using Segform model with keras-hub error. The structure of `inputs` doesn't match the expected structure

Question

I am trying a semantic segmentation task using Segformer model with pretrained model 'mit_b3_cityscapes_1024' .

encoder = keras_hub.models.MiTBackbone.from_preset(
        "mit_b3_cityscapes_1024",  image_shape=(768,768,3)
    )
backbone = keras_hub.models.SegFormerBackbone(
        image_encoder=encoder,
        projection_filters=256,
    )
model = keras_hub.models.SegFormerImageSegmenter(
        backbone=backbone,
        num_classes=12
    )

I read the data using ImageDataGeneraor():

    img_datagen_instance=ImageDataGenerator()
    msk_datagen_instance=ImageDataGenerator()
   
        img_data_gen=img_datagen_instance.flow_from_directory(directory=img_path,target_size=(768,768), batch_size=batch_size,seed=32,
                                                          class_mode=None,color_mode='rgb')

    msk_data_gen=msk_datagen_instance.flow_from_directory(directory=msk_path,target_size=(768,768),batch_size=batch_size,seed=32,
                                                          class_mode=None,color_mode='grayscale')
  

    img_msk_generated=zip(img_data_gen,msk_data_gen)
    
    return img_msk_generated 
train_generator=img_msk_data_generator(img_path=train_img_path,msk_path=train_msk_path,class_num=class_n,batch_size=batch_size)
val_generator=img_msk_data_generator(img_path=val_img_path,msk_path=val_msk_path,class_num=class_n,batch_size=batch_size)

by which when I fit the model I get this erro: history=model.fit(processed_train_generator,steps_per_epoch=step_train,epochs=epoch_num, ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/m/mkeshavarzi/miniconda3/envs/tens-keras-hub/lib/python3.12/site-packages/keras_hub/src/utils/pipeline_model.py", line 163, in fit x = _convert_inputs_to_dataset(x, y, sample_weight, batch_size) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/m/mkeshavarzi/miniconda3/envs/tens-keras-hub/lib/python3.12/site-packages/keras_hub/src/utils/pipeline_model.py", line 69, in _convert_inputs_to_dataset raise e File "/home/m/mkeshavarzi/miniconda3/envs/tens-keras-hub/lib/python3.12/site-packages/keras_hub/src/utils/pipeline_model.py", line 57, in _convert_inputs_to_dataset ds = tf.data.Dataset.from_tensor_slices(inputs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/m/mkeshavarzi/miniconda3/envs/tens-keras-hub/lib/python3.12/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 827, in from_tensor_slices return from_tensor_slices_op._from_tensor_slices(tensors, name) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/m/mkeshavarzi/miniconda3/envs/tens-keras-hub/lib/python3.12/site-packages/tensorflow/python/data/ops/from_tensor_slices_op.py", line 25, in _from_tensor_slices return TensorSliceDataset(tensors, name=name) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/m/mkeshavarzi/miniconda3/envs/tens-keras-hub/lib/python3.12/site-packages/tensorflow/python/data/ops/from_tensor_slices_op.py", line 33, in init element = structure.normalize_element(element) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/m/mkeshavarzi/miniconda3/envs/tens-keras-hub/lib/python3.12/site-packages/tensorflow/python/data/util/structure.py", line 110, in normalize_element ops.convert_to_tensor(t, name="component%d" % i)) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/m/mkeshavarzi/miniconda3/envs/tens-keras-hub/lib/python3.12/site-packages/tensorflow/python/profiler/trace.py", line 183, in wrapped return func(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^ File "/home/m/mkeshavarzi/miniconda3/envs/tens-keras-hub/lib/python3.12/site-packages/tensorflow/python/framework/ops.py", line 732, in convert_to_tensor return tensor_conversion_registry.convert( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/m/mkeshavarzi/miniconda3/envs/tens-keras-hub/lib/python3.12/site-packages/tensorflow/python/framework/tensor_conversion_registry.py", line 234, in convert ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/m/mkeshavarzi/miniconda3/envs/tens-keras-hub/lib/python3.12/site-packages/tensorflow/python/framework/constant_tensor_conversion.py", line 29, in _constant_tensor_conversion_function return constant_op.constant(v, dtype=dtype, name=name) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/m/mkeshavarzi/miniconda3/envs/tens-keras-hub/lib/python3.12/site-packages/tensorflow/python/ops/weak_tensor_ops.py", line 142, in wrapper return op(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/home/m/mkeshavarzi/miniconda3/envs/tens-keras-hub/lib/python3.12/site-packages/tensorflow/python/framework/constant_op.py", line 276, in constant return _constant_impl(value, dtype, shape, name, verify_shape=False, ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/m/mkeshavarzi/miniconda3/envs/tens-keras-hub/lib/python3.12/site-packages/tensorflow/python/framework/constant_op.py", line 289, in _constant_impl return _constant_eager_impl(ctx, value, dtype, shape, verify_shape) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/m/mkeshavarzi/miniconda3/envs/tens-keras-hub/lib/python3.12/site-packages/tensorflow/python/framework/constant_op.py", line 301, in _constant_eager_impl t = convert_to_eager_tensor(value, ctx, dtype) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/m/mkeshavarzi/miniconda3/envs/tens-keras-hub/lib/python3.12/site-packages/tensorflow/python/framework/constant_op.py", line 108, in convert_to_eager_tensor return ops.EagerTensor(value, ctx.device_name, dtype) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ ValueError: Attempt to convert a value (<generator object preprocessing at 0x7f59f0168f40>) with an unsupported type (<class 'generator'>) to a Tensor.

Then I tried to convert the dataset from datagenerator to tf.dataset format with this function :

def preprocessing(img_msk_gen):    
   
    input=img_msk_gen
    def gen():
        for images,masks in input:
            
            masks=to_categorical(masks,num_classes=class_n)
            yield (images, masks)
    dataset=tf.data.Dataset.from_generator(gen
                                           ,output_signature=(tf.TensorSpec(shape=(batch_size,768,768,3), dtype=tf.float32),
                                                      tf.TensorSpec(shape=(batch_size,768,768,class_n), dtype=tf.float32), ))
           
    return dataset

when I print the tensors :

print('processed_train_generator',processed_train_generator)

processed_val_generator=preprocessing(val_generator)
print('processed_val_generator',processed_val_generator)

I get this : processed_train_generator <_FlatMapDataset element_spec=(TensorSpec(shape=(8, 768, 768, 3), dtype=tf.float32, name=None), TensorSpec(shape=(8, 768, 768, 12), dtype=tf.float32, name=None))>

this time running model.fit will give this warning and the process get stock like without being killed or anything. Expected: ['keras_tensor_1'] Received: inputs=Tensor(shape=(8, 768, 768, 3)) warnings.warn(msg)