Split segmentation mask into annotated areas

Hello there,

is there an easy way in PyTorch to split an index mask tensor into multiple masks, each containing only one contiguous labeled area (e.g. [num_areas, H, W])?

Tensor [1, 512, 512] (Index Tensor containing 0 or 1) to a Tensor [6, 512, 512]. 6 is the amount of contiguous areas, which is not static. The colored areas are equal to 1 in my current Tensor:


Greetings, Mukeyii

Hi Kemal!

What you are looking for are the connected components (of an image,
rather than of a graph).

I do not believe that pytorch supplies a built-in connected-components
algorithm, but several pytorch connected-components implementations
appear to be available on the internet. (I haven’t tried any of them, so I
have no opinion as to which of them might work.)

You could also write your own pytorch implementation – doing so would
be a task, but relatively straightforward.

Or you could pre-process your segmentation masks with a non-pytorch
image-processing library that does have a connected-components


K. Frank

1 Like