teamtomo / membrain-seg

membrane segmentation in 3D for cryo-ET
Other
48 stars 12 forks source link

Feature suggestion: adjustable thresholding #24

Closed bbarad closed 9 months ago

bbarad commented 11 months ago

We sometimes want thinner segmentations, particularly for highly appressed membranes. Otherwise we get overlapped segmentations, even though they are clearly separated in the probability view. I think the best way to do this would be to make the threshold level user-selectable (alternatively, could perhaps implement a local thresholding routine based on neighborhood intensities)

Is this something you would be willing to have within the segmentation tooling? It looks like currently threshold is done with a hardcoded setting of zero, but I'd be happy to submit a pull request to convert this to a user-selectable threshold (in dataloading/data_utils.py) with default setting 0.0.

LorenzLamm commented 11 months ago

Hey Ben! Absolutely, it makes sense. When lowering the threshold, some falsely merged membranes may get disconnected again. But probably, it will also lead to other membrane segmentation areas getting worse. Still, probably it would be good to have the threshold as an optional user input to play around with. If you're willing to make PR with this, it would be amazing! Otherwise, I'll also find time to do it the next couple of days -- it wouldn't be much effort to implement it :)

Alternatively, do you think some kind of skeletonization of the membranes may also be useful? The attached "soft" skeletonization operates directly on the score maps. I think it could be useful if you would like to have thinner segmentations, but probably will not perfectly solve the issue of wrongly merged membranes.

skeletonization_snap
bbarad commented 11 months ago

I'll put together a PR today or tomorrow for the user-editable thresholding.

With that said, it looks like the soft skeletonization you attached would work even better - I'd honestly be onboard with just making that the default.

alisterburt commented 11 months ago

Wow! The skelwtonisation looks awesome, how is it done?

+1 on making it default!

LorenzLamm commented 11 months ago

Cool, thanks Ben! :)

@alister The skeletonization above is done by subsequent erosion and dilation of the membranes, and tracking the differences:

  1. You start with thick membranes and erode them, so they are just thinner. So if you dilate it again, it's pretty much the same as the initial membrane. No differences to track.
  2. Then you continue with the eroded, i.e. thinner, membrane
  3. If your membrane is very thin (ideally 1 voxel thick) and you erode it, it's gone. So if you dilate it, it's still gone. The difference between these two gives the skeleton.

The nice thing with this is that I can approximate the dilation and erosion directly on the scores using min- and max pooling and thus use it as a loss function. The implementation is pretty similar to the Centerline Dice implementation (https://openaccess.thecvf.com/content/CVPR2021/papers/Shit_clDice_-_A_Novel_Topology-Preserving_Loss_Function_for_Tubular_Structure_CVPR_2021_paper.pdf)

I am trying to get something similar like a "Surfaceness Dice" score. That would be a better measure of the continuity of membranes and, with the differentiable min- and max-pooling, could also be used directly in the model training.

I'm still working out what's the best way to do the skeletonization (need to check whether some parts are faulty), but I agree it would be cool to have it directly as an output, so users can directly control the thickness of their segmentations. Alternatively, we could also skeletonize the final membrane segmentation if this "soft skeletonization" does not work properly. I'll keep you posted :)

alisterburt commented 11 months ago

Thanks for the paper @LorenzLamm - that's a really cool method!

bbarad commented 11 months ago

I am gonna submit a PR that will close this issue today, perhaps we should open a new feature issue for implementing end user skeletonized output (separate from the loss consideration)?

LorenzLamm commented 11 months ago

Yes, good idea. I'll open a new issue.

LorenzLamm commented 9 months ago

Added the new functionality for adjusting the threshold in the new PR. Now, it's possible to choose a custom threshold (as compared to default 0.0) with the membrain segment function:

membrain segment --segmentation-threshold 5.0 ... more info here: https://github.com/teamtomo/membrain-seg/blob/main/docs/Usage/Segmentation.md#more-membrain-segment-arguments

bbarad commented 9 months ago

Whoops, apologies for letting this slip! Thanks for adding this feature Lorenz! We'll use it heavily.