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

skop/ops/segment/cellpose3.py

@@ -0,0 +1,245 @@
+"""Cellpose 3, the model zoo that came before CellposeSAM.
+
+Cellpose 4 replaced a shelf of specialised models -- cyto3, nuclei, and the
+ones trained on bacteria, yeast and tissue -- with one CellposeSAM model that
+is better on most things and not on everything. ``skop.ops.segment.cellpose``
+is that one; this is the shelf, kept because a model trained on your organism
+still beats a general one that never saw it, and because published results
+were produced with these.
+
+A separate op in a separate environment rather than a version parameter on
+one op. The two are not API-compatible: v3's ``models.Cellpose`` became v4's
+``models.CellposeModel``, and ``eval`` lost a return value on the way. The
+old ``CellposeSegmenter`` in napari-ai-lab carried both, branching on
+``cellpose.version`` at every call site -- which is exactly the runtime
+version sniffing an environment per op is meant to replace. Pinning the
+environment to ``cellpose <4`` means the code below has one path.
+
+``models.Cellpose``, not ``models.CellposeModel``, is the important detail
+for the built-in zoo. In v3 they are different classes: ``Cellpose`` bundles
+a SizeModel, so ``diameter=None`` makes it estimate the diameter before
+segmenting. ``CellposeModel`` has no such thing and would silently use its
+default. That is also why v3's ``eval`` returns four values where v4 returns
+three -- the fourth is the diameter it estimated.
+
+A model you trained yourself has to use ``CellposeModel`` anyway, since a
+SizeModel is trained separately and a finetuned checkpoint does not carry
+one. So the op has both classes, picked by whether ``pretrained_model`` is
+set, and what changes with it is what "estimate the diameter" can mean --
+see ``diameter`` below.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+from functools import partial
+from pathlib import Path
+from typing import Annotated
+
+import numpy as np
+
+from opspec.op import Axes, op
+from opspec.runtime import progress
+from opspec.types import ImageOf, LabelsOf
+
+from .._util import to_gray
+
+
+class PretrainedModel(Enum):
+    """Cellpose 3's built-in models, by what they were trained on.
+
+    The generalist pair first. The rest are the specialists, each finetuned
+    on one imaging domain and worth reaching for when the generalist misses
+    -- which is the reason this op exists at all.
+    """
+
+    cyto3 = "cyto3"
+    nuclei = "nuclei"
+    bacteria_phase = "bact_phase_cp3"
+    bacteria_fluorescence = "bact_fluor_cp3"
+    yeast_phase = "yeast_PhC_cp3"
+    yeast_brightfield = "yeast_BF_cp3"
+    tissue = "tissuenet_cp3"
+    live_cell = "livecell_cp3"
+    deepbacs = "deepbacs_cp3"
+
+
+class CytoplasmChannel(Enum):
+    """Which colour the cells to segment are in.
+
+    Cellpose 3 states this as a number -- 0 for grey, then 1, 2, 3 for red,
+    green and blue -- which is one-based, is not an axis index, and is easy
+    to transpose by accident. Named instead.
+
+    ``grayscale`` is the default and collapses the colours into one plane.
+    """
+
+    grayscale = "grayscale"
+    red = "red"
+    green = "green"
+    blue = "blue"
+
+
+class NucleusChannel(Enum):
+    """Which colour the nuclei are in, if a second channel has them.
+
+    The optional half of Cellpose 3's channel pair: a nuclear stain helps it
+    separate cells that touch. ``none`` is the default and means there isn't
+    one.
+    """
+
+    none = "none"
+    red = "red"
+    green = "green"
+    blue = "blue"
+
+
+#: Cellpose's own numbering for the two channel parameters. One-based over
+#: red/green/blue, with 0 meaning grey in the first slot and absent in the
+#: second.
+_CHANNEL_NUMBER = {"grayscale": 0, "none": 0, "red": 1, "green": 2, "blue": 3}
+
+
+@op(env="cellpose3")
+def cellpose3(
+    image: Annotated[ImageOf[np.ndarray], Axes("y", "x", "c?")],
+    model: PretrainedModel = PretrainedModel.cyto3,
+    pretrained_model: Path | None = None,
+    cytoplasm_channel: CytoplasmChannel = CytoplasmChannel.grayscale,
+    nucleus_channel: NucleusChannel = NucleusChannel.none,
+    diameter: Annotated[
+        float,
+        {"widget_type": "FloatSpinBox", "min": 0.0, "max": 1000.0, "step": 1.0},
+    ] = 30.0,
+    flow_threshold: Annotated[
+        float,
+        {"widget_type": "FloatSlider", "min": 0.0, "max": 3.0, "step": 0.05},
+    ] = 0.4,
+    cellprob_threshold: Annotated[
+        float,
+        {"widget_type": "FloatSlider", "min": -6.0, "max": 6.0, "step": 0.1},
+    ] = 0.0,
+    niter: int = 200,
+    min_size: int = 15,
+    normalize: bool = True,
+    use_gpu: bool = True,
+) -> LabelsOf[np.ndarray]:
+    """Segment cells with one of Cellpose 3's pretrained models.
+
+    Args:
+        image: Plane to segment. A trailing RGB(A) axis is collapsed. A
+            caller naming its axes may hand this a stack instead, and it is
+            segmented plane by plane.
+        model: Which pretrained model to run. cyto3 is the generalist; the
+            specialists are worth trying when it misses on your organism.
+            Ignored when ``pretrained_model`` is set.
+        pretrained_model: A model you trained yourself, which takes
+            precedence over ``model``. It must be a Cellpose 1-3 model, not
+            a CPSAM one -- those load in ``cellpose``, and
+            ``skop.models.cellpose_flavor`` tells the two apart from the
+            file.
+        cytoplasm_channel: Which colour the cells are in. Left at
+            ``grayscale``, the colours are collapsed into one plane before
+            Cellpose sees them -- which is right for a single-channel image
+            and wrong for a stain that only shows in one colour.
+        nucleus_channel: Which colour the nuclei are in, when a nuclear
+            stain is present to help separate touching cells. Ignored unless
+            ``cytoplasm_channel`` names a colour too: Cellpose reads the
+            pair, and a grey first channel makes it drop the second.
+        diameter: Expected cell diameter in pixels, where 0 means "work it
+            out". What that costs depends on the model. For the built-in
+            zoo it runs a size model over the image first, which is the
+            right choice when object size varies between images. A model of
+            your own has no size model, so 0 falls back to the diameter it
+            was trained at -- right when your images resemble your training
+            set, and worth overriding when they do not.
+        flow_threshold: Maximum allowed flow error per mask. Higher accepts
+            more irregular shapes; lower discards them.
+        cellprob_threshold: Cell probability cutoff. Lower finds more, and
+            grows the masks it finds.
+        niter: Dynamics iterations. More is slower and helps long or
+            branched cells, whose pixels need further to travel.
+        min_size: Discard masks smaller than this many pixels. -1 keeps
+            everything.
+        normalize: Whether to percentile-normalize the plane first. Cellpose
+            does this per call, so a caller running it slice by slice over a
+            stack gets each plane stretched to its own range -- which turns a
+            faint plane at the top of a volume into a bright one full of
+            detections. Turn it off and normalize the volume beforehand when
+            that matters; see ``skop.ops.workflows.segment``.
+        use_gpu: Whether to use the GPU when one is available. Falls back to
+            CPU rather than failing if the GPU model cannot be built.
+
+    Returns:
+        A label image, one integer per detected cell.
+
+    Note: weights are downloaded into Cellpose's own cache on first use.
+    """
+    from cellpose import models
+
+    # Naming a colour means Cellpose picks the channels itself, so it has to
+    # be given them: collapsing to grey first leaves nothing to pick from,
+    # and Cellpose does not complain -- its reshape short-circuits on a
+    # single-channel image and never reads `channels` at all.
+    select = {}
+    if cytoplasm_channel is not CytoplasmChannel.grayscale:
+        select["channels"] = [
+            _CHANNEL_NUMBER[cytoplasm_channel.value],
+            _CHANNEL_NUMBER[nucleus_channel.value],
+        ]
+        plane = np.asarray(image)
+        progress(
+            f"Segmenting the {cytoplasm_channel.value} channel"
+            + (
+                f", nuclei in {nucleus_channel.value}"
+                if nucleus_channel is not NucleusChannel.none
+                else ""
+            )
+        )
+    else:
+        plane = to_gray(image)
+
+    # int32 reaches eval from label-image arithmetic and upstream ops, and
+    # Cellpose's normalization does not handle it. Carried over from
+    # napari-ai-lab's CellposeSegmenter, which hit it.
+    if plane.dtype == np.int32:
+        plane = plane.astype(np.float32)
+
+    if pretrained_model is not None:
+        # CellposeModel, since a finetuned checkpoint carries no SizeModel.
+        name = Path(pretrained_model).name
+        build = partial(models.CellposeModel, pretrained_model=str(pretrained_model))
+    else:
+        name = model.value
+        build = partial(models.Cellpose, model_type=model.value)
+
+    progress(f"Loading Cellpose 3 model {name}")
+    try:
+        net = build(gpu=use_gpu)
+    except (AttributeError, ValueError, TypeError) as exc:
+        # A GPU that torch can see but Cellpose cannot use fails here rather
+        # than at eval. Segmenting slowly beats not segmenting.
+        progress(f"GPU model unavailable ({exc}); falling back to CPU")
+        net = build(gpu=False)
+
+    if diameter <= 0 and pretrained_model is not None:
+        # No size model to ask, so the diameter the model was trained at is
+        # the best available answer. Cellpose would otherwise quietly use
+        # CellposeModel's default of 30.
+        diameter = float(getattr(net, "diam_labels", 0.0) or 0.0)
+        progress(f"Using the model's trained diameter, {diameter:.1f} px")
+
+    progress("Running Cellpose 3")
+    # Four returns in v3 -- masks, flows, styles, diams -- where v4 gives
+    # three. Indexed rather than unpacked so the difference cannot bite.
+    result = net.eval(
+        plane,
+        diameter=diameter if diameter > 0 else None,
+        flow_threshold=flow_threshold,
+        cellprob_threshold=cellprob_threshold,
+        niter=niter,
+        min_size=min_size,
+        normalize=normalize,
+        **select,
+    )
+    return np.asarray(result[0]).astype(np.uint16)

skop/ops/segment/cellpose4.py

@@ -0,0 +1,127 @@
+"""Cellpose 4 segmentation, on its transformer models.
+
+Four models, two backbones: CPSAM and the DINOv3 one added in June 2026.
+Same call and same finetuning for all four, so one op with a model choice.
+
+Shares the 'pytorch' environment (0002). Cellpose 3 cannot: version 4
+replaced the zoo and moved the API, so it has an environment of its own.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+from pathlib import Path
+from typing import Annotated
+
+import numpy as np
+
+from opspec.op import Axes, op
+from opspec.runtime import progress
+from opspec.types import ImageOf, LabelsOf
+
+from .._util import channel_axis, to_gray
+
+
+class PretrainedModel(Enum):
+    """Cellpose 4's built-ins, newest of each backbone first. cpdino_vitb is
+    ViT-B, about a third the size of the others and cheaper to finetune."""
+
+    cpsam_v2 = "cpsam_v2"
+    cpsam = "cpsam"
+    cpdino = "cpdino"
+    cpdino_vitb = "cpdino-vitb"
+
+
+@op(env="pytorch")
+def cellpose4(
+    image: Annotated[ImageOf[np.ndarray], Axes("y", "x", "c?")],
+    model: PretrainedModel = PretrainedModel.cpsam_v2,
+    pretrained_model: Path | None = None,
+    diameter: Annotated[
+        float,
+        {"widget_type": "FloatSpinBox", "min": 0.0, "max": 1000.0, "step": 1.0},
+    ] = 0.0,
+    flow_threshold: Annotated[
+        float,
+        {"widget_type": "FloatSlider", "min": 0.0, "max": 3.0, "step": 0.05},
+    ] = 0.4,
+    cellprob_threshold: Annotated[
+        float,
+        {"widget_type": "FloatSlider", "min": -6.0, "max": 6.0, "step": 0.1},
+    ] = 0.0,
+    niter: int = 0,
+    min_size: int = 15,
+    normalize: bool = True,
+    collapse_channels: bool = False,
+    use_gpu: bool = True,
+) -> LabelsOf[np.ndarray]:
+    """Segment cells with Cellpose.
+
+    Args:
+        image: Plane to segment. A trailing RGB(A) axis is passed through
+            to CPSAM, which reads up to three channels itself. A caller
+            naming its axes may hand this a stack instead.
+        model: Which built-in to run. Ignored when ``pretrained_model``
+            is set.
+        pretrained_model: A Cellpose 4 model finetuned on your own data.
+            Empty runs the built-in named by ``model``. It must be a
+            Cellpose 4 model, not one from Cellpose 3 -- those load in
+            ``cellpose3`` instead -- and it carries its own backbone, so
+            ``model`` does not choose one for it.
+        diameter: Expected cell diameter in pixels; 0 lets Cellpose estimate.
+        flow_threshold: Maximum allowed flow error per mask.
+        cellprob_threshold: Cell probability cutoff; lower finds more cells.
+        niter: Dynamics iterations; 0 lets Cellpose scale it to the diameter.
+            More is slower and helps long or branched cells, whose pixels
+            need further to travel.
+        min_size: Discard masks smaller than this many pixels. -1 keeps
+            everything, which is what you want when the objects are small
+            enough that the default is doing the discarding for you.
+        normalize: Whether to percentile-normalize the plane first. Cellpose
+            does this per call, so a caller running it slice by slice over a
+            stack gets each plane stretched to its own range -- which turns a
+            faint plane at the top of a volume into a bright one full of
+            detections. Turn it off and normalize the volume beforehand when
+            that matters; see ``skop.ops.workflows.segment``.
+        collapse_channels: Average a trailing RGB(A) axis to grey before
+            segmenting. Off, because CPSAM takes 1-3 channels in any order
+            and using them beats throwing them away. On is for when one
+            colour is noise the average dilutes -- and for matching
+            ``cellpose3``, which needs the channels named or collapsed.
+        use_gpu: Whether to use the GPU, when one is available.
+
+    Returns:
+        A label image, one integer per detected cell.
+    """
+    from cellpose import models
+
+    if collapse_channels:
+        plane, axis = to_gray(image), None
+    else:
+        plane, axis = np.asarray(image), channel_axis(image)
+
+    # Passed only when set, rather than relying on what the version in this
+    # environment treats as "no model" -- it has been None and False.
+    # Cellpose 4 reads `pretrained_model` as a name or a path.
+    if pretrained_model is not None:
+        source = str(pretrained_model)
+        progress(f"Loading Cellpose model {Path(pretrained_model).name}")
+    else:
+        source = model.value
+        progress(f"Loading Cellpose model {source}")
+    net = models.CellposeModel(gpu=use_gpu, pretrained_model=source)
+
+    progress("Running Cellpose")
+    result = net.eval(
+        plane,
+        channel_axis=axis,
+        diameter=diameter if diameter > 0 else None,
+        flow_threshold=flow_threshold,
+        cellprob_threshold=cellprob_threshold,
+        niter=niter if niter > 0 else None,
+        min_size=min_size,
+        normalize=normalize,
+    )
+    # NB: eval returns (masks, flows, styles) or (masks, flows, styles, diams),
+    # depending on the Cellpose version.
+    return np.asarray(result[0]).astype(np.uint16)

skop/ops/segment/instanseg.py

@@ -0,0 +1,162 @@
+"""InstanSeg segmentation: nuclei, and cells around them.
+
+InstanSeg (instanseg-torch) is plain torch, so it shares the 'pytorch'
+environment (0002). Its fluorescence model takes any number of channels, in
+any order, and labels nuclei and whole cells in one pass.
+
+The scripts InstanSeg ships with also read files -- Bio-Formats through a
+JVM, whole-slide readers -- and none of that is here. The op is handed an
+array; reading the file is the caller's job.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Annotated, NamedTuple
+
+import numpy as np
+
+from opspec.op import Axes, op
+from opspec.runtime import progress
+from opspec.types import ImageOf, LabelsOf
+
+
+class PretrainedModel(Enum):
+    """InstanSeg's published models."""
+
+    fluorescence_nuclei_and_cells = "fluorescence_nuclei_and_cells"
+    brightfield_nuclei = "brightfield_nuclei"
+    single_channel_nuclei = "single_channel_nuclei"
+
+
+class Target(Enum):
+    """What to label, for a model that finds both."""
+
+    nuclei_and_cells = "all_outputs"
+    nuclei = "nuclei"
+    cells = "cells"
+
+
+class NucleiAndCells(NamedTuple):
+    #: One label per nucleus.
+    nuclei: LabelsOf[np.ndarray] | None
+    #: One label per cell, numbered to match its nucleus. None for a model
+    #: that only finds nuclei, or when only nuclei were asked for.
+    cells: LabelsOf[np.ndarray] | None
+
+
+#: Bigger than this, in pixels along either axis, and the image is run in
+#: tiles rather than whole.
+_WHOLE = 2048
+
+
+@op(env="pytorch")
+def instanseg(
+    image: Annotated[ImageOf[np.ndarray], Axes("y", "x", "c?")],
+    cells: Annotated[ImageOf[np.ndarray], Axes("y", "x", "c?")] | None = None,
+    model: PretrainedModel = PretrainedModel.fluorescence_nuclei_and_cells,
+    target: Target = Target.nuclei_and_cells,
+    nuclei_channel: Annotated[int, {"min": 0, "max": 64}] = 0,
+    cells_channel: Annotated[int, {"min": 0, "max": 64}] = 0,
+    pixel_size: Annotated[
+        float,
+        {"widget_type": "FloatSpinBox", "min": 0.0, "max": 10.0, "step": 0.05},
+    ] = 0.0,
+    use_gpu: bool = True,
+) -> NucleiAndCells:
+    """Label nuclei and cells with InstanSeg.
+
+    Args:
+        image: Plane to segment: the nuclei, or every channel when *cells*
+            is not given. A trailing channel axis is passed to the model,
+            which takes any number of fluorescence channels; the brightfield
+            model wants RGB.
+        cells: The cells -- membrane or cytoplasm -- when they are a layer
+            of their own. Goes to the model after the nuclei.
+        model: Which published model to run. Downloaded on first use.
+        target: Nuclei, cells or both, for the model that finds both.
+        nuclei_channel: The nuclei's channel in *image*, counting from 1;
+            0 for all of it.
+        cells_channel: The cells' channel, counting from 1: in *cells* when
+            it is given, else in *image*. 0 for all of *cells*, or with no
+            *cells* and no nuclei channel either, all of *image*. So: two
+            layers, 0 and 0; one layer of channels, given twice or once,
+            the two channel numbers.
+        pixel_size: The image's pixel size in microns. InstanSeg was trained
+            at about 0.5, and rescales the image to that; 0 runs it at the
+            size it is, which is right only if the pixels are near 0.5 um.
+        use_gpu: Whether to use the GPU, when one is available.
+
+    Returns:
+        nuclei: Labels, one per nucleus.
+        cells: Labels, one per cell, or None when there are none to give.
+
+    Note: the weights are downloaded into skop's asset cache on first use.
+    """
+    import os
+
+    import torch
+
+    from ophelper.assets import cache_dir
+
+    # InstanSeg downloads into its own package folder unless told otherwise.
+    os.environ["INSTANSEG_BIOIMAGEIO_PATH"] = str(cache_dir() / "instanseg")
+    from instanseg import InstanSeg
+
+    device = "cuda" if use_gpu and torch.cuda.is_available() else "cpu"
+    progress(f"Loading {model.value} on {device}")
+    net = InstanSeg(model.value, device=device, verbosity=0)
+
+    def channels_first(array: np.ndarray) -> np.ndarray:
+        # InstanSeg wants channels first; a plane with none gets one.
+        return array[np.newaxis] if array.ndim == 2 else np.moveaxis(array, -1, 0)
+
+    def pick(planes: np.ndarray, channel: int, name: str) -> np.ndarray:
+        if not channel:
+            return planes
+        if channel > len(planes):
+            raise ValueError(
+                f"{name} channel {channel} asked for, but it has {len(planes)}; "
+                "channels count from 1"
+            )
+        return planes[channel - 1 : channel]
+
+    if cells is None:
+        # One input: the channel numbers pick within it, nuclei first.
+        planes = channels_first(image)
+        chosen = list(dict.fromkeys(c for c in (nuclei_channel, cells_channel) if c))
+        if chosen:
+            planes = np.concatenate([pick(planes, c, "image") for c in chosen])
+    else:
+        if cells.shape[:2] != image.shape[:2]:
+            raise ValueError(
+                f"cells is {cells.shape[0]} x {cells.shape[1]}, but image is "
+                f"{image.shape[0]} x {image.shape[1]}"
+            )
+        planes = np.concatenate(
+            [
+                pick(channels_first(image), nuclei_channel, "image"),
+                pick(channels_first(cells), cells_channel, "cells"),
+            ]
+        )
+    planes = np.ascontiguousarray(planes, dtype=np.float32)
+    kwargs = {
+        "pixel_size": pixel_size or None,
+        "target": target.value,
+        "return_image_tensor": False,
+    }
+
+    progress(f"Segmenting {image.shape[0]} x {image.shape[1]}")
+    if max(image.shape[:2]) <= _WHOLE:
+        labels = net.eval_small_image(planes, **kwargs)
+    else:
+        labels = net.eval_medium_image(planes, tile_size=512, **kwargs)
+    labels = labels.squeeze(0).cpu().numpy().astype(np.int32)
+
+    # One plane per output: nuclei then cells for the model that finds both
+    # and was asked for both, else just the one asked for.
+    if len(labels) == 2:
+        return NucleiAndCells(nuclei=labels[0], cells=labels[1])
+    if target is Target.cells:
+        return NucleiAndCells(nuclei=None, cells=labels[0])
+    return NucleiAndCells(nuclei=labels[0], cells=None)