[Python] Extracting 2x2 pixel groups from an image into (height, width, 4px, 3rgb) - slice...

I am trying to extract 2x2 pixel groups into a matrix that has a shape (height, width, 4 - one of each pixels, 3 - pixels RGB.

The pixel groups have a 1 PX padding. A sample image might look like this:



What I expect is a matrix with the shape of (2-height, 2-width, 4-pixels per group, 3-RGB).

And this is my code:

def read_color_choices_img(image_mat: tf.Tensor, dimensions_hw: tuple[int, int]) -> NDArray[np.floating]:
"""
Expects an image with squares of 2x2 pixels, each square representing a color choice.

Pixels are expected to be RGB and normalized to [0, 1] range.

Pixel groups are both horizontally and vertically separated by 1 pixel.

Therefore the format should be choices A,B,C,D and separator pixel S
```
A B S A B S ...
C D S C D S ...
S S S S S S
. . . . . . . . .
```

The right most rows and columns must be the separators. Batch dimension is not expected,
input shape should be (height, width, 3 - RGB).
"""
if image_mat.shape[0] != image_mat.shape[1]:
raise ValueError("Input image must be square (height == width), was: "+str(image_mat.shape))
ends_with_padding = (image_mat.shape[0] - dimensions_hw[0] * 3) % 3 == 0
# Either the image ends with padding and then it's a multiple of (2+1)
# Or it does not end with padding and is N*(2+1)-1
color_choice_height = (image_mat.shape[0] // 3 if ends_with_padding else (image_mat.shape[0]+1) // 3)

# This will change once different w vs h is supported
color_choice_dimensions = (color_choice_height, color_choice_height)

print("Input image shape:", image_mat.shape)
# Extract 2x2 patches, skipping 1-pixel separators
patches = tf.image.extract_patches(
images=tf.expand_dims(image_mat, axis=0), # Add batch dimension
sizes=[1, 2, 2, 1], # 2x2 patches
strides=[1, 3, 3, 1], # Move by 3 pixels (2 pixels + 1 separator)
rates=[1, 1, 1, 1],
padding="VALID" # Only fully contained patches - possible padding is dropped
)

# Print all choices as hex table
patches = tf.squeeze(patches, axis=0) # Remove batch dimension
final_patches_shape = (*color_choice_dimensions, 4, 3)
print("Patches shape:", patches.shape)
print("Expected choice shape:", final_patches_shape) # 4 choices per pixel, RGB
# reshape to (height, width, 4, 3) where 4 is the number of choices per pixel
patches = tf.reshape(patches, final_patches_shape) # 4 choices per pixel, RGB

# print choices as hex table to console
for off_h in range(dimensions_hw[0]):
row = []
for off_w in range(dimensions_hw[1]):
choice_group = patches[off_h, off_w]
hex_choices = [f"#{int(c[0]*255):02x}{int(c[1]*255):02x}{int(c[2]*255):02x}" for c in choice_group.numpy()]
row.append(" [".join(hex_choices)+"]")
print(" ".join(row))
return patches


And the code prints:

Input image shape: (6, 6, 3)
Patches shape: (2, 2, 12)
Expected choice shape: (2, 2, 4, 3)
W tensorflow/core/framework/op_kernel.cc:1857] OP_REQUIRES failed at strided_slice_op.cc:117 : INVALID_ARGUMENT: slice index 2 of dimension 1 out of bounds.
I tensorflow/core/framework/local_rendezvous.cc:407] Local rendezvous is aborting with status: INVALID_ARGUMENT: slice index 2 of dimension 1 out of bounds.


I don't see why - my target dimension seems correct to me - 4*3 is 12. What am I doing wrong?

In this case I want my final matrix to be:

[
[ [orange, green, darkblue, pink], [darkred, lightblue, purple, limegreen] ],
[ [cyan, darkgreen, red, blue], [yellow, brown, pink, blue] ],
]

Continue reading...