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

skop/ops/mask/mobilesam.py

@@ -0,0 +1,176 @@
+"""Box-prompted segmentation with MobileSAMv2's prompt-guided decoder.
+
+The other half of the model ``skop.ops.detect.object_aware_yolo`` wraps, and
+the reason both live in the same environment: the object-aware YOLO was
+trained to produce box prompts for exactly this decoder, and both come out of
+the vendored MobileSAMv2 tree that segment-everything carries.
+
+The model is assembled from three pieces, which is what
+``weights_helper.create_mobile_sam_model`` does:
+
+- an empty ``vit_h`` Sam, for the postprocessing and the thresholds,
+- EfficientViT-L2 as its image encoder (``l2.pt``, 234 MB),
+- MobileSAMv2's prompt encoder and mask decoder, which are one 15 MB file.
+
+Assembled here rather than by calling ``create_mobile_sam_model``, for the
+reason ``object_aware_yolo`` gives: ``segment_everything.weights_helper``
+imports toolz and ``napari.utils.progress`` at module scope, and
+segment-everything declares neither. Importing it means installing napari into
+a worker to download a file. Both weight files come from the HuggingFace
+mirror already listed in its own ``WEIGHTS_URLS``, rather than through gdown
+from Google Drive.
+
+Its sibling ``microsam_masks`` is a port; this is a wrapper. That asymmetry is
+the whole argument of docs/design/0007: micro_sam is a conda package, and this
+decoder only exists inside a 28,000-line vendored tree.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+import ophelper.boxes
+import ophelper.masks
+from ophelper.assets import file_from_url
+from opspec.op import op
+from opspec.runtime import progress
+from opspec.types import BoxesOf, ImageOf
+
+from .._util import to_rgb
+from ._result import Masks
+
+# EfficientViT-L2, the image encoder. The same file weights_helper fetches
+# through gdown as "efficientvit_l2"; this is its HuggingFace mirror.
+_ENCODER = "https://huggingface.co/RogerQi/MobileSAMV2/resolve/main/l2.pt"
+
+
+@op(env="segment-everything")
+def mobilesam_masks(
+    image: ImageOf[np.ndarray],
+    boxes: BoxesOf[np.ndarray],
+    batch_size: int = 100,
+) -> Masks:
+    """Segment one object per bounding box.
+
+    Args:
+        image: Image to segment. 2-D or RGB; anything else is stretched to
+            8-bit RGB first, which is what the model was trained on.
+        boxes: (N, 4) as [min_y, min_x, max_y, max_x], from a detector op, a
+            Shapes layer, or ``ophelper.boxes.from_labels``.
+        batch_size: How many boxes to decode at once. The image is encoded
+            once regardless; this only trades GPU memory against the number
+            of decoder calls.
+
+    Returns:
+        masks: (N, Y, X) uint8, one object per plane. These may overlap, which
+            is why they are not a label image; project them with
+            ``ophelper.masks``.
+        boxes: the prompts that produced them.
+
+    Masks come back in prompt order, not sorted -- ``masks[i]`` answers
+    ``boxes[i]``. Boxes the decoder finds nothing in are dropped from both, so
+    this may be shorter than what went in.
+
+    Note: the image encoder is 234 MB, downloaded into skop's asset cache on
+    first use. The decoder ships inside segment-everything itself.
+    """
+    from pathlib import Path
+
+    import segment_everything
+    import torch
+    from segment_everything.vendored.mobilesamv2 import (
+        SamPredictor,
+        sam_model_registry,
+    )
+
+    prompts = ophelper.boxes.as_boxes(boxes)
+    if len(prompts) == 0:
+        return Masks(ophelper.masks.empty(image.shape[:2]), ophelper.boxes.EMPTY.copy())
+
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif torch.backends.mps.is_available():
+        device = "mps"
+    else:
+        device = "cpu"
+
+    progress("Loading MobileSAMv2")
+    encoder_weights = file_from_url(_ENCODER, "mobilesam/l2.pt")
+    # Shipped inside the installed package rather than downloaded -- this is
+    # the one entry in WEIGHTS_URLS marked "local:". Same file as the
+    # Prompt_guided_Mask_Decoder.pt on the mirror above, if it ever moves.
+    decoder_weights = (
+        Path(segment_everything.__file__).parent
+        / "vendored/PromptGuidedDecoder/Prompt_guided_Mask_Decoder.pt"
+    )
+
+    prompt_guided = sam_model_registry["PromptGuidedDecoder"](str(decoder_weights))
+    model = sam_model_registry["vit_h"]()
+    model.prompt_encoder = prompt_guided["PromtEncoder"]  # sic, upstream spelling
+    model.mask_decoder = prompt_guided["MaskDecoder"]
+    model.image_encoder = sam_model_registry["efficientvit_l2"](str(encoder_weights))
+    model.to(device)
+
+    predictor = SamPredictor(model)
+
+    # The expensive half, and it depends only on the image -- which is why
+    # this op takes every box at once rather than one per call (workflow-ops).
+    progress("Encoding the image")
+    predictor.set_image(to_rgb(image))
+
+    # SamPredictor's transform resizes prompts alongside the image it just
+    # encoded, and it reads them as xyxy -- the one conversion this op needs.
+    scaled = predictor.transform.apply_boxes(
+        ophelper.boxes.to_xyxy(prompts), predictor.original_size
+    )
+    scaled = torch.as_tensor(scaled, dtype=torch.float32, device=device)
+
+    # One encoding, broadcast across each batch of prompts. Sliced from the
+    # full repeat rather than reassigned in place: the original narrowed these
+    # tensors on every iteration, which only survived because the short batch
+    # is always the last one.
+    image_embedding = torch.repeat_interleave(predictor.features, batch_size, dim=0)
+    prompt_embedding = torch.repeat_interleave(
+        model.prompt_encoder.get_dense_pe(), batch_size, dim=0
+    )
+
+    decoded = []
+    for start in range(0, len(scaled), batch_size):
+        batch = scaled[start : start + batch_size]
+        progress(
+            f"Segmenting objects {start + 1}-{start + len(batch)} of {len(scaled)}",
+            start,
+            len(scaled),
+        )
+        with torch.no_grad():
+            sparse, dense = model.prompt_encoder(points=None, boxes=batch, masks=None)
+            low_res, _predicted_iou = model.mask_decoder(
+                image_embeddings=image_embedding[: len(batch)],
+                image_pe=prompt_embedding[: len(batch)],
+                sparse_prompt_embeddings=sparse,
+                dense_prompt_embeddings=dense,
+                multimask_output=False,
+                simple_type=True,
+            )
+            # _predicted_iou is a confidence per mask, dropped for the same
+            # reason Boxes drops its scores -- see _result.py and the
+            # per-object-features spec. The
+            # original also computed a stability score here, which cost a
+            # second thresholding pass for a value nothing consumed.
+            full = predictor.model.postprocess_masks(
+                low_res, predictor.input_size, predictor.original_size
+            )
+            decoded.append((full > model.mask_threshold).squeeze(1).to(torch.uint8))
+
+    stacked = torch.cat(decoded).cpu().numpy()
+    del decoded
+    if device == "cuda":
+        torch.cuda.empty_cache()
+
+    # The decoder answers some prompts with nothing at all. Dropping those is
+    # why the surviving boxes are returned alongside the masks.
+    kept = stacked.any(axis=(1, 2))
+    progress(f"Segmented {int(kept.sum())} of {len(prompts)} objects")
+    if not kept.any():
+        return Masks(ophelper.masks.empty(image.shape[:2]), ophelper.boxes.EMPTY.copy())
+    return Masks(np.ascontiguousarray(stacked[kept]), prompts[kept])

skop/ops/morphology.py

@@ -0,0 +1,214 @@
+"""Grayscale morphology, via scikit-image.
+
+Morphology asks a different question from the filters in
+:mod:`skop.ops.smooth`: not "what is the average around here" but "does a
+shape of this size fit here". Every op below slides one structuring element
+-- a footprint of a given radius -- over the image and takes an extremum
+under it, and the whole family follows from the two that do:
+:func:`erosion` takes the minimum, :func:`dilation` the maximum, and the
+rest are those two composed.
+
+Because they choose an existing value rather than computing a new one, they
+introduce no intermediate intensities and preserve the input's dtype -- which
+is why they are the safe thing to run on a label image or a mask, and why
+:func:`white_tophat` is the classical background subtraction.
+
+The radius is the whole parameter: it says what counts as small. Anything
+smaller than the footprint does not survive an opening, and that is the
+knob to turn. Its shape is the second, and :class:`Footprint` -- re-exported
+here, since this is the namespace it belongs to -- is what names the three
+choices. :func:`skop.ops.smooth.median` takes the same enum.
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+import numpy as np
+
+from opspec.op import Axes, op
+from opspec.types import ImageOf
+
+from ._util import Footprint, channel_axis, footprint, per_channel
+
+#: Every op here works on any number of axes, RGB(A) channel by channel.
+_Image = Annotated[ImageOf[np.ndarray], Axes(variadic=True)]
+
+
+def _apply(image: np.ndarray, radius: int, shape: Footprint, fn) -> np.ndarray:
+    """Run a morphological op with a footprint matching the image's rank."""
+    if radius < 1:
+        raise ValueError(f"radius must be at least 1, got {radius}")
+    spatial = image.ndim - (0 if channel_axis(image) is None else 1)
+    fp = footprint(spatial, radius, shape)
+    return per_channel(image, lambda plane: fn(plane, footprint=fp))
+
+
+@op(env="skimage")
+def erosion(
+    image: _Image,
+    radius: int = 1,
+    shape: Footprint = Footprint.ball,
+) -> ImageOf[np.ndarray]:
+    """Shrink bright regions, by taking the minimum under the footprint.
+
+    Bright objects lose a border of ``radius`` and dark ones grow by the
+    same; anything bright and thinner than the footprint disappears. On a
+    binary mask this is how touching objects are pulled apart before a
+    connected-components pass.
+
+    Args:
+        image: Image to erode. Any number of axes; a trailing RGB(A) axis is
+            eroded channel by channel.
+        radius: Radius of the footprint, in pixels.
+        shape: Shape of the footprint.
+
+    Returns:
+        The eroded image, in the input's dtype.
+    """
+    from skimage import morphology
+
+    return _apply(image, radius, shape, morphology.erosion)
+
+
+@op(env="skimage")
+def dilation(
+    image: _Image,
+    radius: int = 1,
+    shape: Footprint = Footprint.ball,
+) -> ImageOf[np.ndarray]:
+    """Grow bright regions, by taking the maximum under the footprint.
+
+    The dual of :func:`erosion`: bright objects gain a border of ``radius``,
+    and gaps narrower than the footprint close up. The usual use is joining a
+    structure that noise or thresholding has broken into pieces.
+
+    Args:
+        image: Image to dilate. Any number of axes; a trailing RGB(A) axis is
+            dilated channel by channel.
+        radius: Radius of the footprint, in pixels.
+        shape: Shape of the footprint.
+
+    Returns:
+        The dilated image, in the input's dtype.
+    """
+    from skimage import morphology
+
+    return _apply(image, radius, shape, morphology.dilation)
+
+
+@op(env="skimage")
+def opening(
+    image: _Image,
+    radius: int = 1,
+    shape: Footprint = Footprint.ball,
+) -> ImageOf[np.ndarray]:
+    """Remove bright detail smaller than the footprint.
+
+    An erosion followed by a dilation. The erosion deletes whatever the
+    footprint does not fit inside, and the dilation restores what survived to
+    roughly its original size -- so unlike a plain erosion this removes small
+    bright specks without shrinking everything else. Sizes and shapes larger
+    than the footprint come through essentially untouched.
+
+    Args:
+        image: Image to open. Any number of axes; a trailing RGB(A) axis is
+            opened channel by channel.
+        radius: Radius of the footprint, in pixels. Bright structure thinner
+            than this is what gets removed.
+        shape: Shape of the footprint.
+
+    Returns:
+        The opened image, in the input's dtype.
+    """
+    from skimage import morphology
+
+    return _apply(image, radius, shape, morphology.opening)
+
+
+@op(env="skimage")
+def closing(
+    image: _Image,
+    radius: int = 1,
+    shape: Footprint = Footprint.ball,
+) -> ImageOf[np.ndarray]:
+    """Remove dark detail smaller than the footprint.
+
+    A dilation followed by an erosion: the mirror of :func:`opening`, filling
+    small dark holes and thin dark gaps without growing the bright regions
+    around them. Between the two, most mask cleanup is a closing to mend
+    objects and then an opening to drop debris.
+
+    Args:
+        image: Image to close. Any number of axes; a trailing RGB(A) axis is
+            closed channel by channel.
+        radius: Radius of the footprint, in pixels. Dark structure thinner
+            than this is what gets filled.
+        shape: Shape of the footprint.
+
+    Returns:
+        The closed image, in the input's dtype.
+    """
+    from skimage import morphology
+
+    return _apply(image, radius, shape, morphology.closing)
+
+
+@op(env="skimage")
+def white_tophat(
+    image: _Image,
+    radius: int = 1,
+    shape: Footprint = Footprint.ball,
+) -> ImageOf[np.ndarray]:
+    """Keep bright detail smaller than the footprint, dropping the background.
+
+    The image minus its :func:`opening` -- that is, exactly what the opening
+    removed. Since an opening keeps whatever is larger than the footprint,
+    the difference is the small bright structure sitting on top of it, with
+    slow background variation subtracted away.
+
+    This is background subtraction, done morphologically: a radius comfortably
+    larger than the objects and smaller than the illumination gradient
+    flattens the field without touching the objects.
+
+    Args:
+        image: Image to filter. Any number of axes; a trailing RGB(A) axis is
+            filtered channel by channel.
+        radius: Radius of the footprint, in pixels. Bright structure smaller
+            than this is what is kept.
+        shape: Shape of the footprint.
+
+    Returns:
+        The bright detail, in the input's dtype.
+    """
+    from skimage import morphology
+
+    return _apply(image, radius, shape, morphology.white_tophat)
+
+
+@op(env="skimage")
+def black_tophat(
+    image: _Image,
+    radius: int = 1,
+    shape: Footprint = Footprint.ball,
+) -> ImageOf[np.ndarray]:
+    """Keep dark detail smaller than the footprint, dropping the background.
+
+    The :func:`closing` minus the image, and the dual of
+    :func:`white_tophat`: it isolates small dark structure on a bright,
+    slowly varying background. What to reach for on a brightfield image,
+    where the objects are the dark things.
+
+    Args:
+        image: Image to filter. Any number of axes; a trailing RGB(A) axis is
+            filtered channel by channel.
+        radius: Radius of the footprint, in pixels. Dark structure smaller
+            than this is what is kept.
+        shape: Shape of the footprint.
+
+    Returns:
+        The dark detail, in the input's dtype.
+    """
+    from skimage import morphology
+
+    return _apply(image, radius, shape, morphology.black_tophat)

skop/ops/normalize.py

@@ -0,0 +1,85 @@
+"""Intensity normalization, separated from the models that assume it.
+
+Every deep learning segmenter wants its input scaled the way its training
+data was, and most of them do it silently on whatever array they are handed.
+That is fine until the array is a slice of something bigger.
+
+StarDist is the case this exists for. ``stardist2d_fluo`` declares
+``Axes("y", "x", "c?")``, so a volume handed to it is looped over plane by
+plane -- and csbdeep normalizes *each plane* against its own percentiles. A
+top or bottom plane holding nothing but out-of-focus blur gets stretched until
+that blur fills the range and reads as nuclei, so the segmenter finds objects
+there that are not there.
+
+Normalizing first, over the whole volume, fixes it: the empty planes stay
+dim relative to the ones with signal in them, because they are measured
+against the same range. Hence ``Axes(variadic=True)`` here, the same
+declaration ``otsu`` carries -- percentiles do not care how many axes they are
+given, so the volume arrives whole and is normalized in one call, and only
+then does the 2-D segmenter run slice by slice.
+
+Pure numpy, so it lives in the 'minimal' environment and composes with a model
+in any other. An op's environment is where it runs, not where its caller runs:
+this normalizing in 'minimal' and feeding StarDist in 'stardist-tf' is two
+worker processes and one shared-memory handoff.
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+import numpy as np
+
+from opspec.op import Axes, op
+from opspec.types import ImageOf
+
+
+@op(env="minimal")
+def percentile(
+    image: Annotated[ImageOf[np.ndarray], Axes(variadic=True)],
+    low: Annotated[
+        float,
+        {"widget_type": "FloatSpinBox", "min": 0.0, "max": 100.0, "step": 0.1},
+    ] = 1.0,
+    high: Annotated[
+        float,
+        {"widget_type": "FloatSpinBox", "min": 0.0, "max": 100.0, "step": 0.1},
+    ] = 99.8,
+    clip: bool = False,
+) -> ImageOf[np.ndarray]:
+    """Rescale intensities so that two percentiles land on 0 and 1.
+
+    The normalization csbdeep applies inside StarDist and CARE, as an op that
+    can be called on its own -- which matters whenever the thing consuming it
+    works on fewer axes than you have. See this module's docstring.
+
+    Percentiles rather than min and max, because one hot pixel or one dead
+    one would otherwise set the whole range.
+
+    Args:
+        image: Image to normalize. Any number of axes; all of them are
+            measured together, so a volume gets one range rather than one per
+            plane. A trailing RGB axis is included in that, which is right for
+            a photograph and wrong for unrelated fluorescence channels --
+            normalize those separately for now.
+        low: Percentile mapped to 0. StarDist's own default is 1.
+        high: Percentile mapped to 1. StarDist's own default is 99.8.
+        clip: Whether to cut everything outside [0, 1]. Off by default,
+            matching csbdeep: values beyond the percentiles are real signal,
+            and a model trained this way saw them.
+
+    Returns:
+        The image as float32, with ``low`` at 0 and ``high`` at 1.
+    """
+    if not 0.0 <= low < high <= 100.0:
+        raise ValueError(f"need 0 <= low < high <= 100, got low={low}, high={high}")
+
+    array = np.asarray(image, dtype=np.float32)
+    lo, hi = np.percentile(array, [low, high])
+
+    # The epsilon is csbdeep's, and it is what keeps a blank image -- every
+    # pixel identical, so lo == hi -- from dividing by zero and returning nan.
+    out = (array - lo) / (hi - lo + 1e-20)
+    if clip:
+        out = np.clip(out, 0.0, 1.0)
+    return out.astype(np.float32)

skop/ops/segment/__init__.py

@@ -0,0 +1,24 @@
+"""Segmentation ops: label the objects in an image.
+
+Each op here gets a module of its own, because each brings its own
+environment, its own pretrained models and its own supporting types. The
+re-exports below are the namespace's public surface -- written out rather
+than generated, so that an IDE can offer them.
+"""
+
+from __future__ import annotations
+
+from .cellpose4 import cellpose4
+from .cellpose3 import cellpose3
+from .stardist2d import stardist2d_fluo, stardist2d_he
+from .starfun3d import segment_nuclei
+from .unseg import unseg
+
+__all__ = [
+    "cellpose4",
+    "cellpose3",
+    "segment_nuclei",
+    "stardist2d_fluo",
+    "stardist2d_he",
+    "unseg",
+]