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

skop/ops/kernels/gibson_lanni.py

@@ -0,0 +1,147 @@
+"""The Gibson-Lanni point spread function, via sdeconv.
+
+A scalar diffraction model of a real objective looking into a real sample: it
+takes the immersion medium and the sample's refractive index, and the depth
+below the coverslip, so it produces the axially asymmetric PSF that spherical
+aberration actually gives you. That asymmetry is the reason to use this over
+:mod:`~skop.ops.kernels.paraxial`.
+
+Ported from tnia-python's ``tnia.deconvolution.psfs``, whose two functions --
+``gibson_lanni_3D`` and ``gibson_lanni_3D_partial_confocal`` -- collapse into
+one op here, the same way the two Gaussians did. The partial-confocal variant
+was the same computation plus recentring, and its ``confocal`` flag was just
+``confocal_factor=2``.
+
+**psfmodels is deliberately absent.** The original could compute this either
+through sdeconv or through psfmodels, chosen by a ``use_psfm`` flag. psfmodels
+is GPL-3.0 and this project is BSD, so that branch is not defaulted off, it is
+gone -- there is no import path here that could reach it.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Annotated
+
+import numpy as np
+
+from opspec.op import op
+from opspec.runtime import progress
+
+from ._fluorophore import Fluorophore
+from ._recenter import recenter_psf_axial
+
+__all__ = ["Fluorophore", "gibson_lanni"]
+
+
+@op(env="sdeconv")
+def gibson_lanni(
+    xy_size: Annotated[int, {"min": 4, "max": 2048}] = 128,
+    z_size: Annotated[int, {"min": 1, "max": 1024}] = 64,
+    voxel_size_xy: Annotated[float, {"min": 0.001, "step": 0.01}] = 0.1,
+    voxel_size_z: Annotated[float, {"min": 0.001, "step": 0.01}] = 0.2,
+    numerical_aperture: Annotated[float, {"min": 0.1, "max": 2.0, "step": 0.05}] = 1.4,
+    ni: Annotated[float, {"min": 1.0, "max": 2.0, "step": 0.001}] = 1.518,
+    ns: Annotated[float, {"min": 1.0, "max": 2.0, "step": 0.001}] = 1.33,
+    pz: Annotated[float, {"min": 0.0, "step": 0.1}] = 0.0,
+    wavelength: Annotated[float, {"min": 0.1, "max": 2.0, "step": 0.01}] = 0.53,
+    confocal_factor: Annotated[float, {"min": 1.0, "max": 2.0, "step": 0.1}] = 1.0,
+    recenter: bool = True,
+) -> np.ndarray:
+    """Generate a 3-D Gibson-Lanni point spread function.
+
+    Args:
+        xy_size: Extent of the PSF in XY, in voxels. Square.
+        z_size: Number of Z planes to return.
+        voxel_size_xy: Lateral voxel size, in microns.
+        voxel_size_z: Axial voxel size, in microns.
+        numerical_aperture: Numerical aperture of the objective.
+        ni: Refractive index of the immersion medium. 1.518 is oil, 1.33
+            water, 1.0 air.
+        ns: Refractive index of the sample.
+        pz: Depth of the point below the coverslip, in microns. Non-zero is
+            what makes the PSF axially asymmetric, and is the whole reason to
+            model spherical aberration.
+        wavelength: Emission wavelength, in microns. ``Fluorophore`` members
+            can be passed here by name.
+        confocal_factor: 1.0 is widefield. 2.0 is confocal -- the PSF squared,
+            since a confocal detects through the same aperture it illuminates
+            through. Values in between are an ad-hoc approximation of a
+            partially closed pinhole and should be used carefully: it is a
+            convenient interpolation, not a model of one.
+        recenter: Whether to put the peak on the centre plane. sdeconv does
+            not centre its output, so a PSF used as a deconvolution kernel
+            wants this on, or the result comes out shifted along Z.
+
+    Returns:
+        A ``(z_size, xy_size, xy_size)`` float32 PSF summing to 1.
+    """
+    from sdeconv.psfs import SPSFGibsonLanni
+
+    # NB: only sdeconv 1.x is supported, and the environment pins it. The
+    # original sniffed sdeconv.__version__ and branched to `PSFGibsonLanni`
+    # with a different constructor for 0.x. 0.x is the only version on
+    # conda-forge but is years old and takes wavelength in nanometres with no
+    # ni/ns at all, so supporting both would mean two different models behind
+    # one signature.
+    compute_z = _compute_z(z_size, voxel_size_z, pz) if recenter else z_size
+
+    progress(f"Computing a {compute_z}x{xy_size}x{xy_size} Gibson-Lanni PSF")
+    model = SPSFGibsonLanni(
+        (compute_z, xy_size, xy_size),
+        NA=numerical_aperture,
+        ni=ni,
+        ni0=ni,
+        ns=ns,
+        res_lateral=voxel_size_xy,
+        res_axial=voxel_size_z,
+        wavelength=wavelength,
+        pZ=pz,
+    )
+
+    # sdeconv computes in torch and hands back a tensor, on the CPU here since
+    # nothing above asked for a device.
+    psf = model().cpu().numpy().astype(np.float32)
+
+    if recenter:
+        progress("Recentring")
+        psf = recenter_psf_axial(psf, z_size)
+
+    # NB: the power comes after the crop and before the normalization, which
+    # is the order the partial-confocal original used. It matters: squaring
+    # first and cropping second would normalize over planes that get thrown
+    # away.
+    if confocal_factor != 1.0:
+        psf = psf**confocal_factor
+
+    return psf / psf.sum()
+
+
+def _compute_z(z_size: int, voxel_size_z: float, pz: float) -> int:
+    """How many planes to compute so that ``z_size`` can be cropped out centred.
+
+    The original used a flat ``z_size + z_size // 2``, and the spec described
+    that as what "makes a non-zero ``pz`` usable". Measured against sdeconv
+    1.0.4, it is not: the PSF's peak marches toward plane zero as the point
+    goes deeper -- about 6.5 planes per micron of ``pz`` at a 0.2 um axial
+    voxel, or 1.3 um of apparent focal shift per micron of depth -- and the
+    flat headroom is used up somewhere around ``pz`` = 1 um. Past that the
+    crop runs off the front of the volume.
+
+    Which the original did not notice, because it sliced without checking and
+    got a silently short array back. :func:`~skop.ops.kernels.recenter_psf_axial`
+    raises instead, so the same PSFs that used to come back quietly wrong
+    would now come back as errors -- hence this.
+
+    The headroom therefore scales with the shift. The factor of 3 is empirical
+    and deliberately generous: the measured shift is ~1.3 um per micron of
+    depth, and doubling that (the peak has to clear ``z_size // 2``, not just
+    stay in the volume) gives ~2.6. At 3 the margin grows rather than shrinks
+    with depth -- 7 planes to spare at ``pz`` = 0, 31 at ``pz`` = 16 um, for
+    ``z_size`` = 32 -- and the cost is only a taller FFT.
+
+    The shift depends on the refractive index mismatch, so a strongly
+    mismatched ``ni``/``ns`` could still outrun this. That case is what the
+    error from ``recenter_psf_axial`` is for; it says to raise ``z_size``.
+    """
+    return z_size + z_size // 2 + math.ceil(3.0 * pz / voxel_size_z)

skop/ops/kernels/paraxial.py

@@ -0,0 +1,128 @@
+"""Paraxial (scalar, in-focus) optical transfer function and point spread function.
+
+The cheapest useful optical model there is: a 2-D diffraction-limited OTF that
+follows from wavelength, numerical aperture and pixel size alone, with no
+immersion medium, no sample refractive index, and no focal depth. Where
+Gibson-Lanni models a real objective looking into a real sample, this models
+the textbook. It costs no environment of its own and takes a millisecond, which
+makes it the right kernel for a quick simulation or a sanity check.
+
+Ported from tnia-python's ``tnia.deconvolution.psfs``, which credits
+https://github.com/jdmanton/rl_positivity_sim by James Manton.
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+import numpy as np
+
+from opspec.op import op
+
+
+@op(env="skimage")
+def paraxial_otf(
+    size: Annotated[int, {"min": 4, "max": 4096}] = 128,
+    wavelength: Annotated[float, {"min": 0.1, "max": 2.0, "step": 0.01}] = 0.53,
+    numerical_aperture: Annotated[float, {"min": 0.1, "max": 2.0, "step": 0.05}] = 1.4,
+    pixel_size: Annotated[float, {"min": 0.001, "step": 0.01}] = 0.1,
+) -> np.ndarray:
+    """Generate a paraxial optical transfer function.
+
+    Args:
+        size: Extent of the OTF, in pixels. Square.
+        wavelength: Emission wavelength, in microns. ``Fluorophore`` members
+            can be passed here by name.
+        numerical_aperture: Numerical aperture of the objective.
+        pixel_size: Pixel size in the sample plane, in microns.
+
+    Returns:
+        A ``(size, size)`` OTF, 1 at DC and falling to 0 at the cutoff
+        frequency, beyond which it is exactly 0.
+    """
+    return _otf(size, wavelength, numerical_aperture, pixel_size)
+
+
+@op(env="skimage")
+def paraxial_psf(
+    size: Annotated[int, {"min": 4, "max": 4096}] = 128,
+    wavelength: Annotated[float, {"min": 0.1, "max": 2.0, "step": 0.01}] = 0.53,
+    numerical_aperture: Annotated[float, {"min": 0.1, "max": 2.0, "step": 0.05}] = 1.4,
+    pixel_size: Annotated[float, {"min": 0.001, "step": 0.01}] = 0.1,
+) -> np.ndarray:
+    """Generate a paraxial point spread function.
+
+    The inverse transform of :func:`paraxial_otf`, which is what a
+    deconvolution op wants as a kernel.
+
+    Args:
+        size: Extent of the PSF, in pixels. Square.
+        wavelength: Emission wavelength, in microns. ``Fluorophore`` members
+            can be passed here by name.
+        numerical_aperture: Numerical aperture of the objective.
+        pixel_size: Pixel size in the sample plane, in microns.
+
+    Returns:
+        A ``(size, size)`` float32 PSF summing to 1.
+
+    Note:
+        **This PSF has negative side lobes, and Richardson-Lucy will diverge
+        on it.** An ideal paraxial PSF is the Airy intensity pattern and is
+        non-negative; this one is sampled on a grid and inverted with a DFT,
+        which leaves ringing, some of it below zero. No single negative value
+        is large -- each is around 1e-4 of the peak -- but together they hold
+        about 1% of the PSF's mass. That is invisible for simulation, and
+        fatal for deconvolution, which is built on ratios and has no meaning
+        for a negative kernel: fifty iterations of ``richardson_lucy`` on an
+        unclipped one reaches 1e53.
+
+        Clip and renormalize before deconvolving::
+
+            psf = np.clip(psf, 0.0, None)
+            psf /= psf.sum()
+
+        The op does not do this itself. At 1% of the mass the correction is
+        too large to apply silently -- it would move every number the
+        tnia-python original produced by more than a rounding error -- and the
+        simulation callers who are the reason this exists do not need it.
+    """
+    otf = _otf(size, wavelength, numerical_aperture, pixel_size)
+
+    # NB: the original cast the complex result of ifftn straight to float32,
+    # which discards the imaginary part behind a ComplexWarning. Taking .real
+    # is the same number, said out loud. The imaginary part is numerical noise
+    # either way -- the OTF is real and symmetric, so its transform is real.
+    psf = np.fft.fftshift(np.fft.ifftn(np.fft.ifftshift(otf)).real).astype(np.float32)
+    return psf / psf.sum()
+
+
+def _otf(
+    size: int, wavelength: float, numerical_aperture: float, pixel_size: float
+) -> np.ndarray:
+    """The OTF itself, shared by both ops so neither has to call the other.
+
+    An op calling an op would mean a round trip through the worker for what is
+    twenty lines of numpy.
+    """
+    resolution = 0.5 * wavelength / numerical_aperture
+
+    # NB: the centre is size/2 + 1, not size/2, and the radial coordinate is
+    # normalized by the largest coordinate *after* that offset is subtracted --
+    # so the normalization runs to size/2 - 2 rather than size/2. Both are
+    # inherited verbatim. They amount to a sub-pixel shift of the centre and a
+    # few-percent scaling of the cutoff, which matters to nobody using this as
+    # a simulation kernel, and changing either would silently move every result
+    # this has ever produced.
+    centre = size / 2 + 1
+    axis = np.linspace(0, size - 1, size) - centre
+    x, y = np.meshgrid(axis, axis)
+
+    filter_radius = 2 * pixel_size / resolution
+    radius = np.sqrt(x * x + y * y) / axis.max()
+
+    inside = radius <= filter_radius
+    # Zeroed outside the passband before the arccos, so the square root never
+    # sees a negative and the arccos never sees an argument above 1.
+    v = np.where(inside, radius / filter_radius, 0.0)
+
+    return 2 / np.pi * (np.arccos(v) - v * np.sqrt(1 - v * v)) * inside

skop/ops/kernels/psf.py

@@ -0,0 +1,59 @@
+"""Point spread functions.
+
+Ported from tnia-python's ``tnia.deconvolution.gaussian_psf``, whose separate
+``gaussian_2d`` and ``gaussian_3d`` are collapsed into one op here: they differ
+only in whether there is a Z axis.
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+import numpy as np
+
+from opspec.op import op
+
+
+@op(env="skimage")
+def gaussian_psf(
+    xy_dim: Annotated[int, {"min": 1, "max": 4096}] = 64,
+    xy_sigma: Annotated[float, {"min": 0.0, "step": 0.1}] = 2.0,
+    z_dim: Annotated[int, {"min": 0, "max": 4096}] = 0,
+    z_sigma: Annotated[float, {"min": 0.0, "step": 0.1}] = 2.0,
+) -> np.ndarray:
+    """Generate a Gaussian point spread function, in 2D or 3D.
+
+    Args:
+        xy_dim: Extent of the PSF in XY, in pixels.
+        xy_sigma: Standard deviation in XY, in pixels.
+        z_dim: Extent along Z, in pixels. 0 or 1 gives a 2D PSF.
+        z_sigma: Standard deviation along Z, in pixels. Ignored in 2D.
+
+    Returns:
+        A float64 PSF summing to 1, shaped (Z, Y, X) or (Y, X).
+    """
+    # NB: the original built this with a Python loop per voxel. The Gaussian is
+    # separable, so the outer product below is the same array for a 64^3 PSF's
+    # 260k iterations of nothing.
+    profile_xy = _profile(xy_dim, xy_sigma)
+
+    if z_dim > 1:
+        profile_z = _profile(z_dim, z_sigma)
+        psf = (
+            profile_z[:, np.newaxis, np.newaxis]
+            * profile_xy[np.newaxis, :, np.newaxis]
+            * profile_xy[np.newaxis, np.newaxis, :]
+        )
+    else:
+        psf = profile_xy[:, np.newaxis] * profile_xy[np.newaxis, :]
+
+    # NB: gaussian_3d added 1e-12 before normalizing and gaussian_2d did not.
+    # Neither does now -- after normalization the floor is ~1e-12 of the peak,
+    # which no caller can act on, and having the two paths agree is worth more.
+    return psf / psf.sum()
+
+
+def _profile(dim: int, sigma: float) -> np.ndarray:
+    """One axis of a separable Gaussian, centered on the array."""
+    coords = np.linspace(-(dim // 2), dim // 2, dim)
+    return np.exp(-(coords**2) / (2.0 * sigma**2))

skop/ops/labels.py

@@ -0,0 +1,144 @@
+"""Ops that operate on a label image, rather than producing one.
+
+Where a segmentation op turns intensities into objects, these take the objects
+as given and fix up their numbering. They share the ``skimage`` environment
+because none of them needs more than numpy and scipy.
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+import numpy as np
+
+from opspec.op import Axes, op
+from opspec.runtime import progress
+from opspec.types import LabelsOf
+
+
+@op(env="skimage")
+def connect(
+    labels: Annotated[LabelsOf[np.ndarray], Axes("z", "y", "x")],
+    threshold: float = 0.1,
+    criterion: Annotated[str, {"choices": ["iou", "iop"]}] = "iou",
+) -> LabelsOf[np.ndarray]:
+    """Give one object one label along an axis, by linking overlaps.
+
+    Segmenting a volume slice by slice numbers each slice from scratch, so one
+    cell becomes a stack of unrelated labels -- rainbow pancakes. This walks
+    the stack and relabels each slice so that an object inherits the label of
+    the object it overlaps in the slice before, turning the pancakes back into
+    solid objects.
+
+    The link is one-to-one per adjacent pair: each object may claim at most one
+    predecessor, chosen by maximizing total overlap across the whole pair
+    (:func:`scipy.optimize.linear_sum_assignment`). That is what keeps two
+    labels that merely *touch* within a slice from being fused, which is why
+    this cannot be done by running connected components over the binary union.
+    An object matching nothing starts a new label, so a label that disappears
+    for a slice does not reclaim its old number when it comes back.
+
+    The same algorithm is tracking, with Z where T usually goes; it is
+    StarDist's ``group_matching_labels`` and Cellpose's ``stitch3D``, and
+    scikit-image has no equivalent.
+
+    Args:
+        labels: Label image to connect, one object numbering per slice along
+            the first declared axis. The other axes are the slice, and are
+            never mixed into the matching.
+        threshold: Minimum overlap for two objects to be considered the same,
+            in [0, 1]. StarDist links on any overlap at all, Cellpose defaults
+            to 0.25; low is safe here because the one-to-one constraint already
+            rules out the pathological links.
+        criterion: How overlap is scored. ``"iou"`` is intersection over union,
+            the usual choice. ``"iop"`` is intersection over the *smaller* of
+            the two areas, which is more forgiving when an object grows or
+            shrinks sharply from slice to slice -- as it does at the top and
+            bottom of a cell, or in an anisotropic stack.
+
+    Returns:
+        A label image of the same shape, numbered consecutively from 1.
+    """
+    from scipy.optimize import linear_sum_assignment
+
+    if criterion not in ("iou", "iop"):
+        raise ValueError(f"criterion must be 'iou' or 'iop', got {criterion!r}")
+    if not 0.0 <= threshold <= 1.0:
+        raise ValueError(f"threshold must be in [0, 1], got {threshold}")
+
+    out = np.zeros(labels.shape, dtype=np.int32)
+    if out.size == 0:
+        return out
+    # The first slice sets the numbering everything after it inherits, so it
+    # is compacted rather than copied: that alone makes the result sequential.
+    ids, inverse = np.unique(labels[0].ravel(), return_inverse=True)
+    lut = np.arange(len(ids), dtype=np.int32)
+    if ids[0] != 0:
+        lut += 1  # No background in this slice, so nothing may map to 0.
+    out[0] = lut[inverse].reshape(labels[0].shape)
+    next_id = int(out[0].max()) + 1
+
+    for z in range(1, len(labels)):
+        progress(f"Connecting slice {z + 1} of {len(labels)}", z, len(labels))
+        prev, cur = out[z - 1], labels[z]
+        prev_ids, cur_ids, inter, prev_area, cur_area = _overlaps(prev, cur)
+
+        # Every object gets a fresh label until a match says otherwise.
+        lut = np.zeros(len(cur_ids) + 1, dtype=np.int32)
+        matches: dict[int, int] = {}
+        if len(prev_ids) and len(cur_ids):
+            if criterion == "iou":
+                score = inter / (prev_area[:, None] + cur_area[None, :] - inter)
+            else:
+                score = inter / np.minimum(prev_area[:, None], cur_area[None, :])
+            rows, cols = linear_sum_assignment(score, maximize=True)
+            keep = (score[rows, cols] >= threshold) & (score[rows, cols] > 0)
+            matches = {int(j): int(i) for j, i in zip(cols[keep], prev_ids[rows[keep]])}
+        for j in range(len(cur_ids)):
+            if j in matches:
+                lut[j + 1] = matches[j]
+            else:
+                lut[j + 1] = next_id
+                next_id += 1
+
+        # Index 0 of the lut is background; object j sits at j + 1.
+        flat = cur.ravel()
+        compact = np.searchsorted(cur_ids, flat) + 1
+        compact[flat == 0] = 0
+        out[z] = lut[compact].reshape(cur.shape)
+
+    return out
+
+
+def _overlaps(
+    prev: np.ndarray, cur: np.ndarray
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """Pixel counts shared by each pair of objects across two slices.
+
+    Returns the object labels of each slice (background excluded, sorted), the
+    dense intersection matrix between them, and each object's area. Dense
+    because the matrix is objects-by-objects, not pixels-by-pixels: a slice
+    with a thousand cells still only costs a megabyte.
+    """
+    prev_ids, prev_idx = np.unique(prev.ravel(), return_inverse=True)
+    cur_ids, cur_idx = np.unique(cur.ravel(), return_inverse=True)
+    prev_area = np.bincount(prev_idx, minlength=len(prev_ids)).astype(np.float64)
+    cur_area = np.bincount(cur_idx, minlength=len(cur_ids)).astype(np.float64)
+
+    both = (prev.ravel() != 0) & (cur.ravel() != 0)
+    pairs = prev_idx[both] * len(cur_ids) + cur_idx[both]
+    inter = np.bincount(pairs, minlength=len(prev_ids) * len(cur_ids)).reshape(
+        len(prev_ids), len(cur_ids)
+    )
+
+    # Drop the background row and column, which are only there because 0 sorts
+    # first among the labels np.unique found.
+    prev_from = 1 if len(prev_ids) and prev_ids[0] == 0 else 0
+    cur_from = 1 if len(cur_ids) and cur_ids[0] == 0 else 0
+    return (
+        prev_ids[prev_from:],
+        cur_ids[cur_from:],
+        inter[prev_from:, cur_from:].astype(np.float64),
+        prev_area[prev_from:],
+        cur_area[cur_from:],
+    )

skop/ops/mask/__init__.py

@@ -0,0 +1,29 @@
+"""Mask detection: given boxes, draw the object inside each one.
+
+Stage two of a detect-then-segment workflow, fed by ``skop.ops.detect``. The
+distinction from ``skop.ops.segment`` is what decides where the objects are:
+Cellpose and StarDist decide for themselves and return a label image, while a
+mask detector is *told*, by a prompt, and returns one mask per prompt. That
+separation is what lets a workflow pair any detector with any mask detector,
+which ``skop.ops.workflows.mask`` does.
+
+They return masks rather than labels because SAM-family masks overlap -- a
+cell and the debris on it, a coin behind a coin -- and a label image cannot
+say so. ``ophelper.masks`` holds the projections onto something a viewer shows.
+
+Every op here shares a signature and returns ``Masks``, so they are
+substitutable. One op per environment, since ``@op(env=...)`` is fixed per
+function: ``microsam_masks`` in the shared 'pytorch' environment, because
+micro_sam is a conda package that coexists with every other torch library
+there, and ``mobilesam_masks`` in segment-everything's own, alongside the
+object-aware detector it was trained to pair with and inside the vendored
+tree it cannot leave.
+"""
+
+from __future__ import annotations
+
+from ._result import Masks
+from .microsam import microsam_masks
+from .mobilesam import mobilesam_masks
+
+__all__ = ["Masks", "microsam_masks", "mobilesam_masks"]

skop/ops/mask/_result.py

@@ -0,0 +1,40 @@
+"""What every mask detector returns.
+
+One class, shared by every mask detector op, so that "these ops are
+substitutable" is a fact rather than a convention -- a workflow choosing
+between them gets the same fields and the same output names either way. The
+same reasoning as ``skop.ops.detect._result``, one stage downstream.
+
+Two fields rather than one, because a mask detector does not always answer
+every prompt. SAM returns an empty prediction for a box it cannot make sense
+of, both implementations drop those, and so N out is not N in. Without the
+surviving prompts there is no way back to which detection each mask came
+from.
+
+The masks are a stack rather than a label image because SAM masks overlap,
+and rather than the list of annotation dicts the SAM ecosystem passes around
+because that cannot cross the Appose boundary: ``oprunner._codec`` rejects
+``bool`` outright, and N dicts of arrays would become N shared-memory blocks
+where one stack is one. See docs/design/0008-mask-detector-ops.md.
+
+No area, no stability score, no predicted IoU. Area is
+``masks.sum(axis=(1, 2))`` and does not need carrying; the other two are
+MobileSAM's and have no micro_sam counterpart, so they cannot live in a shared
+return type until docs/spec/per-object-features.md says how per-object
+values travel. The same reason ``Boxes`` carries no scores.
+"""
+
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import numpy as np
+
+from opspec.types import BoxesOf, MasksOf
+
+
+class Masks(NamedTuple):
+    #: (N, Y, X) uint8, 0 or 1, one object per plane. These may overlap.
+    masks: MasksOf[np.ndarray]
+    #: (N, 4) as [min_y, min_x, max_y, max_x] -- the prompt each mask answers.
+    boxes: BoxesOf[np.ndarray]

skop/ops/mask/microsam.py

@@ -0,0 +1,114 @@
+"""Box-prompted segmentation with micro_sam.
+
+micro_sam is Segment Anything finetuned for microscopy, by the Computational
+Cell Analytics group at Göttingen. The finetuning is the point: plain SAM was
+trained on photographs and is unremarkable on a field of nuclei, where
+``vit_b_lm`` was trained on light microscopy and is not.
+
+This op prompts it with boxes -- stage two of the detect-then-segment workflow
+in docs/spec/workflow-ops.md, fed by an op from ``skop.ops.detect``.
+
+Ported rather than wrapped, unlike its MobileSAM sibling: micro_sam is a
+normal conda package, and the code between it and here is a loop. The original
+is segment_everything/mask_detectors/microsam.py, which drove micro_sam
+through ``sam_annotator._state.AnnotatorState`` -- a singleton belonging to
+micro_sam's napari plugin. The public path below does the same work without
+reaching into another package's private module.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+
+import numpy as np
+
+import ophelper.boxes
+import ophelper.masks
+from opspec.op import op
+from opspec.runtime import progress
+from opspec.types import BoxesOf, ImageOf
+
+from ._result import Masks
+
+
+class PretrainedModel(Enum):
+    """micro_sam's finetuned checkpoints, by the imaging they were trained on.
+
+    All ViT-B. Larger encoders (``vit_l``, ``vit_h``) exist for some of these
+    and are better and slower; they are not offered here until someone wants
+    one.
+    """
+
+    light_microscopy = "vit_b_lm"
+    electron_microscopy = "vit_b_em_organelles"
+    histopathology = "vit_b_histopathology"
+    #: Plain Segment Anything, unfinetuned -- the baseline to compare against.
+    natural_image = "vit_b"
+
+
+@op(env="pytorch")
+def microsam_masks(
+    image: ImageOf[np.ndarray],
+    boxes: BoxesOf[np.ndarray],
+    model: PretrainedModel = PretrainedModel.light_microscopy,
+) -> Masks:
+    """Segment one object per bounding box.
+
+    Args:
+        image: Image to segment. 2-D or RGB. micro_sam normalizes it to 8-bit
+            itself, by min and max -- so a single hot pixel costs contrast
+            everywhere else, and clipping such an image first is worth it.
+        boxes: (N, 4) as [min_y, min_x, max_y, max_x], from a detector op, a
+            Shapes layer, or ``ophelper.boxes.from_labels``.
+        model: Which finetuned checkpoint to prompt.
+
+    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 SAM finds nothing in are dropped from both, so this may
+    be shorter than what went in.
+
+    Note: the weights are ~375 MB, downloaded into micro_sam's own cache on
+    first use, which MICROSAM_CACHEDIR relocates.
+    """
+    # Checked before the imports: a detector that found nothing is an ordinary
+    # workflow outcome, and there is no reason to load a 375 MB model to
+    # answer no prompts.
+    prompts = ophelper.boxes.as_boxes(boxes)
+    if len(prompts) == 0:
+        return Masks(ophelper.masks.empty(image.shape[:2]), ophelper.boxes.EMPTY.copy())
+
+    from micro_sam.prompt_based_segmentation import segment_from_box
+    from micro_sam.util import get_sam_model, precompute_image_embeddings
+
+    progress(f"Loading {model.value}")
+    predictor = get_sam_model(model_type=model.value)
+
+    # 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("Computing image embeddings")
+    embeddings = precompute_image_embeddings(predictor, image, ndim=2, verbose=False)
+
+    # micro_sam's own box prompts are [min_y, min_x, max_y, max_x]: its
+    # _process_box does the swap to SAM's xyxy internally. That is skop's
+    # canonical order, so nothing is converted here.
+    found = []
+    kept = []
+    for i, box in enumerate(prompts):
+        progress(f"Segmenting object {i + 1} of {len(prompts)}", i, len(prompts))
+        mask = np.squeeze(segment_from_box(predictor, box, image_embeddings=embeddings))
+        # SAM answers some prompts with nothing at all. Dropping those is why
+        # the surviving boxes are returned alongside the masks.
+        if not mask.any():
+            continue
+        found.append(mask.astype(np.uint8))
+        kept.append(i)
+
+    progress(f"Segmented {len(found)} of {len(prompts)} objects")
+    if not found:
+        return Masks(ophelper.masks.empty(image.shape[:2]), ophelper.boxes.EMPTY.copy())
+    return Masks(np.stack(found), prompts[kept])