I'm trying to inference a TFLite model that was originally built in PyTorch. I have been following along the lines of the PyTorch implementation and have to preprocess images along the RGB channels. I found the closest TensorFlow equivalent of transforms.Normalize() to be tf.image.per_image_standardization() (documentation). Although this is a pretty good match, tf.image.per_image_standardization() does this by taking mean and std across the channels and applies it to them. Here's their full implementation from here
def per_image_standardization(image):
"""Linearly scales `image` to have zero mean and unit norm.
This op computes `(x - mean) / adjusted_stddev`, where `mean` is the average
of all values in image, and
`adjusted_stddev = max(stddev, 1.0/sqrt(image.NumElements()))`.
`stddev` is the standard deviation of all values in `image`. It is capped
away from zero to protect against division by 0 when handling uniform images.
Args:
image: 3-D tensor of shape `[height, width, channels]`.
Returns:
The standardized image with same shape as `image`.
Raises:
ValueError: if the shape of 'image' is incompatible with this function.
"""
image = ops.convert_to_tensor(image, name='image')
_Check3DImage(image, require_static=False)
num_pixels = math_ops.reduce_prod(array_ops.shape(image))
image = math_ops.cast(image, dtype=dtypes.float32)
image_mean = math_ops.reduce_mean(image)
variance = (math_ops.reduce_mean(math_ops.square(image)) -
math_ops.square(image_mean))
variance = gen_nn_ops.relu(variance)
stddev = math_ops.sqrt(variance)
# Apply a minimum normalization that protects us against uniform images.
min_stddev = math_ops.rsqrt(math_ops.cast(num_pixels, dtypes.float32))
pixel_value_scale = math_ops.maximum(stddev, min_stddev)
pixel_value_offset = image_mean
image = math_ops.subtract(image, pixel_value_offset)
image = math_ops.div(image, pixel_value_scale)
return image
whereas PyTorch's transforms.Normalize() allows us to mention the mean and std to be applied across each channel like below.
# transformation
pose_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
What would be a way to get this functionality in TensorFlow 2.x?
Edit: I created a quick botch that seems to solve this by defining a function as such:
def normalize_image(image, mean, std):
for channel in range(3):
image[:,:,channel] = (image[:,:,channel] - mean[channel])/std[channel]
return image
I'm not sure how efficient this is but seems to get the job done. I still have to convert the output to a tensor before inputing to the model.
