From 072db9647e3fb69959fda428b82a9264a8937711 Mon Sep 17 00:00:00 2001 From: Kwabena W Agyeman Date: Tue, 16 Jul 2024 12:48:57 -0700 Subject: [PATCH] scripts/libraries: Update normalization to handle 3D/4D image tensors. (#2258) * scripts/libraries: Add sanity checks to image Normalization. --- scripts/libraries/ml/ml/preprocessing.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/scripts/libraries/ml/ml/preprocessing.py b/scripts/libraries/ml/ml/preprocessing.py index 4619b7844..7844040d6 100644 --- a/scripts/libraries/ml/ml/preprocessing.py +++ b/scripts/libraries/ml/ml/preprocessing.py @@ -35,8 +35,15 @@ class Normalization: buffer, shape, dtype = args # Create an image using the input tensor as buffer. - pixfmt = image.GRAYSCALE if shape[3] == 1 else image.RGB565 - img = image.Image(shape[2], shape[1], pixfmt, buffer=buffer) + if len(shape) != 4: + raise ValueError("Expected input tensor with shape: (1, H, W, C)") + b, h, w, c = shape + if b != 1: + raise ValueError("Expected batches to be 1") + if c != 1 and c != 3: + raise ValueError("Expected channels to be 1 or 3") + pixfmt = image.GRAYSCALE if c == 1 else image.RGB565 + img = image.Image(w, h, pixfmt, buffer=buffer) # Copy and scale (if needed) the input image to the input buffer. hints = image.BILINEAR | image.CENTER | image.SCALE_ASPECT_EXPAND | image.BLACK_BACKGROUND