Is there way to embed non-tf functions to a tf.Keras model graph as SavedModel Signature?

I want to add preprocessing functions and methods to the model graph as a SavedModel signature.

example:

# suppose we have a keras model
# ...
# defining the function I want to add to the model graph
@tf.function
def process(model, img_path):
    # do some preprocessing using different libs. and modules...
    outputs = {"preds": model.predict(preprocessed_img)}
    return outputs
# saving the model with a custom signature
tf.saved_model.save(new_model, dst_path, 
                    signatures={"process": process})

or we can use tf.Module here. However, the problem is I can not embed custom functions into the saved model graph.

Is there any way to do that?

I think you slightly misunderstand the purpose of save_model method in Tensorflow.

As per the documentation the intent is to have a method which serialises the model's graph so that it can be loaded with load_model afterwards.

The model returned by load_model is a class of tf.Module with all it's methods and attributes. Instead you want to serialise the prediction pipeline.

To be honest, I'm not aware of a good way to do that, however what you can do is to use a different method for serialisation of your preprocessing parameters, for example pickle or a different one, provided by the framework you use and write a class on top of that, which would do the following:

class MyModel:
    def __init__(self, model_path, preprocessing_path):
        self.model = load_model(model_path)
        self.preprocessing = load_preprocessing(preprocessing_path)
    def predict(self, img_path):
        return self.model.predict(self.preprocessing(img_path))