Как создать serve_input_fn в Tensorflow 2.0 для предварительной обработки изображений?

Question

Я использую Tensorflow 2.0 и могу обучить CNN для классификации изображений 3-канальных изображений.Я выполняю предварительную обработку изображений в конвейере ввода данных (показано ниже) и хотел бы включить функциональность предварительной обработки в саму обслуживаемую модель.Моя модель обслуживается с помощью контейнера TF Serving Docker и API Predict.

Конвейер ввода данных для обучения основан на документации по https://www.tensorflow.org/alpha/tutorials/load_data/images.

Моя функция предварительной обработки изображения конвейера - 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

Со списками путей к изображениям конвейер предварительно выбирает и выполняет предварительную обработку изображений с использованием tf-функций в 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)

Примеры Tensorflow, которые я нашел, используют serve_input_fn () и используют tf.placeholder , который, похоже, больше не существует в Tensorflow 2.0.

Пример serve_input_fn в Tensorflow 2.0 показан на https://www.tensorflow.org/alpha/guide/saved_model. Поскольку я использую Predict API, похоже, мне нужно что-то похожее на:

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)

В идеале обслуживаемая модель должна принимать 4-мерный тензор 3-канальных выборок изображения любого размера и выполнять начальную предварительную обработку изображения на них (декодировать изображение, нормализовать, преобразовывать в HSV и изменять размер) перед классификацией.

Как я могу создать serve_input_fn в Tensorflow 2.0 с функцией предварительной обработки, аналогичной моей load_and_preprocess_path функции?