scripts/libraries: Update normalization to handle 3D/4D image tensors. (#2258)

* scripts/libraries: Add sanity checks to image Normalization.
This commit is contained in:
Kwabena W Agyeman 2024-07-16 12:48:57 -07:00 committed by GitHub
parent b94e8528f1
commit 072db9647e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -35,8 +35,15 @@ class Normalization:
buffer, shape, dtype = args buffer, shape, dtype = args
# Create an image using the input tensor as buffer. # Create an image using the input tensor as buffer.
pixfmt = image.GRAYSCALE if shape[3] == 1 else image.RGB565 if len(shape) != 4:
img = image.Image(shape[2], shape[1], pixfmt, buffer=buffer) 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. # Copy and scale (if needed) the input image to the input buffer.
hints = image.BILINEAR | image.CENTER | image.SCALE_ASPECT_EXPAND | image.BLACK_BACKGROUND hints = image.BILINEAR | image.CENTER | image.SCALE_ASPECT_EXPAND | image.BLACK_BACKGROUND