Add micro-sam (μSAM) watershed postprocessing - #501
Draft
nilsmechtel wants to merge 2 commits into
Draft
Conversation
Ports the seeded watershed from micro_sam's InstanceSegmentationWithDecoder (torch_em's watershed_from_center_and_boundary_distances) as bioimageio.core.microsam_watershed: foreground + center distance + boundary distance maps in, instance labels out. Matches the upstream defaults (thresholds 0.5, foreground smoothing 1.0, distance smoothing 1.6) and the seeded-watershed-on-boundary-distances formulation. scikit-image is required only for this function and added as the optional extra bioimageio.core[microsam], following the stardist extra.
Member
|
just some guidance: Any added op should inherite from |
Restructures the module from a standalone numpy callable into MicroSamWatershed(SamplewiseOperator) with a from_proc_descr classmethod, and registers it in proc_ops (import, Processing union, get_proc dispatch), following the CellposeFlowDynamics precedent. Requires MicroSamWatershedDescr from bioimageio.spec.model.v0_5.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
micro-sam models with the additional instance segmentation (AIS) decoder predict three dense maps (foreground probability, center distance, inverted boundary distance). Exported as bioimageio packages, such models run a plain prompt-free prediction pass, but turning the maps into instance labels requires micro-sam's seeded watershed, which so far only exists inside the
micro_sampackage. This PR makes that final step available to any tool building onbioimageio.core, following the precedent of the stardist and cellpose postprocessing already in this repo.What is added
MicroSamWatershed, aSamplewiseOperatorthat replaces the three stacked decoder maps on its sample member with instance labels. Seeds are connected components where both smoothed distance maps are below their thresholds inside the foreground mask; the seeded watershed runs on the smoothed boundary distances restricted to that mask. Faithful port ofmicro_sam.instance_segmentation.InstanceSegmentationWithDecoder/torch_em.util.segmentation.watershed_from_center_and_boundary_distances, with the upstream parameter defaults (thresholds 0.5, foreground smoothing 1.0, distance smoothing 1.6,min_size0).from_proc_descr(MicroSamWatershedDescr, member_id)plus registration inproc_ops— import,Processingunion, and aget_procdispatch branch — mirroringCellposeFlowDynamics.bioimageio.core[microsam]providingscikit-image(only needed for the watershed itself; gaussian smoothing and connected components use the existing scipy dependency), mirroring thestardistextra. scikit-image is imported lazily so the module stays importable without it.min_sizefiltering, channel-count validation, andfrom_proc_descrkwargs mapping.Depends on bioimage-io/spec-bioimage-io#768
from_proc_descrbinds toMicroSamWatershedDescr, which that PR adds tobioimageio.spec.model.v0_5. Until it is released, thebioimageio.spec ==0.5.12.0pin here does not contain the descriptor and CI cannot pass; the pin needs bumping to whichever release carries it. Verified locally against a co-developed spec checkout (extraPaths = ["../spec-bioimage-io/src"]).Test plan
pytest tests/test_ops_microsam.py— 4 passed, against the spec branch of #768.get_procdispatch checked end to end: anOutputTensorDescrcarrying aMicroSamWatershedDescrpostprocessing entry resolves to aMicroSamWatershedwith the declared kwargs and the tensor's member id.