How to do preprocessing inside a Keras model for new inference-only end-to-end model

Question

With Tensorflow 2.3, I've trained a tf.keras model with EfficientNet and now I want to export an inference-only end-to-end model that will include the preprocessing layers (essentially decoding an image encoded as a base64 string and potentially processing the normalization). My code to define the trained model and the inference model:

imageSize = (224,224,3)
inputs = layers.Input(shape=imageSize)
eff_net_model = EfficientNetB0(input_tensor=inputs)
eff_net_model.compile(...)
eff_net_model.fit(...)
# training is finished now wrap the model with preprocessing for inference model
model = b64_image_model_wrapper(eff_net_model,imageSize)

My wrapper fct is the following:

def b64_image_model_wrapper(model,imageSize,method=tf.image.ResizeMethod.BILINEAR, 
                        mean=0.0,std=1.0,input_name="b64_image"):

  def preprocess_and_decode(img_str, new_shape=imageSize):
    img = tf.io.decode_base64(img_str)
    img = tf.image.decode_jpeg(img, channels=3)
    img = tf.image.resize(img, new_shape, method=method)
    return img

  input64 = tf.keras.layers.Input(shape=(None,), dtype="string", name=input_name)
  output_tensor = tf.keras.layers.Lambda(
    lambda img: tf.map_fn(lambda im: preprocess_and_decode(im[0]), img, fn_output_signature=tf.float32,swap_memory=True))(input64)
  x = (output_tensor - mean) / std
  x = model(x)
  new_model = tf.keras.Model(input64, x, name=model.name)
  new_model.output_names = model.output_names
  return new_model

Later if I want to predict with this new model (the one that takes as input a string encoded image) I do:

model.predict(np.array(["hello"])) # just for testing purposes

I get the error:

Exception has occurred: InvalidArgumentError
Expected image (JPEG, PNG, or GIF), got unknown format starting with '\261\330_\201\250\354\205\327\340i\327'
 [[{{node EfficientNet/lambda/map/while/body/_1/EfficientNet/lambda/map/while/DecodeJpeg}}]]

Likewise if I save this new model (tf.saved_model.save(model)) and try to use it with Tensorflow Serving I get the error:

Error while reading resource variable block6c_bn/moving_mean from Container: localhost. This could mean that the variable was uninitialized. Not found: Container localhost does not exist. (Could not find resource: localhost/block6c_bn/moving_mean)\n\t [[{{node functional_1/EfficientNet/block6c_bn/FusedBatchNormV3/ReadVariableOp}}]]

I don't know what it means exactly but it seems that some weights are not initialized? Am I doing the model wrapping the right way?

Answer 1

The reason you get error when you call the model.predict(np.array(["hello"])) is because:

1. The input format is not base64 string
2. your input after tf.io.decode_base64 is not .jpg file

Following code show you how to test the model using .jpg file:

with open("./homersimpson.0.0.jpg", "rb") as imageFile:
    base64_bytes = base64.urlsafe_b64encode(imageFile.read())
base64_strings = base64_bytes.decode('ascii')
img = tf.io.decode_base64(base64_strings)
prediction = model.predict(np.array([base64_strings]))

Also I using model.save('./test_save_model') to save the entire model which don't have any problem