skimage.segmentation.watershed implementation?

[ma-sadeghi]

Is your feature request related to a problem? Please describe.
It's be great if cucim implements skimage's watershed algorithm. watershed is really useful for segmentation and it'd be great if we had a GPU implementation of it.

Describe the solution you'd like
Implement skimage's watershed algorithm.

Describe alternatives you've considered
None

Additional context
None

[gigony]

Thanks @ma-sadeghi for the issue! We had a discussion on this issue before and we found that there are many ways to implement GPU-accelerated watershed algorithms and implementing scikit-image's watershed algorithm directly with CuPy is not feasible or tricky to implement.

We would keep follow up this issue and please let us know if you have any idea/knowledge on GPU-accelerated watershed algorithm.

Thanks!

/cc @grlee77

[jakirkham]

Greg likely has more thoughts here 🙂

Though one thing worth mentioning is we do have a random_walker implementation. Admittedly not exactly what you are looking for, but may still be useful

[grlee77]

Looking at the scikit-image watershed demo, the seed points for the watershed are determined via scipy.distance_transform_edt. I did start a CuPy-based implementation for the distance transform based on the code in https://github.com/orzzzjq/Parallel-Banding-Algorithm-plus. I started with the 3D variant and have that working, but wanted to make things a bit more general in terms of not requiring all dimensions of the array to have the same size.

I had not started working on the watershed itself. Some reading turned up some GPU-based implementations in the literature, but I haven't looked into specifics of which might be best to implement. My initial impression of the scikit-image algorithm was that the priority queue data structure approach that is used there would likely not adapt well to the GPU.

[jakirkham]

This reminds me, @JoOkuma pointed me to this watershed implementation, which might also be of interest here 🙂

[JoOkuma]

Hey @jakirkham.

I agree with @gigony, an implementation that produces results equivalent to skimage's watershed is tricky and might not even be that fast.

@alanpeixinho might have the GPU implementation from the article above. However, from my understanding, it does not produce the same results as the original watershed.

[grlee77]

Two other GPU implementations:

1.) 2D implementation based on cellular automata in recent NPP

2.) A different GPU-based algorithm with corresponding citations at watershed-cuda, but I do not see any license info in that repository. Perhaps @louismullie can indicate if it could be used under an Apache 2.0 or similar license?

[louismullie]

The implementation on my profile can be used under the Apache 2.0 license.

[grlee77]

The implementation on my profile can be used under the Apache 2.0 license.

Great, thanks @louismullie

[rijobro]

Any progress on this? @louismullie your implementation only works for 2d, correct?

[jakirkham]

Based on recent discussion, another option might be to build on NPP's FloodFill

[ma-sadeghi]

@jakirkham Would it work for 3D as well?

[jakirkham]

No it is 2D only

[ma-sadeghi]

@gigony Could you explain a little bit why porting scikit-image's watershed using cupy is not feasible? Thanks!

[jakirkham]

Actually Greg's probably the best one to explain that given he's familiar with the scikit-image codebase.

That said, we talked about this again today (hence the comment on this issue). AIUI the primary challenge with the algorithm used in scikit-image is it leverages a priority queue to go through and explore each neighborhood. This isn't particularly friendly for parallelism or working on the GPU. The other algorithms mentioned above may be better for parallelism, but could render results that are qualitatively different from what scikit-image does.

[JoOkuma]

I would like to add that the solution to the watershed is not unique in the tie-zones (objects' boundaries where pixels have the same value). The tie-breaking strategy is usually an implementation detail and will probably yield different results in the GPU implementation.

[jakirkham]

Just noting this is something MONAI had expressed interest in using here.

cc @drbeh