I’m trying to convert a .h5 model, which is a simplified version of Unet for image segmentation, into a .tflite model. Does anyone know the correct way to do this? So far, the .h5 model seems to make good predictions, but when I try the .tflite model, the predictions are literally an image completely in purple color.

This is what I’ve done:

converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()


with open('unet_model_prueba_sin_int8.tflite', 'wb') as f:
    f.write(tflite_model)

And then to test de .tflite model:

interpreter = tf.lite.Interpreter(model_path="unet_model_prueba_sin_int8.tflite")
interpreter.allocate_tensors()


input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()


def load_and_preprocess_image(image_path, target_size):
    image = Image.open(image_path).resize(target_size)
    image_array = np.array(image, dtype=np.float32) / 255.0  # Normalizar la imagen
    image_array = np.expand_dims(image_array, axis=0)
    return image_array


def predict_with_tflite(image_array):
    input_tensor = image_array.astype(np.float32)
    interpreter.set_tensor(input_details[0]['index'], input_tensor)
    interpreter.invoke()
    output = interpreter.get_tensor(output_details[0]['index'])
    return output


image_path = "image_path/image.jpg"


image_array = load_and_preprocess_image(image_path, target_size=(96, 96))


prediction = predict_with_tflite(image_array)



def display_segmentation_mask(pred):
    """Display the segmentation mask."""
    mask = np.argmax(pred, axis=-1) 
    plt.imshow(mask[0], cmap='viridis')  
    plt.axis('off')  
    plt.show()

display_segmentation_mask(prediction)

Any guidance or suggestions would be greatly appreciated!