w-k-jones
commented
1 year ago Had a look deeper into it and it turns out that modifying the scikit-image watershed function is a lot simpler than I expected. I've got it running with Semi-Lagrangian offsets, and a similar system should handle periodic and polar boundary conditions. @freemansw1 / @galexsky do you have any example datasets for periodic boundaries that I could use to test this?
Currently we use the watershed function from the scikit-image library (https://github.com/scikit-image/scikit-image/blob/v0.19.2/skimage/segmentation/_watershed.py#L95-L229) for segmentation. This is fast, because it uses a cython backend (https://github.com/scikit-image/scikit-image/blob/v0.19.2/skimage/segmentation/_watershed_cy.pyx) but inflexible for use cases such as periodic BCs and semi-Lagrangian frameworks. I've been putting some thought into whether we could add our own watershed routine to tobac to account for these cases without workarounds.
The core problem with the scikit-learn watershed is that it uses a fixed neighbourhood for all pixels, which means that we can't treat pixels at the boundary differently to those in the centre of the image. I've had a couple of thoughts about how we could go about implementing an improved version:
Python watershed implementation. I have this mostly complete, with semi-Lagrangian support and periodic BCs. But, it's about 100x slower than the scikit-image version. How does this compare to the current workaround for periodic BCs?
Pass a large array of per-pixel neighbourhoods to the scikit-image watershed. This should be relatively simple to implement as we can do most of the modification outside of cython, but would come with a memory hit due to large array sizes required (4 x field.size for 2d, 6x field.size for 3d and that's only considering square connectivity)
Handle boundary conditions within the cython function. This would be more challenging to implement, but have the best performance. I should look at how periodic BCs are handled in other scikit-image/ndimage functions that support them to see if there is an existing implementation we could use.
Primarily opening this issue just to get my thoughts written down to follow up on later, but if you have any thoughts (particularly regarding periodic BCs and comparison to the current implementation) I would be happy to hear them!