CODEWITHSUNDEEP
pythonnumpystruct
Akihiko
Akihiko

Reputation: 372

Python: struct.pack and struct.unpack for NumPy array

I am exploring NumPy version of struct.pack and struct.unpack. My particular problem is converting an RGB value (r, g, b) into a binary.

For example:

f = lambda r, g, b: struct.unpack('I', struct.pack('BBBB', r, g, b, 255))[0]
f(155, 156, 148)
#--> 4287929499

When I apply this conversion to an image (2-d array of RGB), it becomes very slow.

For example:

import numpy as np
import struct

img_rgb = np.random.randint(0, 256, (480, 640, 3))
%%time
np.apply_along_axis(lambda rgb: struct.unpack('I', struct.pack('BBBB', rgb[0], rgb[1], rgb[2], 255))[0], 2, img_rgb)
#--> Wall time: 5.23 s

My question is that is there NumPy version of struct.pack and struct.unpack? Or, how can I make this code faster with NumPy?

Upvotes: 1

Views: 2737

Answers (1)

Mad Physicist
Mad Physicist

Reputation: 114440

You don't need pack and unpack operations in numpy because it uses raw binary arrays instead of individually wrapping each value in a separate object. In general, the operations you use are reshape, transpose, ndarray.view, and sometimes ndarray.astype.

In your particular case, you have and (M, N, 3) image, which you want to add a channel of 255 to and convert to np.uint8:

img_rgb = np.random.randint(0, 256, size=(480, 640, 3))
img_rgba = np.concatenate((img_rgb.astype(np.uint8), np.full((*img_rgb.shape[:-1], 1), 255, dtype=np.uint8)), axis=-1)

By setting the last dimension to size 4, you ensure that you can now view the data directly as np.uint32:

result = img_rgba.view(np.uint32)

The result will be of shape (M, N, 1). You can remove the extra dimension with np.squeeze:

result = np.squeeze(img_rgba.view(np.uint32))

This doesn't really buy you much: it removes the convenient 3-channel representation without altering the memory layout. Dumping either img_rgba or result to a raw binary file would produce identical results.

In general, functions like apply_along_axis are generally code-smell in numpy. They are no better than running an explicit for-loop. Your goal should always be to vectorize your code to perform operations on all elements at once.

Upvotes: 3

Related Questions