Move data normalization from a Python function to TensorFlow graph

demo-resnet.py

# This demonstrates how to move normalization of data from a Python function to inside a
# TensorFlow graph, demo'ing using resnet50.preprocess_input.
#
# For some background, see https://docs.edgeimpulse.com/docs/edge-impulse-studio/learning-blocks/adding-custom-learning-blocks
# For questions, see https://forum.edgeimpulse.com
#
# Tested with TensorFlow 2.7

import tensorflow as tf
import numpy as np
from tensorflow import keras
from tensorflow.keras import layers

# RGB array (0..1), like what you'll get from Edge Impulse
arr = np.array([
    [ 1.0, 0.9, 0.8 ],
    [ 0.9, 0.8, 0.7 ],
    [ 0.8, 0.7, 0.6 ],
    [ 0.2, 0.3, 0.4 ],
    [ 0.1, 0.2, 0.3 ],
    [ 0.0, 0.1, 0.2 ],
])

# Reshape into 3D
INPUT_SHAPE = (1, 6, 3)
arr = arr.reshape(INPUT_SHAPE)

# mean = [103.939, 116.779, 123.68]
# std = None

# Resnet50 expects the input in 0..255, so multiply here
input_for_resnet_preprocess = (arr * 255).copy()
# then we pass it through `preprocess_input` which normalizes the data
processed = tf.keras.applications.resnet50.preprocess_input(input_for_resnet_preprocess)
# this is what our data should look like.
print('processed according to resnet50.preprocess_input\n', processed)

# The resnet50.preprocess_input does two things:
# 1. RGB => BGR conversion
# 2. Scaling using: mean = [103.939, 116.779, 123.68], std = None (found this by looking in the Keras source code)

# The easiest way to do the scaling (at least that I know of) is to make a new matrix with the same
# shape as the input and fill in the mean values. We'll then subtract these two matrixes inside the graph.
scale_matrix = np.ones(INPUT_SHAPE)
scale_matrix[:,:,0] = 103.939
scale_matrix[:,:,1] = 116.779
scale_matrix[:,:,2] = 123.68

# Now construct the graph
input = keras.Input(shape=INPUT_SHAPE)
x = input
# RGB => BGR conversion using a lambda. This actually converts to a StridedSlice op in the TFLite graph
x = layers.Lambda(function=lambda x: x[..., -1::-1])(x)
# Rescale from 0..255
x = layers.Rescaling(scale=255)(x)
# Normalization using the scale matrix above
x = layers.Subtract()([ x, scale_matrix ])
output = x

# Compile the model (we don't need to train it, there's no trainable parameters here, just math)
model = keras.Model(inputs=input, outputs=output)
model.compile()

# print('summary', model.summary())

# test out the model. The predict function expects a batch so reshape to 4D
to_pred = arr.reshape((1,) + INPUT_SHAPE)
# predict
res = model.predict(to_pred)

# this matches the output from resnet50.preprocess_input above
print('predict result (should match resnet50.preprocess_input\n', res)

# Expected output:
# processed according to resnet50.preprocess_input
#  [[[ 100.061  112.721  131.32 ]
#   [  74.561   87.221  105.82 ]
#   [  49.061   61.721   80.32 ]
#   [  -1.939  -40.279  -72.68 ]
#   [ -27.439  -65.779  -98.18 ]
#   [ -52.939  -91.279 -123.68 ]]]
# predict result (should match resnet50.preprocess_input
#  [[[[ 100.061     112.721     131.32    ]
#    [  74.561      87.221     105.82    ]
#    [  49.060997   61.721      80.32    ]
#    [  -1.939003  -40.279     -72.68    ]
#    [ -27.439003  -65.779     -98.18    ]
#    [ -52.939003  -91.279    -123.68    ]]]]