Skip to content

Add micro-sam (μSAM) watershed postprocessing - #501

Draft
nilsmechtel wants to merge 2 commits into
mainfrom
add-microsam-watershed-postprocessing
Draft

Add micro-sam (μSAM) watershed postprocessing#501
nilsmechtel wants to merge 2 commits into
mainfrom
add-microsam-watershed-postprocessing

Conversation

@nilsmechtel

@nilsmechtel nilsmechtel commented Aug 17, 2026

Copy link
Copy Markdown
Member

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_sam package. This PR makes that final step available to any tool building on bioimageio.core, following the precedent of the stardist and cellpose postprocessing already in this repo.

What is added

  • MicroSamWatershed, a SamplewiseOperator that 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 of micro_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_size 0).
  • from_proc_descr(MicroSamWatershedDescr, member_id) plus registration in proc_ops — import, Processing union, and a get_proc dispatch branch — mirroring CellposeFlowDynamics.
  • Optional extra bioimageio.core[microsam] providing scikit-image (only needed for the watershed itself; gaussian smoothing and connected components use the existing scipy dependency), mirroring the stardist extra. scikit-image is imported lazily so the module stays importable without it.
  • Unit tests with synthetic two-cell maps: instance separation across a boundary ridge, min_size filtering, channel-count validation, and from_proc_descr kwargs mapping.

Depends on bioimage-io/spec-bioimage-io#768

from_proc_descr binds to MicroSamWatershedDescr, which that PR adds to bioimageio.spec.model.v0_5. Until it is released, the bioimageio.spec ==0.5.12.0 pin 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_proc dispatch checked end to end: an OutputTensorDescr carrying a MicroSamWatershedDescr postprocessing entry resolves to a MicroSamWatershed with the declared kwargs and the tensor's member id.
  • Behavior cross-checked against the micro_sam / torch_em source and the parameter defaults used by the micro-sam BioEngine deployment.

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.
@FynnBe

FynnBe commented Aug 25, 2026

Copy link
Copy Markdown
Member

just some guidance: Any added op should inherite from SamplewiseOperator or BlockwiseOperator and should have a from_proc_descr method that is called with an appropriate bioimageio.spec.model.v0_5 object.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants