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:
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