Skip to content

Instantly share code, notes, and snippets.

@janjongboom
Created Nov 17, 2022
Embed
What would you like to do?
Move data normalization from a Python function to TensorFlow graph
# 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 ]]]]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment