Numpy - fastest way to filter volumes from 3D array

user_15481857 0
Viewed 223

I have relatively large 3D arrays where I'm attempting to filter out specific labeled sets of pixels (operationally defined here as volumes). Labeled volumes have 26-connectivity.

I start with:

  • A labeled array (based on 26-connectivity), where each connected component is assigned a value
  • A array of label values that need to be removed from the labeled/binary array
  • An array with single seed indices from each of the labeled regions that are to be deleted.

A mock version of these starting variables can be constructed below.

import numpy as np

labeled = np.zeros((600,600,600)) # Empty labeled array
filter_coords = np.zeros((1000, 3),dtype=int) # Empty coords array
filter_ids = np.arange(3,3001, step=3)# Values associated with each label coord
id = 0
for i in range(1,3001):
    # Random bottom and top coords to label and create a "volume".
    b = np.random.randint(0,600,3,dtype=int) 
    t = np.random.randint(0,50,3,dtype=int)
    top = b + t

    #Fill the labeled array with 'mock' labeled volumes
    labeled[b[0]:top[0],
            b[1]:top[1],
            b[2]:top[2]] = i

    # Find coords from 1000 of these volumes
    if i % 3 == 0:
        filter_coords[id] = b
        id += 1

In my current approach, I find a single seed coordinate from a single element of the labeled volume that needs to be filtered. I then start expanding a box around that region and convert all labeled elements into zeros until the expanding box meets no more labeled elements of interest.

# Remove labeled regions from volume
def filter(label, filter_ids, fcoords):
    for i in range(fcoords.shape[0]):
        filtered = False
        b = 25 # Arbitrary starting box size
        while filtered == False:
            min = fcoords[i] - b
            min[min < 0] = 0 # Make sure lower bound isn't negative
            mask = label[min[0]:fcoords[i,0]+b+1,
                         min[1]:fcoords[i,1]+b+1,
                         min[2]:fcoords[i,2]+b+1]
            if filter_ids[i] in mask:
                # If the id is present in our sub-set, replace it w/ 0
                mask[mask == filter_ids[i]] = 0
                # Then replace the original region with the mask
                label[min[0]:fcoords[i,0]+b+1,
                      min[1]:fcoords[i,1]+b+1,
                      min[2]:fcoords[i,2]+b+1] = mask
                b += 5 # Increase box size
            else:
                filtered = True
                
    return label

labeled = filter(labeled, filter_ids, filter_coords)

Could anyone provide suggestions on whether my approach to this problem is robust? If not, are there alternative methods that I could implement?

Thanks.

numpy arrays python
2 Answers

Here is an alternative approach based on scipy.ndimage.labeled_comprehension. I get nearly the same results as with your code, but not exactly. By the way, I think the last line in the definition of your filter function should read return label as opposed to return labeled.

import numpy as np
from scipy.ndimage import labeled_comprehension
np.random.seed(0)

# used to get aronud special treatment of zero label by scipy
labeled1 = labeled + 1

# get mapping from old labels to new labels
remove_id_map = np.arange(3001 + 1, dtype=int)
remove_id_map[filter_ids + 1] = 1

# get boolean mask of lienar indices to remove
remove_index_mask = np.zeros(labeled.size, dtype=bool)
remove_index_mask[np.ravel_multi_index(filter_coords.T, dims=labeled.shape)] = True

# replaces labels with ones if conditions are met
def f(arr, ind):
  remove_ind = remove_index_mask[ind].any()
  if remove_ind:
    return remove_id_map[arr]
  else:
    return arr

# apply labeled comprehension
result = labeled_comprehension(labeled1, labeled1, None, f, 
                               out_dtype=labeled.dtype, default=0, 
                               pass_positions=True)

result = (result - 1).reshape(labeled.shape)
user_45859630

Ok - I've found a solution that works quite well and was inspired by the alternative to np.isin provided here by norok2.

Essentially, the 3D labeled volume and a 1D np.array of values that should be removed from the labeled volume are passed into an njit function. The function starts by converting the 1D array into a set.

The function then scans across the volume and converts all elements of the volume that are in the set of ids to zero. I keep my volumes in 3D rather than flattening them because sometimes they are np.memmaps, and numba seems to want to pull the memmamps into memory following .ravel() or .flatten(). The speed seems similar between 3D vs flattened arrays that are in memory.

Here is my best solution to this problem so far.

@njit(parallel=True, cache=True)
def filter_segments(labeled, remove_ids):
    """Remove specific ids from a labeled np.3darray
        labeled: a 3D np.array or np.memmap that is labeled with specific ids
        remove_ids: a 1D np.array of values that will be removed from the volume
        returns the labeled volume with all elements of the remove_ids converted to zero
    """
    remove_ids = set(remove_ids)
    for z in prange(labeled.shape[0]):
        for y in range(labeled.shape[1]):
            for x in range(labeled.shape[2]):
                p = labeled[z,y,x]
                if p and p in remove_ids:
                    labeled[z,y,x] = 0
    return labeled

Hopefully this can be useful to someone else in the future!

user_154818570
Related