How do I remove loop for numpy subtraction of 2d and 3d arrays?

Viewed 127

Is it possible to speed up the following code by removing the for-loop?

all_t = np.array([])    
for color in color_list:
    t = np.subtract(image, color)
    all_t = np.append(all_t, t)
all_t = all_t.reshape(len(color_list), image.shape[0], image.shape[1])

How I'm using the code:

  • color_list is a numpy array of shape (z,3). It consists of z colors with with RGB values
  • image is a numpy array of shape (x,y,3). It is an RGB image with dimension (x,y)
  • all_t is a numpy array of shape (z,x,y) consisting of the z images with subtracted color values

Edit:

  • all_t is of shape (z,x,y,3)
    all_t = all_t.reshape(len(color_list), image.shape[0], image.shape[1], 3)
  • invalid color values are allowed for this operation as of now
3 Answers

Adding

Another approach, but also by comprehension list, in which a mask is made for each color in the list, then subtract is performed once

color_map = np.array([[[color_list[i]] * image.shape[0]] * image.shape[1] for i in range(len(color_list))])
images = [image] * len(color_list)
all_t = np.subtract(images,color_map)

you can measure compare the performance with your code in question, the results will depend on the length of color_list and image size.

Previouse Assumption Temporary, If we use a list comprehension version of your code:

all_t = [np.subtract(image,color) for color in color_list] 
all_t = np.array(all_t)
all_t = all_t.reshape(len(color_list), image.shape[0], image.shape[1]

Here is a solution for removing the for loop from your code. I am going to work with your method of subtracting two pixels (which can result in invalid pixel values) Please see bellow for alternative solutions

You have two numpy arrays: image of shape (x, y, 3) color_list of shape (z, 3)

Your goal is to subtract each color in color_list from the image yielding z different images that can be stored in a array of shape (z, x, y, 3)

# First we repeat the pixel values according to the number of colors we would like to subtract ie. `len(color_list) = z`
# The shape of rep_images will be(x, y, z, 3)
rep_images = np.repeat(image[:,:,None], len(color_list), axis=2)

# We can no subtract `color_list` from `rep_images`. 
# This results in the first color being subtracted from the pixels along `z` at index 0, second color form pixels along `z` index 1 etc.
sub_image = rep_images - color_list

# Now we can move the `z` axis to the front in order to have `z` images, of `x` by `y` pixels of 3 RGB colors
# `result` now has a shape of (z, x, y, 3) and contains the z subtracted images
result = np.moveaxis(sub_image, 2, 0)

Note on subtracting colors

When subtracting colors in a trivial way by taking the difference in the individual R, G and B channels, you can easily reach invalid colors.

For example, subtracting (255, 0, 0) from (0, 0, 0) would result in (-255, 0 ,0)

The naive approach would be to set negative values to zero with:

result[result < 0] = 0

However other methods can be employed. For example:

  • Subtracting Gray scale images
  • Subtracting value (essentially equivalent to subtraction of gray scale images but somewhat preserving color)
  • Subtracting color saturation

Most of these methods can be employed by converting the image and colors to HSL and using the same algorithm. Here are a few useful resources:

In [189]: color_list = np.array([[1,2,3],[4,5,6]])
In [190]: color_list.shape
Out[190]: (2, 3)
In [191]: image = np.ones((4,5,3),int)
In [192]: all_t = np.array([])
     ...: for color in color_list:
     ...:     t = np.subtract(image, color)
     ...:     all_t = np.append(all_t, t)
     ...: 
In [193]: all_t.shape
Out[193]: (120,)

all_t is a 1d array. np.append is clear that without axis it flattens/ravels all inputs.

Your reshape is wrong, ignoring the size 3 dimension:

In [194]: all_t.reshape(len(color_list), image.shape[0], image.shape[1])
Traceback (most recent call last):
  File "<ipython-input-194-c246f37a52d0>", line 1, in <module>
    all_t.reshape(len(color_list), image.shape[0], image.shape[1])
ValueError: cannot reshape array of size 120 into shape (2,4,5)

A working reshape:

In [195]: x=all_t.reshape(len(color_list), image.shape[0], image.shape[1],3)
In [196]: x.shape
Out[196]: (2, 4, 5, 3)

If you are going to build an an array iteratively, use list append. It is faster, and easier to use correctly:

In [197]: all_t = []
     ...: for color in color_list:
     ...:     t = np.subtract(image, color)
     ...:     all_t.append(t)
     ...: 
In [198]: len(all_t)
Out[198]: 2
In [199]: y = np.array(all_t)
In [200]: y.shape
Out[200]: (2, 4, 5, 3)
In [201]: np.allclose(x,y)
Out[201]: True

But there's no need to iterate on the first dimension of color_list:

In [202]: z = image - color_list[:,None,None,:]
In [203]: z.shape
Out[203]: (2, 4, 5, 3)
In [204]: np.allclose(x,z)
Out[204]: True

By broadcasting rules this is

(4,5,3) array - a (2,1,1,3)
(1,4,5,3)
(2,4,5,3)

The other answer, https://stackoverflow.com/a/65550381/901925, constructs:

In [205]: np.repeat(image[:,:,None], len(color_list), axis=2).shape
Out[205]: (4, 5, 2, 3)
In [206]: color_list.shape
Out[206]: (2, 3)

which by broadcasting also works (206 expands to (1,1,2,3)); he then changes the dimension order.

Related