Connected components labeling in JAX for metrics

[pabloduque0]

Im looking into adding connected components labeling for one of my JAX programs. This is used for precision/recall for object label segmentation. Im fairly new to JAX so please correct me if anything of the following is not correct.

Bottom line is that I have not found a "jaxy" way of properly doing this (or calling something).

• jax2tf.call_tf would be one option to call tensorflow_addons.image.connected_components but since it is a custom op that would not allow me to jit this part (and since this sits in my metrics it would be a bummer to leave it out of the jitted part). Also this function only supports 4 way CCL.
• scipy.ndimage.measurements.label, skimage.measure.label or opencv equivalent would leave me in a similar in a similar position without being able to jit.
• Connected components labeling is not too complicated and I have a couple of self implementations but they make use of disjoint-sets or other data structure or approaches that make it not very friendly to JAX. I believe something like this is what generally is implemented in most places.

Any thoughts/ideas?

Im happy to submit something to jax.scipy.ndimage to work as the equivalent of scipy.ndimage.measurements.label if I come with something that is worthy (and if there is interest for it).

Thank you very much in advance!

[jakevdp]

Hi - sorry for the lack of response here. I don't know of any good answer to your question: I'm not aware of any JAX-compatible connected components routine, and because the algorithms for this are iterative in nature, I don't think they map well onto the XLA primitives that are available to JAX. That said, an implementation of scipy.ndimage.measurements.label would be a welcome addition to JAX, if a suitable approach were found.

[pabloduque0]

Hey Jake! Thanks for getting back to me, very much appreciated.

Let me think for a bit if I can come up with something that is somewhat decent but all the ideas that I got for now are not very efficient. Will get back to you if something interesting comes to mind.

Thanks again! :)