Add support to keep more than 1 component in KeepLargestConnected

[Nic-Ma]

Is your feature request related to a problem? Please describe.
Feature request from Ahmed: is there a way to keep the largest 2 components? I faced this when segmenting the kidneys and wanted to keep both kidneys in the segmentation.

[a-parida12]

@Nic-Ma are you refering to extending this https://docs.monai.io/en/stable/transforms.html#keeplargestconnectedcomponent to accept generic n largest components?

If no one is working on it may be I could put a PR on this?

[Nic-Ma]

Hi @rijobro ,

I remember we had some discussion about this ticket before, but I don't remember clearly the plan now..
Could you please help share some info?

Thanks in advance.

[rijobro]

Hm, an easy way would be to start with a label image of all zeros. We loop across as many n times as desired and set the label to 1 if those voxels are part of the biggest island. Then we set that island to 0 in the input image and iterate again. This may be slow, I'm not sure if measure.label is a heavy operation or not.

Not looked into it much, but perhaps scipy.ndimage.label might be able to do what we want more easily?

[rijobro]

We can certainly make use of the return_num argument to measure.label to determine the number objects.

[rijobro]

from skimage.measure import label
import numpy as np

def get_largest_connected_component_mask(img, connectivity=None, num_components=1, return_num=False):
    # features will be an image, 0 for background and then each different
    # feature will have its own index.
    features, num_features = label(img, connectivity=connectivity, return_num=True)

    # if num features less than max desired, nothing to do.
    if num_features <= num_components:
        out = img.astype(bool)
        num_components = num_features
    else:
        # get number voxels per feature (bincount). np.argsort[::-1] to get indices
        # of largest components. Convert to list for ease
        features_to_keep = list(np.argsort(np.bincount(features.flat))[::-1])
        # remove 0 (background)
        features_to_keep.remove(0)
        # only keep the first n non-background indices
        features_to_keep = features_to_keep[:num_components]
        # generate labelfield. True if in list of features to keep
        out = np.isin(features, features_to_keep)
    return (out, num_features, num_components) if return_num else out

if __name__ == "__main__":
    img = np.array([[[0,0,1,1,0,0,1],
                    [0,0,0,1,0,0,1],
                    [1,1,0,0,1,0,1],
                    [0,0,0,1,0,0,1]]])

    for num_components in range(5):
        out, num_input_features, num_output_features = get_largest_connected_component_mask(img, num_components=num_components, return_num=True)
        print(f"Should have <={num_components} component(s). Input had {num_input_features}, output has {num_output_features}.")
        print(out)

Output:

Should have <=0 component(s). Input had 3, output has 0.
[[[False False False False False False False]
  [False False False False False False False]
  [False False False False False False False]
  [False False False False False False False]]]
Should have <=1 component(s). Input had 3, output has 1.
[[[False False  True  True False False False]
  [False False False  True False False False]
  [False False False False  True False False]
  [False False False  True False False False]]]
Should have <=2 component(s). Input had 3, output has 2.
[[[False False  True  True False False  True]
  [False False False  True False False  True]
  [False False False False  True False  True]
  [False False False  True False False  True]]]
Should have <=3 component(s). Input had 3, output has 3.
[[[False False  True  True False False  True]
  [False False False  True False False  True]
  [ True  True False False  True False  True]
  [False False False  True False False  True]]]
Should have <=4 component(s). Input had 3, output has 3.
[[[False False  True  True False False  True]
  [False False False  True False False  True]
  [ True  True False False  True False  True]
  [False False False  True False False  True]]]

[rijobro]

The above is backward compatible, for brevity I removed numpy/tensor conversion and the cusim implementation, but these would have to be taken care of.

[a-parida12]

@rijobro Thanks for sharing the implementation and the ideas on how to get started!

I looked into monai.transforms.utils.get_largest_connected_component_mask @ https://github.com/Project-MONAI/MONAI/blob/dev/monai/transforms/utils.py#L954 it looks like this function needs to be extended to accept num_components and return_num as parameters. The cusim and measure.labels are already implemented there.
Does that sound good?

Another question, do we really need the return_num? as the functions are called get_largest_connected_component_mask so as a user will expect only a binary mask to be returned.

[rijobro]

Yes, I meant that my implementation needs to be added to the original. I'll get it done tomorrow and get a PR done.