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

skop/ops/segment/stardist2d.py

@@ -0,0 +1,192 @@
+"""StarDist 2D inference, pretrained or from a model on disk.
+
+Ported from src/imgops/implementations/stardist2d.py. Shares the
+'stardist-tf' environment with starfun3d.py.
+
+Two ops for the pretrained models rather than one with a ``model`` switch,
+because they do not accept the same thing: the H&E model is trained on stain
+colour and always wants a channel axis, while the fluorescence model wants a
+single channel and collapses one if it is given it. That is a difference in
+what the op *consumes*, so it belongs in ``Axes`` where a front end can see
+it, not in an enum parameter where it is invisible until StarDist raises.
+
+``stardist2d_custom`` is the third, and takes a directory. It cannot make the
+same declaration -- a model on disk is single- or multi-channel depending on
+how it was trained -- so it declares ``c?`` and reads the model's own
+``config.json`` to decide. A caller with many trained models points this one
+op at each of them in turn; see design 0011.
+"""
+
+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 ImageOf, LabelsOf
+
+from .._util import to_gray
+
+_FLUO = "2D_versatile_fluo"
+_HE = "2D_versatile_he"
+
+
+@op(env="stardist-tf")
+def stardist2d_fluo(
+    image: Annotated[ImageOf[np.ndarray], Axes("y", "x", "c?")],
+    prob_thresh: Annotated[
+        float,
+        {"widget_type": "FloatSlider", "min": 0.0, "max": 1.0, "step": 0.05},
+    ] = 0.5,
+    nms_thresh: Annotated[
+        float,
+        {"widget_type": "FloatSlider", "min": 0.0, "max": 1.0, "step": 0.05},
+    ] = 0.4,
+    normalize: bool = True,
+) -> LabelsOf[np.ndarray]:
+    """Detect objects in a fluorescence image with pretrained StarDist.
+
+    Args:
+        image: Plane to segment. A trailing RGB(A) axis is collapsed, since
+            this model wants a single channel.
+        prob_thresh: Object probability threshold.
+        nms_thresh: Non-maximum suppression threshold.
+        normalize: Whether to percentile-normalize the input first.
+
+    Returns:
+        A label image, one integer per detected object.
+
+    Note: the model weights are downloaded on first use.
+    """
+    return _predict(to_gray(image), _FLUO, prob_thresh, nms_thresh, normalize)
+
+
+@op(env="stardist-tf")
+def stardist2d_he(
+    image: Annotated[ImageOf[np.ndarray], Axes("y", "x", "c")],
+    prob_thresh: Annotated[
+        float,
+        {"widget_type": "FloatSlider", "min": 0.0, "max": 1.0, "step": 0.05},
+    ] = 0.5,
+    nms_thresh: Annotated[
+        float,
+        {"widget_type": "FloatSlider", "min": 0.0, "max": 1.0, "step": 0.05},
+    ] = 0.4,
+    normalize: bool = True,
+) -> LabelsOf[np.ndarray]:
+    """Detect objects in an H&E-stained image with pretrained StarDist.
+
+    Args:
+        image: Plane to segment, with its colour channels intact -- this model
+            is trained on stain colour, so the channel axis is required.
+        prob_thresh: Object probability threshold.
+        nms_thresh: Non-maximum suppression threshold.
+        normalize: Whether to percentile-normalize the input first.
+
+    Returns:
+        A label image, one integer per detected object.
+
+    Note: the model weights are downloaded on first use.
+    """
+    return _predict(image, _HE, prob_thresh, nms_thresh, normalize)
+
+
+@op(env="stardist-tf")
+def stardist2d_custom(
+    image: Annotated[ImageOf[np.ndarray], Axes("y", "x", "c?")],
+    model_dir: Annotated[str, {"widget_type": "FileEdit", "mode": "d"}] = "",
+    prob_thresh: Annotated[
+        float,
+        {"widget_type": "FloatSlider", "min": 0.0, "max": 1.0, "step": 0.05},
+    ] = 0.5,
+    nms_thresh: Annotated[
+        float,
+        {"widget_type": "FloatSlider", "min": 0.0, "max": 1.0, "step": 0.05},
+    ] = 0.4,
+    normalize: bool = True,
+) -> LabelsOf[np.ndarray]:
+    """Detect objects with a StarDist 2D model loaded from disk.
+
+    The counterpart to the two pretrained ops: same prediction, but the model
+    comes from a directory rather than StarDist's download cache. That is what
+    a model trained by ``skop.ops.train`` -- or by anything else that writes a
+    StarDist model -- is loaded with.
+
+    Args:
+        image: Plane to segment. The trailing channel axis is collapsed when
+            the model expects a single channel, and kept when it does not;
+            the model's own config decides, not the caller.
+        model_dir: Directory holding ``config.json`` and the weights. The
+            directory name is the model name, as StarDist stores it.
+        prob_thresh: Object probability threshold.
+        nms_thresh: Non-maximum suppression threshold.
+        normalize: Whether to percentile-normalize the input first.
+
+    Returns:
+        A label image, one integer per detected object.
+    """
+    import json
+    import os
+
+    from stardist.models import StarDist2D
+
+    if not model_dir:
+        raise ValueError(
+            "model_dir is required -- this op loads a model from disk. Use "
+            "stardist2d_fluo or stardist2d_he for the pretrained models."
+        )
+
+    model_dir = os.path.normpath(model_dir)
+    basedir, name = os.path.split(model_dir)
+    config_path = os.path.join(model_dir, "config.json")
+    if not os.path.isfile(config_path):
+        raise ValueError(f"No StarDist model in {model_dir}: no config.json")
+
+    with open(config_path) as f:
+        n_channel_in = json.load(f).get("n_channel_in", 1)
+
+    x = image if n_channel_in > 1 else to_gray(image)
+
+    progress(f"Loading model {name}")
+    net = StarDist2D(None, name=name, basedir=basedir)
+
+    return _run(net, x, prob_thresh, nms_thresh, normalize)
+
+
+def _predict(
+    x: np.ndarray,
+    model: str,
+    prob_thresh: float,
+    nms_thresh: float,
+    normalize: bool,
+) -> np.ndarray:
+    """Run one pretrained StarDist 2D model over an already-shaped array."""
+    from stardist.models import StarDist2D
+
+    progress(f"Loading pretrained model {model}")
+    net = StarDist2D.from_pretrained(model)
+    return _run(net, x, prob_thresh, nms_thresh, normalize)
+
+
+def _run(
+    net,
+    x: np.ndarray,
+    prob_thresh: float,
+    nms_thresh: float,
+    normalize: bool,
+) -> np.ndarray:
+    """Normalize if asked, then predict instances with an already-loaded net."""
+    from csbdeep.utils import normalize as normalize_percentile
+
+    if normalize:
+        x = normalize_percentile(x, 1, 99.8)
+
+    progress("Predicting instances")
+    labels, _ = net.predict_instances(
+        x,
+        prob_thresh=prob_thresh,
+        nms_thresh=nms_thresh,
+    )
+    return np.asarray(labels).astype(np.uint16)

skop/ops/segment/starfun3d.py

@@ -0,0 +1,104 @@
+"""3D nuclei segmentation with pretrained StarDist models.
+
+Ported from https://github.com/ctrueden/starfun3d. Shares the 'stardist-tf'
+environment with stardist2d.py.
+
+The models come from Galindo et al. (2023), "3D Nuclei Segmentation By
+Combining GAN Based Image Synthesis and Existing 3D Manual Annotations",
+https://doi.org/10.1101/2023.12.06.570366, published at
+https://zenodo.org/records/10518151. They total 36 MB, so they are fetched
+into skop's asset cache on first use rather than carried in the repository.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Annotated, NamedTuple
+
+import numpy as np
+
+from opspec.op import op
+from opspec.runtime import progress
+from opspec.types import ImageOf, LabelsOf, PointsOf
+
+_ZENODO = "https://zenodo.org/records/10518151/files/model_{name}.zip?download=1"
+
+
+class Model(Enum):
+    """Pretrained StarDist3D models, one per acquisition modality.
+
+    Average nucleus size in the training data, in pixels:
+    confocal [x=39, y=39, z=7], sospim [x=27, y=28, z=10],
+    spinning [x=39, y=39, z=7].
+    """
+
+    confocal = "confocal"
+    sospim = "sospim"
+    spinning = "spinning"
+
+
+class Nuclei(NamedTuple):
+    labels: LabelsOf[np.ndarray]
+    points: PointsOf[np.ndarray]
+
+
+@op(env="stardist-tf")
+def segment_nuclei(
+    image: ImageOf[np.ndarray],
+    model: Model = Model.confocal,
+    prob_thresh: Annotated[
+        float | None,
+        {"widget_type": "FloatSlider", "min": 0.0, "max": 1.0, "step": 0.05},
+    ] = None,
+    nms_thresh: Annotated[
+        float | None,
+        {"widget_type": "FloatSlider", "min": 0.0, "max": 1.0, "step": 0.05},
+    ] = None,
+    normalize: bool = True,
+) -> Nuclei:
+    """Segment nuclei in a 3D volume.
+
+    Args:
+        image: Volume to segment, in (Z, Y, X) order.
+        model: Which pretrained model to use; pick the one matching how the
+            data was acquired.
+        prob_thresh: Object probability threshold. None uses the model's own.
+        nms_thresh: Non-maximum suppression threshold. None uses the model's.
+        normalize: Whether to percentile-normalize the input first.
+
+    Returns:
+        A label volume, and the detected nucleus centers as (N, 3) points.
+    """
+    from csbdeep.utils import normalize as normalize_percentile
+    from stardist.models import StarDist3D
+
+    from ophelper import assets
+
+    model_dir = assets.unzip_from_url(
+        _ZENODO.format(name=model.value),
+        f"starfun3d/{model.value}",
+        marker="config.json",
+    )
+
+    progress(f"Loading model {model.value}")
+    # NB: StarDist3D locates a model as basedir/name, so the model's own
+    # directory name must be used -- the original starfun3d passed a constant
+    # here, and so loaded the confocal model whichever model was asked for.
+    net = StarDist3D(None, name=model_dir.name, basedir=str(model_dir.parent))
+
+    x = normalize_percentile(image, 1, 99.8) if normalize else image
+
+    predict_kwargs = {}
+    if prob_thresh is not None:
+        predict_kwargs["prob_thresh"] = prob_thresh
+    if nms_thresh is not None:
+        predict_kwargs["nms_thresh"] = nms_thresh
+
+    progress("Predicting instances")
+    labels, details = net.predict_instances(x, **predict_kwargs)
+
+    points = np.asarray(details.get("points", np.zeros((0, 3))))
+    return Nuclei(
+        labels=np.asarray(labels).astype(np.uint16),
+        points=points.astype(np.int32),
+    )

skop/ops/segment/unseg/__init__.py

@@ -0,0 +1,138 @@
+"""UNSEG: unsupervised segmentation of nuclei and cells.
+
+Segments tissue images carrying both a nucleus marker and a cell membrane
+marker, without training data.
+
+Ported from https://github.com/ctrueden/unseg-fiji. The algorithm lives in
+_algorithm.py, which is imported inside the op body because its own imports
+(OpenCV, scikit-learn, scikit-image) are far too heavy for discovery.
+
+  Uttam Lab. UNSEG: unsupervised segmentation of cells and their nuclei in
+  complex tissue samples.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Annotated, NamedTuple
+
+import numpy as np
+
+from opspec.op import op
+from opspec.runtime import progress
+from opspec.types import ImageOf, LabelsOf
+
+
+class DistanceTransform(Enum):
+    """How to compute distances when splitting touching objects."""
+
+    geodesic = "GDT"
+    euclidean = "DT"
+
+
+class TernaryMethod(Enum):
+    """How to resolve the three-way foreground/background/boundary decision."""
+
+    argmax = "Argmax"
+    posterior = "Posterior"
+
+
+class Segmentation(NamedTuple):
+    nuclei: LabelsOf[np.ndarray]
+    cells: LabelsOf[np.ndarray]
+    n_nuclei: int
+    n_cells: int
+
+
+@op(env="unseg-cv")
+def unseg(
+    image: ImageOf[np.ndarray],
+    channel_axis: int = 0,
+    nuclei_channel: int = 2,
+    membrane_channel: int = 0,
+    area_threshold: Annotated[int, {"min": 1, "max": 10000}] = 20,
+    convexity_threshold: Annotated[int, {"min": 0, "max": 100}] = 4,
+    cell_marker_threshold: Annotated[int, {"min": 0, "max": 255}] = 25,
+    dist_tr: DistanceTransform = DistanceTransform.geodesic,
+    sigma0: Annotated[int, {"min": 0, "max": 20}] = 3,
+    k0: Annotated[int, {"min": 0, "max": 20}] = 1,
+    r0: Annotated[int, {"min": 1, "max": 50}] = 5,
+    pct: tuple[float, float] = (0.01, 0.01),
+    nk: tuple[int, ...] = (5, 10, 20, 40),
+    t0: Annotated[float, {"min": 0.0, "max": 1.0, "step": 0.05}] = 0.5,
+    ternary_met: TernaryMethod = TernaryMethod.argmax,
+    area_ratio_threshold: Annotated[
+        float, {"min": 0.0, "max": 1.0, "step": 0.05}
+    ] = 0.65,
+    dilation_radius: Annotated[int, {"min": 1, "max": 50}] = 5,
+) -> Segmentation:
+    """Segment nuclei and cells in a two-marker tissue image.
+
+    Args:
+        image: Multichannel 2D image carrying both markers.
+        channel_axis: Which axis of ``image`` holds the channels. Fiji and
+            napari usually hand over (C, Y, X), hence the default of 0.
+        nuclei_channel: Index of the nucleus marker channel, e.g. DAPI.
+        membrane_channel: Index of the cell membrane marker channel, e.g.
+            Na+/K+ ATPase.
+        area_threshold: Smallest object area to keep, in pixels.
+        convexity_threshold: How non-convex an object may be before UNSEG
+            tries to split it.
+        cell_marker_threshold: Membrane marker intensity below which a pixel
+            is not considered part of a cell.
+        dist_tr: Distance transform used when splitting touching objects.
+        sigma0: Smoothing width for the empirical probability estimate.
+        k0: Gradient-adaptive smoothing strength.
+        r0: Radius of the local binary mask neighborhood.
+        pct: Background removal percentiles, one per marker.
+        nk: Cluster counts to try when building the global binary mask.
+        t0: Threshold on the global binary mask.
+        ternary_met: How to resolve the three-way mask.
+        area_ratio_threshold: Nucleus-to-cell area ratio above which a cell is
+            rejected.
+        dilation_radius: How far to grow nuclei when seeding cells.
+
+    Returns:
+        Nucleus and cell label images, and the count of each.
+    """
+    from ._algorithm import nuclei_cell_segmentation
+
+    nuclei_marker = np.take(image, nuclei_channel, axis=channel_axis)
+    membrane_marker = np.take(image, membrane_channel, axis=channel_axis)
+    if nuclei_marker.ndim != 2:
+        raise ValueError(
+            f"UNSEG segments 2D images, but selecting a channel of a "
+            f"{image.ndim}D image left {nuclei_marker.ndim} dimensions. "
+            f"Check channel_axis (currently {channel_axis})."
+        )
+
+    intensity = np.zeros((*nuclei_marker.shape, 2), dtype="float64")
+    intensity[:, :, 0] = nuclei_marker
+    intensity[:, :, 1] = membrane_marker
+
+    progress(f"Segmenting {nuclei_marker.shape} image")
+    nuc_seg, cel_seg, n_nuclei, n_cells = nuclei_cell_segmentation(
+        intensity,
+        area_threshold=area_threshold,
+        convexity_threshold=convexity_threshold,
+        cell_marker_threshold=cell_marker_threshold,
+        dist_tr=dist_tr.value,
+        sigma0=sigma0,
+        k0=k0,
+        r0=r0,
+        pct=list(pct),
+        nk=list(nk),
+        t0=t0,
+        ternary_met=ternary_met.value,
+        # NB: no display in a worker process.
+        visualization=False,
+        area_ratio_threshold=area_ratio_threshold,
+        dilation_radius=dilation_radius,
+    )
+
+    return Segmentation(
+        nuclei=np.asarray(nuc_seg).astype(np.int32),
+        cells=np.asarray(cel_seg).astype(np.int32),
+        n_nuclei=int(n_nuclei),
+        n_cells=int(n_cells),
+    )