将计算值元素添加到多维numpy数组的快速方法

Question

I've got a numpy array 'image' that is a two dimensional array where each element has two components. I want to convert this to another two dimensional array where each element has three components. The first two and a third one calculated from the first two, like so:

for x in range(0, width):
    for y in range(0, height):
        horizontal, vertical = image[y, x]

        annotated_image[y, x] = (horizontal, vertical, int(abs(horizontal) > 1.0 or abs(vertical) > 1.0))

This loop works as expected, but is very slow when compared to other numpy functions. For a medium-sized image this takes an unacceptable 30 seconds.

Is there a different way to do the same calculation but faster? The original image array does not have to be preserved.

Answer 1

You could just separate the components of the image and work with multiple images instead:

image_component1 = image[:, :, 0]
image_component2 = image[:, :, 1]

result = (np.abs(image_component1) > 1.) | (np.abs(image_component2) > 1.)

If you for some reason need the layout you specified you could as well construct another three dimensional image:

result = np.empty([image.shape[0], image.shape[1], 3], dtype=image.dtype)

result[:, :, 0] = image[:, :, 0]
result[:, :, 1] = image[:, :, 1]
result[:, :, 2] = (np.abs(image[:, :, 0]) > 1.) | (np.abs(image[:, :, 1]) > 1.)