Element check and modification in image segmentation numpy array while maintaining shape-CodePudding

I am designing a real-time system to filter out certain regions of an image based on their class ID from a semantic segmentation network. From the network, I have a numpy array (13,16,1) of the class IDs (numbers 0 through 20) that correspond to regions of the image (i.e. (1,1) is the top leftmost region, etc). What I would like to end up with is something that I can reference later in the program to decide whether to filter out a given region of an image based on class ID.

I want to filter out a "blacklist" of classes (ex. class IDs 0, 2, 13), and have later analysis only consider regions of the image where those classes are not detected.

The solution that I had envisioned originally was creating a copy of the class ID array by performing a check on each element of the class ID array, checking if the class ID matches any in the "blacklist", and changing that value to 0 or 1 based on that check. I haven't been able to get any of my go-to solutions to work, and I'm slightly worried that an element-by-element check will be too slow for real time performance, especially if I increase the image resolution.

For reference, I'm starting with an array like:

class_ID_array =      1   2   3
                      4   5   6
                      7   8   9

And I want to perform an operation along these lines that produces an output in the format of desired_array:

desired_array = class_ID_array.copy()
blacklist = [2, 4, 5]
for x in desired_array:
    if x is in blacklist:
       x = 0
    else:
       x = 1

desired_array =      1   0   1
                     0   0   1
                     1   1   1

What would be an efficient way to do this?

CodePudding user response：

def operation( a, blacklist ):
    s = a.shape
    l = np.ndarray.flatten( a )
    ol = [ e not in blacklist for e in l ]
    return np.reshape( ol, s )