CODEWITHSUNDEEP
tensorflow2.0
Donald Bourque
Donald Bourque

Reputation: 41

How to create a serving_input_fn in Tensorflow 2.0 for image preprocessing?

I am using Tensorflow 2.0 and am able to train a CNN for image classification of 3-channel images. I perform image preprocessing within the data input pipeline (shown below) and would like to include the preprocessing functionality in the served model itself. My model is served with a TF Serving Docker container and the Predict API.

The data input pipeline for training is based on the documentation at https://www.tensorflow.org/alpha/tutorials/load_data/images.

My pipeline image preprocessing function is load_and_preprocess_from_path_label:

def load_and_preprocess_path(image_path):

    # Load image
    image = tf.io.read_file(image_path)
    image = tf.image.decode_png(image)

    # Normalize to [0,1] range
    image /= 255

    # Convert to HSV and Resize
    image = tf.image.rgb_to_hsv(image)
    image = tf.image.resize(image, [HEIGHT, WIDTH])

    return image

def load_and_preprocess_from_path_label(image_path, label):

    return load_and_preprocess_path(image_path), label

With lists of image paths, the pipeline prefetches and performs image preprocessing using tf functions within load_and_preprocess_from_path_label:

all_image_paths, all_image_labels = parse_labeled_image_paths()
x_train, x_test, y_train, y_test = sklearn.model_selection.train_test_split(all_image_paths, all_image_labels, test_size=0.2)

# Create a TensorFlow Dataset of training images and labels
ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
image_label_ds = ds.map(load_and_preprocess_from_path_label)

BATCH_SIZE = 32
IMAGE_COUNT = len(all_image_paths)

ds = image_label_ds.apply(tf.data.experimental.shuffle_and_repeat(buffer_size=IMAGE_COUNT))
ds = ds.batch(BATCH_SIZE)
ds = ds.prefetch(buffer_size=AUTOTUNE)

# Create image pipeline for model
image_batch, label_batch = next(iter(ds))
feature_map_batch = model(image_batch)

# Train model
model.fit(ds, epochs=5)

Previous Tensorflow examples I've found use serving_input_fn(), and utilized tf.placeholder which seems to no longer exist in Tensorflow 2.0.

An example for serving_input_fn in Tensorflow 2.0 is shown on https://www.tensorflow.org/alpha/guide/saved_model. Since I am using the Predict API, it looks like I would need something similar to:

serving_input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn(...)

# Save the model with the serving preprocessing function
model.export_saved_model(MODEL_PATH, serving_input_fn)

Ideally, the served model would accept a 4D Tensor of 3-channel image samples of any size and would perform the initial image preprocessing on them (decode image, normalize, convert to HSV, and resize) before classifying.

How can I create a serving_input_fn in Tensorflow 2.0 with a preprocessing function similar to my load_and_preprocess_path function?

Upvotes: 4

Views: 994

Answers (2)

Yuri
Yuri

Reputation: 241

Expanding and simplifying @harry-salmon answer. For me the following worked:

def save_model_with_serving_signature(model, model_path):
    @tf.function(input_signature=[tf.TensorSpec(shape=[None, ], dtype=tf.string)])
    def serve_load_and_preprocess_path(image_paths):
        return model(tf.map_fn(load_and_preprocess_path, image_paths, dtype=tf.float32))

    tf.saved_model.save(
        model,
        model_path,
        signatures=serve_load_and_preprocess_path
    )

Note: dtype=tf.float32 in map function was important and didn't work without it. I found solution here. Also I simplified the concrete function work by simply adding a decorator (see this for details).

Upvotes: 2

Harry Salmon
Harry Salmon

Reputation: 189

I faced a similar issue when upgrading. It appears that the way to achieve this in Tensorflow 2 is to provide a function which the saved model can use to make the predictions, something like: 

def serve_load_and_preprocess_path(image_paths: tf.Tensor[tf.string]):
    # loaded images may need converting to the tensor shape needed for the model        
    loaded_images = tf.map_fn(load_and_preprocess_path, image_paths, dtype=tf.float32)
    predictions = model(loaded_images)
    return predictions


serve_load_and_preprocess_path = tf.function(serve_load_and_preprocess_path)
serve_load_and_preprocess_path = serve_load_and_preprocess_path.get_concrete_function(
    image_paths=tf.TensorSpec([None,], dtype=tf.string))

tf.saved_model.save(
    model,
    MODEL_PATH,
    signatures=serve_load_and_preprocess_path
)

# check the models give the same output
loaded = tf.saved_model.load(MODEL_PATH)
loaded_model_predictions = loaded.serve_load_and_preprocess_path(...)
np.testing.assert_allclose(trained_model_predictions, loaded_model_predictions, atol=1e-6)

Upvotes: 4

Related Questions