简体   繁体   中英

How to find if numpy array contains vector along 3rd dimension?

I want to find if a 3D numpy array contains a specific 1D vector along the 3rd dimension. (I need to check if an image contains pixel(s) of a specific color.)

I need to return true if and only if any of the pixels match exactly with the target.

I've tried the following:

import numpy as np
target = np.array([255, 0, 0])
search_area = np.array([[[0,0,0],[1,1,1],[2,2,2]],
                        [[3,3,3],[4,4,4],[5,5,5]],
                        [[6,6,6],[7,7,7],[8,8,255]]])
contains_target = np.isin(target, search_area).all():  # Returns True

Which returns True since each element can be found individually somewhere within the entire array.

Next I tried:

target = np.array([255, 0, 0])
search_area = np.array([[[0,0,0],[1,1,1],[2,2,2]],
                        [[3,3,3],[4,4,4],[5,5,5]],
                        [[6,6,6],[7,7,7],[8,0,255]]])
contains_target = (target == search.all(2)).any()  # Returns True

This works better since it matches the elements of target for each pixel individually, but it still returns True when they aren't in order or in the right numbers.

Lastly, I tried:

def pixel_matches_target(self, pixel_to_match):
        return (target == pixel_to_match).all()

contains_target = np.apply_along_axis(self.pixel_matches_target, 2, search_area).any()

But it's too slow to be used (about 1 second per pass).

How can I find if a numpy array contains a vector along a specific axis?


EDIT:

I ended up circumventing the issue by converting the RGB images to binary masks using cv2.inRange() and checking if the resulting 2D array contains True values. This resulted in several orders of magnitude faster execution.

One somewhat decent possibility to solve your problem would be (if you can afford the extra temp-memory):

import numpy as np

target = np.array([255, 0, 0])
search_area = np.array([[[0,0,0],[1,1,1],[2,2,2]],
                        [[3,3,3],[4,4,4],[5,5,5]],
                        [[6,6,6],[7,7,7],[8,0,255]]])

# works for general N-D sub-arrays
adjusted_shape = search_area.reshape((-1, *target.shape))

contains_target = target.tolist() in adjusted_shape.tolist() # False

If your arrays are integers you may use numpy.array_equal() to check that the arrays match (if using float s see numpy.allclose() instead). Assuming the sub-arrays to match are always in the 3rd sub-row you may do the following:

if sum(np.array_equal(target,a) for a in arr[:,2]):
    # Contains the target!

Should the sub-array occur anywhere you can use:

sum(np.array_equal(target,item) for sublist in arr for item in sublist))

Note: This doesn't answer the generalized question, but is several orders of magnitude faster for the specific problem of finding if an image contains pixels of a specific color.

import cv2
import numpy

target_lower = np.array([250, 0, 0])
target_upper = np.array([255, 5, 5])
search_area = np.array([[[0,0,0],[1,1,1],[2,2,2]],
                        [[3,3,3],[4,4,4],[5,5,5]],
                        [[6,6,6],[7,7,7],[8,0,255]]])

mask = cv2.inRange(search_area, target_lower, target_upper)
mask = mask.astype(bool)

contains_target = (True in mask)

Additionally, it has the benefit of allowing for a bit of flexibility for the target color.

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM