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

skop/ops/deconvolve/_edges.py

@@ -0,0 +1,92 @@
+"""Edge handling: deciding where the image has data, and extending it.
+
+Padding and bad-pixel masking look like two chores but are one idea. An edge
+is a place where no data was acquired; so is a saturated pixel. Both are
+handled the same way -- held out of the HTones array, which is what tells the
+iteration how much signal each location should have received.
+
+This is identical whichever backend runs the iteration, and happens on the CPU
+either way, so both share it. The two originals ported here had drifted apart
+in exactly these steps; sharing them is what keeps that from happening again.
+"""
+
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import numpy as np
+
+from ._pad import get_next_smooth, pad, pad_psf, pad_to_largest, unpad
+
+
+class Padded(NamedTuple):
+    """Arrays at their working size, and what it takes to get back."""
+
+    image: np.ndarray
+    psf: np.ndarray
+    htones: np.ndarray
+    original_size: tuple[int, ...]
+    mask: np.ndarray | None
+    mask_values: np.ndarray | None
+
+    def crop_and_restore(self, estimate: np.ndarray) -> np.ndarray:
+        """Undo ``pad_and_mask``: crop back, and put masked pixels back.
+
+        A method rather than a free function because the padding is what
+        knows how to reverse itself -- there is no way to pair the wrong
+        result with the wrong padding.
+        """
+        if estimate.shape != self.original_size:
+            estimate = unpad(estimate, self.original_size)
+        if self.mask is not None:
+            estimate = estimate * self.mask + self.mask_values
+        return estimate
+
+
+def pad_and_mask(
+    image: np.ndarray,
+    psf: np.ndarray,
+    noncirc: bool,
+    mask: np.ndarray | None,
+    dtype=np.float64,
+) -> Padded:
+    """Hold out masked pixels, and pad the arrays to their working size.
+
+    Note:
+        When ``noncirc`` is off and the image and psf already agree, there is
+        nothing to pad and only the masking applies.
+    """
+    image = np.asarray(image, dtype=dtype)
+    psf = np.asarray(psf, dtype=dtype)
+    original_size = image.shape
+
+    mask_values = None
+    if mask is None:
+        htones = np.ones_like(image)
+    else:
+        mask = np.asarray(mask, dtype=dtype)
+        if mask.shape != original_size:
+            raise ValueError(
+                f"mask shape {mask.shape} does not match image shape {original_size}"
+            )
+        # Masked pixels are set aside and restored afterwards. Folding the mask
+        # into HTones is what makes them behave like edges.
+        mask_values = image * (1.0 - mask)
+        image = image * mask
+        htones = mask.copy()
+
+    if noncirc:
+        extended = get_next_smooth(
+            [image.shape[i] + 2 * (psf.shape[i] // 2) for i in range(image.ndim)]
+        )
+        image, _ = pad(image, extended, "constant")
+        htones, _ = pad(htones, extended, "constant")
+        psf = pad_psf(psf, extended)
+    elif image.shape != psf.shape:
+        # NB: the numpy original padded the psf up to the image, which
+        # truncates a psf wider than the image. Both backends now pad each to
+        # the larger of the two, as the cupy version did.
+        image, psf = pad_to_largest(image, psf, "constant")
+        htones, _ = pad(htones, image.shape, "constant")
+
+    return Padded(image, psf, htones, original_size, mask, mask_values)

skop/ops/deconvolve/_pad.py

@@ -0,0 +1,117 @@
+"""Padding helpers for FFT-based deconvolution.
+
+Vendored from ``tnia.deconvolution.pad`` (tnia-python), reformatted but
+semantically unchanged. Only what the ops here need came across.
+
+Private, so discovery never imports it as an op module. Pure numpy and
+``math``, which is what lets both backends share it: ``unpad`` only slices, so
+it works on a cupy array as readily as a numpy one.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+
+def _handle_prime(p: int, x: int, a: np.ndarray) -> None:
+    """Accumulate log(p) into every index reachable as a multiple of p."""
+    log = math.log(p)
+    power = p
+
+    while power <= x + a.shape[0]:
+        j = x % power
+        if j > 0:
+            j = power - j
+
+        while j < a.shape[0]:
+            a[j] += log
+            j += power
+
+        power *= p
+
+
+def next_smooth(x: int) -> int:
+    """Return the next number >= x divisible only by primes up to 7.
+
+    FFTs are dramatically faster on smooth sizes, so a padded extent is
+    rounded up to one. Based on A. Granville, *Finding smooth numbers
+    computationally*.
+
+    Authors: Johannes Schindelin, Brian Northan.
+    """
+    if x == 1:
+        return 1
+
+    z = int(16 * math.log2(x))
+    delta = 0.000001
+
+    a = np.zeros(z)
+
+    _handle_prime(2, x, a)
+    _handle_prime(3, x, a)
+    _handle_prime(5, x, a)
+    _handle_prime(7, x, a)
+
+    log = math.log(x)
+    for i in range(a.shape[0]):
+        if a[i] >= log - delta:
+            return x + i
+
+    return -1
+
+
+def get_next_smooth(size) -> tuple[int, ...]:
+    """Apply ``next_smooth`` to every element of an n-dimensional size."""
+    return tuple(next_smooth(i) for i in size)
+
+
+def pad(img: np.ndarray, paddedsize, mode: str):
+    """Pad an image up to ``paddedsize``, centered.
+
+    Returns:
+        The padded array, and the per-axis padding that was applied.
+    """
+    padding = tuple(
+        (math.ceil((i - j) / 2), math.floor((i - j) / 2))
+        for i, j in zip(paddedsize, img.shape)
+    )
+    return np.pad(img, padding, mode), padding
+
+
+def pad_psf(psf: np.ndarray, paddedsize) -> np.ndarray:
+    """Pad a PSF with zeros up to ``paddedsize``, its centre where FFTs want it.
+
+    The centre, index ``n // 2`` of the PSF, goes to index ``N // 2`` of the
+    padded array, which ``ifftshift`` then moves to 0. ``pad`` centres the
+    array instead, and for a PSF of even size that puts the centre half a
+    pixel off one way or the other depending on whether the padded size is
+    odd: the result shifts by a pixel. Untiled, every run shifts the same way;
+    tiled, tiles of different sizes disagree, by 20% or more on sharp points.
+    """
+    padding = tuple(
+        (big // 2 - small // 2, big - small - (big // 2 - small // 2))
+        for big, small in zip(paddedsize, psf.shape)
+    )
+    return np.pad(psf, padding, "constant")
+
+
+def pad_to_largest(img: np.ndarray, psf: np.ndarray, mode: str):
+    """Pad both arrays so each axis matches the larger of the two."""
+    largest = tuple(max(i, p) for i, p in zip(img.shape, psf.shape))
+    img, _ = pad(img, largest, mode)
+    return img, pad_psf(psf, largest)
+
+
+def unpad(padded, imgsize):
+    """Crop a padded array back to ``imgsize``, undoing ``pad``.
+
+    Slicing only, so this accepts a cupy array as well as a numpy one.
+    """
+    padding = tuple(
+        (math.ceil((i - j) / 2), math.floor((i - j) / 2))
+        for i, j in zip(padded.shape, imgsize)
+    )
+    slices = tuple(slice(p[0], p[0] + s) for p, s in zip(padding, imgsize))
+    return padded[slices]

skop/ops/deconvolve/richardson_lucy.py

@@ -0,0 +1,100 @@
+"""Richardson-Lucy deconvolution on the CPU, with numpy.
+
+Ported from tnia-python's ``tnia.deconvolution.richardson_lucy_np``. Where it
+differed from the cupy version, the cupy behaviour won -- see the NB comments
+below and ``richardson_lucy_cupy``, which this op is kept numerically
+equivalent to.
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+import numpy as np
+
+from opspec.op import Axes, Overlap, PeakMemory, op
+from opspec.runtime import cancel_requested, progress
+
+from ._edges import pad_and_mask
+
+# Floor for the divisions and the estimate: an iteration must not divide by
+# zero, and intensities are non-negative by construction.
+DELTA = 1e-6
+
+
+#: Non-circulant: half the PSF, per axis, how far it reaches at most; safe,
+#: if often more than the light needs. Circulant: 10 px, as clij2-fft and
+#: tnia-python use. A Tiler can set either.
+_OVERLAP = Overlap(param="psf", of="shape", scale=0.5, only_if="noncirc", otherwise=10)
+#: Non-circulant pads by half the PSF on each side; circulant doesn't pad.
+_PAD = Overlap(param="psf", of="shape", scale=0.5, only_if="noncirc")
+
+
+@op(
+    env="skimage",
+    # As richardson_lucy_cupy: the mask tiled the same as the image, z never
+    # cut. In RAM and in float64, numpy's FFT precision; 10x is a guess, not
+    # yet measured, counted on the padded size.
+    tile=("image", "mask"),
+    split=("y", "x"),
+    overlap=_OVERLAP,
+    peak_memory=PeakMemory(10, "float64", pad=_PAD),
+    merge="blend",
+)
+def richardson_lucy(
+    image: Annotated[np.ndarray, Axes("z?", "y", "x")],
+    psf: np.ndarray,
+    num_iters: Annotated[int, {"min": 1, "max": 1000}] = 10,
+    noncirc: bool = False,
+    mask: np.ndarray | None = None,
+) -> np.ndarray:
+    """Deconvolve an image by the Richardson-Lucy algorithm.
+
+    Args:
+        image: The image to deconvolve.
+        psf: The point spread function. Need not match the image's shape.
+        num_iters: How many iterations to run.
+        noncirc: Whether to use non-circulant edge handling, which pads the
+            image out rather than treating it as wrapping around.
+        mask: Optional array the shape of the image, zero at pixels that
+            should not be considered -- saturated ones, say. They are held out
+            of the deconvolution and their original values restored at the end.
+
+    Returns:
+        The deconvolved image, as float64.
+
+    Note:
+        numpy's FFT promotes float32 to float64, so passing 32-bit input to
+        save memory does not work here; the result is float64 regardless.
+    """
+    fftn, ifftn, ifftshift = np.fft.fftn, np.fft.ifftn, np.fft.ifftshift
+
+    padded = pad_and_mask(image, psf, noncirc, mask)
+    image, psf, htones = padded.image, padded.psf, padded.htones
+
+    otf = fftn(ifftshift(psf))
+    otf_conj = np.conjugate(otf)
+
+    htones = np.real(ifftn(fftn(htones) * otf_conj))
+    htones[htones < DELTA] = 1.0
+
+    # NB: the numpy original seeded the circulant case with the image itself.
+    # Both backends now start from a flat sheet at the image's mean, which is
+    # what the cupy version did and what the non-circulant case always needed.
+    estimate = np.full_like(image, image.mean())
+
+    for i in range(num_iters):
+        if cancel_requested():
+            break
+        progress(f"Iteration {i + 1} of {num_iters}", i, num_iters)
+
+        reblurred = np.real(ifftn(fftn(estimate) * otf))
+        reblurred[reblurred < DELTA] = DELTA
+        ratio = image / reblurred
+        correction = np.real(ifftn(fftn(ratio) * otf_conj))
+        correction[correction < 0] = DELTA
+        estimate = estimate * correction / htones
+        estimate[estimate < 0] = DELTA
+
+    progress("Deconvolution complete", num_iters, num_iters)
+    return padded.crop_and_restore(estimate)

skop/ops/deconvolve/richardson_lucy_cupy.py

@@ -0,0 +1,116 @@
+"""Richardson-Lucy deconvolution on an NVIDIA GPU, with cupy.
+
+Ported from tnia-python's ``tnia.deconvolution.richardson_lucy``. The
+iteration is the same one ``richardson_lucy`` runs, and the two are held
+numerically equivalent by a test; what differs is where the arrays live.
+
+Dropped in the port: the RMSE-against-truth tracking (a benchmarking concern,
+not an op's), ``do_unpad`` (an op returns something the caller can use), and
+the memory-pool diagnostics.
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+import numpy as np
+
+from opspec.op import Axes, Overlap, PeakMemory, op
+from opspec.runtime import cancel_requested, progress
+
+from ._edges import pad_and_mask
+from .richardson_lucy import DELTA
+
+#: Non-circulant: half the PSF, per axis, how far it reaches at most; safe,
+#: if often more than the light needs. Circulant: 10 px, as clij2-fft and
+#: tnia-python use. A Tiler can set either.
+_OVERLAP = Overlap(param="psf", of="shape", scale=0.5, only_if="noncirc", otherwise=10)
+#: Non-circulant pads by half the PSF on each side; circulant doesn't pad.
+_PAD = Overlap(param="psf", of="shape", scale=0.5, only_if="noncirc")
+
+
+@op(
+    env="cupy",
+    # The mask is tiled the same as the image; the PSF is given whole to every
+    # tile. z is never cut: a PSF is usually long in z, and cutting there
+    # costs the most overlap for the least gain.
+    tile=("image", "mask"),
+    split=("y", "x"),
+    overlap=_OVERLAP,
+    # On the GPU, in float32, counted on the padded size. About 13 float32
+    # buffers are live at once by the code -- image, PSF, weights, estimate,
+    # two complex OTFs, and an iteration's complex FFTs -- plus cuFFT's
+    # workspace. Not yet measured.
+    peak_memory=PeakMemory(16, "float32", pad=_PAD, device="gpu"),
+    merge="blend",
+)
+def richardson_lucy_cupy(
+    image: Annotated[np.ndarray, Axes("z?", "y", "x")],
+    psf: np.ndarray,
+    num_iters: Annotated[int, {"min": 1, "max": 1000}] = 10,
+    noncirc: bool = False,
+    mask: np.ndarray | None = None,
+) -> np.ndarray:
+    """Deconvolve an image by the Richardson-Lucy algorithm, on the GPU.
+
+    Args:
+        image: The image to deconvolve.
+        psf: The point spread function. Need not match the image's shape.
+        num_iters: How many iterations to run.
+        noncirc: Whether to use non-circulant edge handling, which pads the
+            image out rather than treating it as wrapping around.
+        mask: Optional array the shape of the image, zero at pixels that
+            should not be considered -- saturated ones, say. They are held out
+            of the deconvolution and their original values restored at the end.
+
+    Returns:
+        The deconvolved image, as float32, back in host memory.
+
+    Note:
+        cupy's FFT works in single precision, so this runs in float32 where
+        ``richardson_lucy`` runs in float64. Expect the two to agree to
+        single-precision tolerance, not exactly.
+    """
+    import cupy as cp
+
+    # Masking and padding happen on the host, so only what the iteration needs
+    # is uploaded -- and in float32, which is all cupy's FFT will give back.
+    padded = pad_and_mask(image, psf, noncirc, mask, dtype=np.float32)
+    image = cp.asarray(padded.image)
+    psf = cp.asarray(padded.psf)
+    htones = cp.asarray(padded.htones)
+
+    otf = cp.fft.fftn(cp.fft.ifftshift(psf))
+    otf_conj = cp.conjugate(otf)
+
+    htones = cp.real(cp.fft.ifftn(cp.fft.fftn(htones) * otf_conj))
+    htones[htones < DELTA] = 1.0
+
+    estimate = cp.full_like(image, float(image.mean()))
+
+    for i in range(num_iters):
+        if cancel_requested():
+            break
+        progress(f"Iteration {i + 1} of {num_iters}", i, num_iters)
+
+        reblurred = cp.real(cp.fft.ifftn(cp.fft.fftn(estimate) * otf))
+        reblurred[reblurred < DELTA] = DELTA
+        ratio = image / reblurred
+        correction = cp.real(cp.fft.ifftn(cp.fft.fftn(ratio) * otf_conj))
+        correction[correction < 0] = DELTA
+        estimate = estimate * correction / htones
+        estimate[estimate < 0] = DELTA
+
+        # Wait for the iteration actually to finish. Every cupy call above
+        # queues work and returns at once, so without this the loop issues all
+        # num_iters iterations in a few milliseconds -- reporting each -- and
+        # then blocks on the transfer at the end. The bar would sweep to 100%
+        # and sit there while the GPU did the work, and Cancel would arrive
+        # long after every iteration had been queued and so do nothing.
+        #
+        # The cost is that the queue can no longer run ahead. For FFTs this
+        # size that is small, and a progress bar that lies is worse.
+        cp.cuda.runtime.deviceSynchronize()
+
+    progress("Deconvolution complete", num_iters, num_iters)
+    return padded.crop_and_restore(cp.asnumpy(estimate))

skop/ops/detect/__init__.py

@@ -0,0 +1,28 @@
+"""Object detection, class-agnostic or with a user's trained YOLO model.
+
+FastSAM and object-aware YOLO return the same thing, so a
+caller can swap one for the other -- the first stage of a detect-then-segment
+workflow, feeding boxes to a mask detector.
+
+Those two are class-agnostic on purpose. A COCO-pretrained detector has no
+category for a cell or a coin and reports nothing; these have a single class,
+called object, and report everything.
+
+One op per environment, since ``@op(env=...)`` is fixed per function:
+``fastsam`` in the shared 'pytorch' environment, ``object_aware_yolo`` in one
+of its own, because segment-everything vendors an ultralytics fork that must
+not meet the real one.
+
+``yolo`` runs a user's detection checkpoint on overlapping tiles, returning
+boxes, confidences and class IDs. It shares the 'pytorch' environment.
+"""
+
+from __future__ import annotations
+
+from ._result import Boxes, Detections
+from .fastsam import fastsam
+from .jdll_yolo import jdll_yolo
+from .object_aware_yolo import object_aware_yolo
+from .yolo import yolo
+
+__all__ = ["Boxes", "Detections", "fastsam", "jdll_yolo", "object_aware_yolo", "yolo"]

skop/ops/detect/_result.py

@@ -0,0 +1,45 @@
+"""Return types for box detectors.
+
+``Boxes`` is shared by the class-agnostic ops, so "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.
+
+A one-field NamedTuple rather than a bare return annotation, because the field
+name becomes the output name (0001): a front end labels the layer "boxes"
+rather than "result", and adding a second output later is additive.
+
+No classes: these detectors are class-agnostic by design, and a column of
+zeros pretending to be class IDs would be worse than its absence.
+
+No scores either, for now, though every detector here computes them. A
+confidence per box is not a layer -- it is a *feature* of one, and skop has no
+way to say that yet. Returning it as a bare array made a front end guess, and
+napari guessed "image", which a 1-D array cannot be. Rather than encode a
+guess, the value is dropped until the question is answered. See
+docs/spec/per-object-features.md.
+"""
+
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import numpy as np
+
+from opspec.types import BoxesOf
+
+
+class Boxes(NamedTuple):
+    #: (N, 4) as [min_y, min_x, max_y, max_x], in image coordinates.
+    boxes: BoxesOf[np.ndarray]
+
+
+class Detections(NamedTuple):
+    """Classified boxes, with one confidence and class ID per row.
+
+    Lists keep per-object values out of image layers until skop has a
+    features role (docs/spec/per-object-features.md).
+    """
+
+    boxes: BoxesOf[np.ndarray]
+    confidences: list[float]
+    classes: list[int]

skop/ops/detect/fastsam.py

@@ -0,0 +1,103 @@
+"""Class-agnostic object detection with FastSAM.
+
+FastSAM is a YOLOv8 segmentation model trained on 2% of SA-1B. Ultralytics
+describes it as recognising and segmenting "all objects as the same class" --
+it has one class, called object, so its confidence score means "is this a
+thing" rather than "is this one of 80 COCO categories".
+
+That is the distinction that matters here. A COCO-pretrained detector
+localises coins or cells perfectly well -- box regression is class-agnostic --
+and then scores every class near zero and throws them all away below the
+confidence threshold. A single-class model keeps them.
+
+Only the boxes are returned. FastSAM also produces masks, from a prototype
+branch cropped by these same boxes; a mask detector op built on that is
+possible later (see docs/design/0007) but is not this op's job.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Annotated
+
+import numpy as np
+
+import ophelper.boxes
+from ophelper.assets import file_from_url
+from opspec.op import Axes, op
+from opspec.runtime import progress
+from opspec.types import ImageOf
+
+from .._util import to_rgb
+from ._result import Boxes
+
+_WEIGHTS = "https://github.com/ultralytics/assets/releases/download/v8.3.0/{name}"
+
+
+class PretrainedModel(Enum):
+    """FastSAM's two published checkpoints."""
+
+    small = "FastSAM-s.pt"
+    large = "FastSAM-x.pt"
+
+
+@op(env="pytorch")
+def fastsam(
+    image: Annotated[ImageOf[np.ndarray], Axes("y", "x", "c?")],
+    model: PretrainedModel = PretrainedModel.small,
+    conf: Annotated[
+        float,
+        {"widget_type": "FloatSlider", "min": 0.0, "max": 1.0, "step": 0.05},
+    ] = 0.4,
+    iou: Annotated[
+        float,
+        {"widget_type": "FloatSlider", "min": 0.0, "max": 1.0, "step": 0.05},
+    ] = 0.9,
+    max_det: int = 300,
+    imgsz: int = 1024,
+) -> Boxes:
+    """Find every object in an image, without regard to what it is.
+
+    Args:
+        image: Image to detect in. 2-D or RGB; anything else is stretched to
+            8-bit RGB first, which is what the model was trained on.
+        model: Which checkpoint to use. The small one is 24 MB, the large
+            145 MB.
+        conf: Confidence threshold. Lower finds more, and more spuriously.
+        iou: NMS threshold. Two detections overlapping by more than this are
+            treated as one object, so dense or touching objects need it high.
+        max_det: Cap on how many objects to report.
+        imgsz: Size the image is resized to for inference.
+
+    Returns:
+        boxes: (N, 4) as [min_y, min_x, max_y, max_x], in image coordinates.
+
+    Note: the weights are downloaded into skop's asset cache on first use.
+    """
+    from ultralytics import FastSAM
+
+    weights = file_from_url(_WEIGHTS.format(name=model.value), f"fastsam/{model.value}")
+
+    progress(f"Loading {model.value}")
+    net = FastSAM(str(weights))
+
+    progress("Detecting objects")
+    results = net.predict(
+        to_rgb(image),
+        conf=conf,
+        iou=iou,
+        max_det=max_det,
+        imgsz=imgsz,
+        retina_masks=False,
+        verbose=False,
+    )
+
+    # An image with nothing in it gives back an empty result list, not a
+    # result holding zero boxes.
+    if not results or results[0].boxes is None:
+        return Boxes(ophelper.boxes.EMPTY.copy())
+
+    found = results[0].boxes
+    progress(f"Found {len(found)} objects")
+    # found.conf holds a confidence per box. Not returned -- see _result.py.
+    return Boxes(ophelper.boxes.from_xyxy(found.xyxy.cpu().numpy()))