Speed up loops over numpy arrays that modify the array

Question

I have numpy arrays with the shape (X, Y, 4), where X and Y are large. I want to do the following operation on axis 2 of the array (i.e., A[x, y] for each x and y): for each 1-dimensional vector A[x, y], if each of the first 3 elements is within some distance d to the corresponding element in constant vector C, then set the 4th element to 0. This is how I do it with Python loops:

for i in range(X): for j in range(Y): if np.all(np.abs(A[i, j, 0:3] - C) <= d): A[i, j, 3] = 0

These loops are very slow, obviously. I suspect there's a faster, ufunc-based way of doing this operation, but for the life of me can't figure out how to do it. Does anyone have any suggestions?

(Note on application: this is to add alpha channel to an image so that colors close to C become transparent.)

Hi, this might be offtopic but I would like to ask a question. Which shape of images is advisable to use? Is it (X, Y, 3) or (3, X, Y)? — mathfux
– mathfux, Commented Oct 8, 2020 at 23:41

Quang Hoang · Accepted Answer · 2020-10-08 23:32:23Z

2

Try:

mask = (np.abs(A[:,:,:3]-C)<=d).all(-1) A[mask,3] = 0

answered Oct 8, 2020 at 23:32

Quang Hoang

151k11 gold badges64 silver badges86 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Mark Wexler Over a year ago

Thanks a lot! Exactly the kind of thing I was looking for. For a 120x120 array of unsigned bytes, your solution takes 0.42 ms, while mine takes 77 ms. Good reason to study numpy indexing more closely!

Mark Wexler Over a year ago

Be careful if you have unsigned data! I do, since my data is RGB values between 0 and 255, so my data is uint8. With unsigned data, when you do the A[:,:,:3]-C subtraction in Quang's code, for instance when you do 60-61 you wind up with 255, rather than the expected -1. You need to change the data to a signed type after doing the subtraction. For exampe if your data is uint8, the first line in Quang's code should be changed to mask = (np.abs((A[:,:,:3]-C).astype(np.int8))<=d).all(-1).

Collectives™ on Stack Overflow

Speed up loops over numpy arrays that modify the array

1 Answer 1

2 Comments

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Related