scikit-ops 0.1.0.dev0__py3-none-any.whl

skop/ops/workflows/mask/detect_then_mask.py

@@ -0,0 +1,97 @@
+"""Find the objects, then draw each one: a detector paired with a masker.
+
+The two-stage segmenter from docs/spec/workflow-ops.md, and the reason the
+detectors and mask detectors were built to share signatures in the first
+place. Any of the two detectors pairs with either of the two mask detectors,
+which is four segmenters from four ops -- and the point of the pairing is that
+they are genuinely different: the object-aware YOLO finds objects a
+photograph-trained SAM misses, and micro_sam draws cells that MobileSAM
+smears.
+
+Both stages take the image. Only one asks for it: ``binds`` on each stage's
+arguments names what this workflow supplies, so a front end renders the rest
+and nothing gets typed twice. The masker's ``boxes`` are bound for the same
+reason, though they come from the detector rather than from the panel.
+
+Note that a pairing may cross environments -- ``fastsam`` in 'pytorch' with
+``mobilesam_masks`` in 'segment-everything' means two workers, and the image
+crosses shared memory twice. That is the cost of composing on the host, and
+it is why this is a workflow rather than an op.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Annotated, NamedTuple
+
+import numpy as np
+
+from opspec.op import Choices, ParamsFor, op
+from opspec.runtime import progress, run
+from opspec.types import BoxesOf, ImageOf, MasksOf
+from skop.ops.detect.fastsam import fastsam
+from skop.ops.detect.object_aware_yolo import object_aware_yolo
+from skop.ops.mask.microsam import microsam_masks
+from skop.ops.mask.mobilesam import mobilesam_masks
+
+
+class Detected(NamedTuple):
+    """What each stage produced, in the order the stages ran.
+
+    Stage order rather than importance: a front end adds these to a viewer in
+    declaration order, so the pipeline reads down the layer list, and the last
+    stage -- the answer -- ends up on top where a new layer belongs. This is
+    why the fields sit the other way round from ``skop.ops.mask.Masks``, which
+    every mask detector returns. There, the masks are the point and the boxes
+    record which prompts survived; here they are stage one and stage two.
+    """
+
+    #: The prompts that produced the masks: the detector's boxes, minus any
+    #: the mask detector answered with nothing.
+    boxes: BoxesOf[np.ndarray]
+    #: (N, Y, X) uint8, one object per plane, possibly overlapping.
+    masks: MasksOf[np.ndarray]
+
+
+@op()
+def detect_then_mask(
+    image: ImageOf[np.ndarray],
+    detector: Annotated[
+        Callable,
+        Choices(object_aware_yolo=object_aware_yolo, fastsam=fastsam),
+    ] = object_aware_yolo,
+    detector_args: Annotated[dict | None, ParamsFor("detector", binds="image")] = None,
+    masker: Annotated[
+        Callable,
+        Choices(mobilesam=mobilesam_masks, microsam=microsam_masks),
+    ] = mobilesam_masks,
+    masker_args: Annotated[
+        dict | None, ParamsFor("masker", binds=("image", "boxes"))
+    ] = None,
+) -> Detected:
+    """Detect objects and segment each one.
+
+    Args:
+        image: The image to segment.
+        detector: Which op finds the bounding boxes. ``object_aware_yolo`` was
+            trained on microscopy; ``fastsam`` is a general detector.
+        detector_args: Settings for the chosen detector, minus the image.
+        masker: Which op draws a mask inside each box. ``microsam`` is
+            finetuned for microscopy and slower; ``mobilesam`` is the one the
+            object-aware detector was trained to pair with.
+        masker_args: Settings for the chosen mask detector, minus the image
+            and the boxes, which this workflow supplies.
+
+    Returns:
+        boxes: the box each mask answers.
+        masks: (N, Y, X) uint8, one object per plane, possibly overlapping.
+
+    A detector that finds nothing is an ordinary outcome, not a failure: the
+    mask detector is skipped and the result is empty.
+    """
+    progress("Detecting objects")
+    boxes = run(detector, image=image, **(detector_args or {})).boxes
+
+    progress(f"Segmenting {len(boxes)} objects")
+    found = run(masker, image=image, boxes=boxes, **(masker_args or {}))
+    return Detected(found.boxes, found.masks)

skop/ops/workflows/segment/__init__.py

@@ -0,0 +1,7 @@
+"""Workflows that produce a label image."""
+
+from __future__ import annotations
+
+from .connect_2d_in_3d import Connected, connect_2d_in_3d
+
+__all__ = ["Connected", "connect_2d_in_3d"]

skop/ops/workflows/segment/connect_2d_in_3d.py

@@ -0,0 +1,145 @@
+"""Segment a volume plane by plane, then join the planes back into objects.
+
+The workflow ``skop.ops.normalize`` was written for. A 2-D segmenter handed a
+volume is looped over its planes, and every one of these segmenters normalizes
+whatever array it is given -- so each plane gets stretched to its own range,
+and a top or bottom plane holding nothing but out-of-focus blur is stretched
+until that blur fills the range and reads as nuclei. The segmenter then finds
+objects there that are not there.
+
+Three ops in a row fix it, and the order is the whole point:
+
+1. normalize over the **whole volume**, so an empty plane stays dim relative
+   to the planes with signal in them.
+2. the segmenter, **plane by plane and with its own normalization off**, so
+   nothing puts the stretch back.
+3. connect, which links objects that overlap between adjacent planes, so one
+   cell stops being a stack of unrelated labels.
+
+Step 2 is why ``normalize`` is bound rather than offered. Leaving a checkbox
+there would let someone switch back on precisely the behaviour this exists to
+avoid, and get a result that looks like a failure of the idea rather than of
+the setting. Running a segmenter with its own normalization on is still a
+perfectly good thing to do -- from the Ops panel, where that is what you asked
+for.
+
+**All three stages are choosers, including the two with one choice each.**
+There is only one normalizer and one connector today and there will be more,
+and a stage that starts as a chooser gains its second option by having a name
+added to a list rather than by being rewritten. It also keeps the shape
+uniform: every stage's settings come from the chosen op's own signature, so
+nothing here restates a parameter that lives somewhere else and can drift from
+it.
+
+**Three dimensions, and the dimensions are held fixed.** The connector needs
+an axis to walk along, so there is no 2-D version of this. The workflow
+declares ``Axes("z", "y", "x")`` on its own input, which is what lets a front
+end ask which of *your* axes is z -- but it does not pass that question down.
+Each stage's axis handling is pinned here, because getting it wrong is the bug
+the workflow is for.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Annotated, NamedTuple
+
+import numpy as np
+
+from opspec.op import Axes, Choices, ParamsFor, op
+from opspec.runtime import progress, run
+from opspec.types import ImageOf, LabelsOf
+from skop.ops.labels import connect
+from skop.ops.normalize import percentile
+from skop.ops.segment.cellpose3 import cellpose3
+from skop.ops.segment.cellpose4 import cellpose4
+from skop.ops.segment.stardist2d import stardist2d_fluo
+
+#: What every stage is handed, and what the workflow refuses to make adjustable.
+_ZYX = list("zyx")
+
+
+class Connected(NamedTuple):
+    """What each stage produced, in the order the stages ran.
+
+    Stage order rather than importance: a front end adds these to a viewer in
+    declaration order, so the pipeline reads down the layer list, and the last
+    stage -- the answer -- ends up on top where a new layer belongs.
+    """
+
+    #: The normalized volume the segmenter actually saw. Returned because it
+    #: is the evidence: if the top plane is still dim here, the normalization
+    #: did its job, and a spurious object there is the segmenter's doing.
+    normalized: ImageOf[np.ndarray]
+    #: One label per object, consistent through the stack.
+    labels: LabelsOf[np.ndarray]
+
+
+@op()
+def connect_2d_in_3d(
+    image: Annotated[ImageOf[np.ndarray], Axes("z", "y", "x")],
+    normalizer: Annotated[Callable, Choices(percentile=percentile)] = percentile,
+    normalizer_args: Annotated[
+        dict | None, ParamsFor("normalizer", binds="image")
+    ] = None,
+    segmenter: Annotated[
+        Callable,
+        Choices(stardist=stardist2d_fluo, cellpose4=cellpose4, cellpose3=cellpose3),
+    ] = stardist2d_fluo,
+    segmenter_args: Annotated[
+        dict | None, ParamsFor("segmenter", binds=("image", "normalize"))
+    ] = None,
+    connector: Annotated[Callable, Choices(connect=connect)] = connect,
+    connector_args: Annotated[
+        dict | None, ParamsFor("connector", binds="labels")
+    ] = None,
+) -> Connected:
+    """Segment a volume plane by plane and join the planes into objects.
+
+    Args:
+        image: The volume, as z, y, x. Normalized once, whole, before any
+            plane is segmented.
+        normalizer: Which op scales the volume. Whichever it is, it is run on
+            the whole volume rather than plane by plane -- that is the step
+            this workflow exists to get right.
+        normalizer_args: Settings for the chosen normalizer, minus the volume.
+        segmenter: Which 2-D segmenter to run on each plane. Each brings its
+            own environment, and picking one for the first time builds it.
+        segmenter_args: Settings for the chosen segmenter, minus the plane and
+            its normalization, both of which this workflow supplies.
+        connector: Which op links objects between adjacent planes.
+        connector_args: Settings for the chosen connector, minus the labels.
+
+    Returns:
+        normalized: The volume as the segmenter saw it.
+        labels: The connected label volume, numbered from 1.
+
+    Note: the three stages live in three environments, so a run crosses shared
+    memory three times. That is the cost of composing on the host, and it is
+    paid per call rather than per plane -- skop loops the segmenter over z
+    inside its own worker.
+    """
+    progress("Normalizing the whole volume")
+    # Naming the axes hands a variadic op the volume entire rather than a
+    # plane at a time. This is the step the whole workflow exists for.
+    normalized = run(
+        normalizer, image=image, axes={"image": _ZYX}, **(normalizer_args or {})
+    )
+
+    progress("Segmenting plane by plane")
+    # The segmenter declares two axes and is given three, so skop loops it over
+    # z. normalize=False is what keeps it from undoing the step above.
+    labels = run(
+        segmenter,
+        image=normalized,
+        normalize=False,
+        axes={"image": _ZYX},
+        **(segmenter_args or {}),
+    )
+
+    progress("Connecting objects between planes")
+    joined = run(
+        connector, labels=labels, axes={"labels": _ZYX}, **(connector_args or {})
+    )
+
+    return Connected(normalized, joined)