CODEWITHSUNDEEP
pythonnumpyimage-processingvectorization
burgerhex
burgerhex

Reputation: 1048

Vectorizing code of looping through NumPy 3D array

I'm currently writing a program that processes an image and returns the same image except some pixels that are identified as a certain feature (found via truth_table) are blackened. It works, but since the images are around 5000 by 5000 pixels, it takes about 1 whole minute just to go through this part:

image_DN_no_12s = np.copy(image_DN)
for y in range(y_size):
    for x in range(x_size):
        if truth_table[y, x] == 12:
            image_DN_no_12s[y, x, :] = np.zeros((1, 1, 4))

(The 3rd dimension in np.zeros() is 4 because I am reading red, green, blue, and infrared channels.)

I would like to vectorize this code so that it will run faster. I have tried using np.where() and boolean indexing, but I can't seem to find a way to make it run faster.

Thanks in advance for the help!

Upvotes: 0

Views: 44

Answers (1)

Mateen Ulhaq
Mateen Ulhaq

Reputation: 27201

If img and truth_table are arrays where truth_table is of broadcastable shape:

img[truth_table == 12] = 0

You don't need to explicitly create another axis of 4 channels since broadcasting handles it for you.

Upvotes: 1

Related Questions