Labeling components in a binary tensor

Hi Everyone, I have a binary image which is the output of a segmentation net. As part of the post processing I have to find all the connected components in this image, which I do with the openCV connectedComponents function. However I was wondering if there is a pytorch equivalent so I can do this operation while the tensor is still on the GPU. Does anyone have or know of something equivalent or similar?


1 Like

Hi David!

I don’t know one way or the other whether there is a pytorch / gpu
connected-components function. (I sort of doubt it.) But, in any
event, a connected-components algorithm would be unlikely to
benefit much from a gpu’s pipelined architecture.

I suppose that if after constructing the connected components
you needed to move the result back to the gpu, avoiding the
gpu-cpu-gpu memory transfers could be a benefit of running
connected components on the gpu.


K. Frank

Hi K, thanks for your reply. I haven’t yet found anything similar in pytorch so perhaps you’re right. Will keep looking though.


A GPU connected component implementation

I think this can easily be taken in.