Filtering out zero matrices from tensor in numpy

I'm given a tensor of segmentation masks for ML algorithm of size (30, 256, 256). The problem is, that some of those entries are zero matrices and I have to filter them out. Now I'm using a naive for loop based technique with np.array_equal function to manually filter them out.

Is there a way to do this more efficiently in NumPy way using some fancy indexing?

Upvotes: 1

Views: 47

Answers (2)

fuglede
fuglede

Reputation: 18201

If a is your array, you could do

a[np.any(a != 0, (1, 2))]

Upvotes: 1

Eduardo Soares
Eduardo Soares

Reputation: 1000

Just iterate over your matrices and use the function count_non_zero() to efficiently check if a matrix is a zero-matrix.

import numpy as np

for matrix in tensor:
  if np.count_nonzero(matrix) != 0:
     #keep in your tensor
  else:
     #remove from your tensor

Upvotes: 1

Related Questions