I would use bincount and threshold the result to make a lookup table:
import numpy as np
threshold = 121
size = np.bincount(labels.ravel())
keep_labels = size <= threshold
# Make sure the background is left as 0/False
keep_labels[0] = 0
filtered_labels = keep_labels[labels]
On the last above I index the array keep_labels
with the array labels
. This is called advanced indexing in numpy and it requires that labels
be an integer array. Numpy then uses the elements of labels
as indices to keep_labels
and produces an array the same shape as labels
.