2

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.)

Mark Wexler
  • 81
  • 5

1 Answers1

2

Try:

mask = (np.abs(A[:,:,:3]-C)<=d).all(-1)
A[mask,3] = 0
Quang Hoang
  • 146,074
  • 10
  • 56
  • 74
  • 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 Oct 09 '20 at 00:10
  • 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)`. – Mark Wexler Oct 09 '20 at 09:12