How to remove matrices in a numpy array of matrices?-CodePudding

I have a numpy array arr_seg_labs which has the following shape: (1735, 128, 128). It contains pixel masks between 1 and 10 and also contains zeros and 255 (background). I want to remove those (128, 128) matrices which not contain the given category identifier (9) and to keep those which contain at least one 9. I made a mask (horse_mask) for this, but I don't know how can I continue this thread to filter this numpy array

CAT_IDX_HORSE = 9
horse_mask = arr_seg_labs == CAT_IDX_HORSE

CodePudding user response：

IIUC you can use masks and indexing as:

CAT_IDX_HORSE = 9
mask = (a == CAT_IDX_HORSE ).sum((1, 2))
result = a[mask != 0]