Numpy: Filter list of 2D indices by 2D mask-CodePudding

I am having a numpy array of 2D points data (x, y, other_properties), where x and y are integer pixel coordinates. Additionally, I am having a binary 2D segmentation mask mask. I would like to filter the list, to obtain only the points where the mask is one/true.

I thought about doing something like:

valid_indices = np.argwhere(mask_2D)

to then filter the data based on the valid indices and I would like to do that using numpy acceleration. Here is a toy example

# the data representing x and y coordinates
data = np.arange(10).reshape((5, 2))
data = np.concatenate((data, data), axis=0)
print(f"data: {data}")

# Let's say we already obtained the indices from the segmentation mask
valid_indices = [(6, 7), (0, 1)]
print(f"valid_indices: {valid_indices}")

filtered = []
for point in data:
    if tuple(point) in valid_indices:
        filtered.append(tuple(point))
filtered = np.array(filtered)
print(f"filtered: {filtered}")

Output:

data:
[[0 1]
 [2 3]
 [4 5]
 [6 7]
 [8 9]
 [0 1]
 [2 3]
 [4 5]
 [6 7]
 [8 9]]
valid_indices:
[(6, 7), (0, 1)]
filtered:
[[0 1]
 [6 7]
 [0 1]
 [6 7]]

Process finished with exit code 0

Is there a way to obtain the behavior above using numpy? The solution could also directly use the binary 2D segmentation mask. If not, do you have any suggestions how to speed up the process? Thank you!

CodePudding user response：

You can do it with broadcast compare,

data[(data == np.array(valid_indices)[:,None]).all(-1).any(0)]