from __future__ import annotations
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Annotated, Any, List
import numpy as np
from phenotypic.sdk_.typing_ import GpuInputLayer, GpuOutputKind, TuneSpec
if TYPE_CHECKING:
from phenotypic._core._image import Image
from ._object_detector import ObjectDetector
# <<Interface>>
[docs]
class GpuDetector(ObjectDetector, ABC):
"""Interface ABC for GPU-accelerated object detectors (batched/streaming).
Subclass GpuDetector when your detection algorithm depends on a GPU
(e.g., deep-learning foundation models like SAM2 or micro-sam).
GpuDetector provides a concrete ``_operate`` built from a small set of
protected, overridable hooks — ``_preprocess`` (raw ``input_layer`` array →
model-ready sample), ``_collate`` (samples → batch), ``_infer_batch``
(batch → per-sample results), and ``_write_object_output`` (result →
``objmap``/``objmask``). The single-image notebook path and the batched CLI
engine drive the *same* hooks, so a detector implemented once runs in both.
Capability is declared via three fields — ``input_layer``
(``rgb``/``gray``/``detect_mat``; defaults to the layer the model was
trained on), ``supports_batching``, and ``output_kind``
(``instance``/``semantic``).
When a pipeline contains a GpuDetector, the CLI enforces:
- **Local execution:** Sequential processing (n_jobs=1) to avoid
multiple workers competing for the same GPU.
- **SLURM execution:** Automatically requests GPU resources
(``--gpus-per-node=1``) if the user hasn't specified GPU args.
Raises an error if the target partition has no GPUs.
- **No GPU available:** Raises RuntimeError at pipeline validation
time with a clear message.
**When to subclass GpuDetector vs ObjectDetector**
Subclass GpuDetector if your detector relies on GPU-accelerated inference:
- Deep-learning models (SAM2, micro-sam, or custom neural networks).
- Any algorithm that requires ``torch``, ``tensorflow``, or similar
GPU-backed frameworks at inference time.
- Detectors where CPU fallback is technically possible but
impractically slow for production use.
Subclass ObjectDetector directly if your algorithm is CPU-based:
- Classical computer vision (thresholding, edge detection, watershed).
- Algorithms implemented with NumPy, SciPy, or scikit-image.
- Detectors that run in milliseconds on CPU.
**Lazy model loading**
GpuDetector subclasses should defer model construction to the first
``apply()`` call rather than ``__init__()``. This enables:
- Fast construction and serialization round-trips without GPU/torch.
- Pipeline ``to_json()``/``from_json()`` without importing heavy
dependencies.
- Parameter inspection and validation before committing GPU memory.
Use a ``_ensure_model_loaded()`` pattern::
from pydantic import PrivateAttr
class MyGpuDetector(GpuDetector):
model_size: str = "small" # Annotated class-level fields
device: str = "auto"
# underscore-prefixed private attr → skipped by serialization
_model: object = PrivateAttr(default=None)
def _ensure_model_loaded(self):
if self._model is not None:
return
import torch # lazy import
# ... build model ...
def _infer_one(self, sample):
# ``sample`` is a preprocessed (H, W, 3) uint8 array. Return a
# uint16 labeled objmap (output_kind="instance") or a bool mask
# (output_kind="semantic"). The base _operate/_infer_batch wire
# this into the image; do NOT override _operate.
# ... run inference ...
return objmap
Notes:
``_operate`` is concrete here and should not be overridden. Non-batchable
subclasses (SAM2, micro-sam) implement just ``_ensure_model_loaded`` +
``_infer_one``; the default ``_infer_batch`` loops ``_infer_one`` and is
the sole caller of ``_ensure_model_loaded``. Batchable subclasses
(Spec 2 foundation models) instead override ``_infer_batch`` with a true
``(N, C, H, W)`` forward — no engine changes needed. The class also lets
the CLI make informed GPU resource-allocation decisions.
"""
# Capability / routing markers — pydantic FIELDS (not ClassVar) so they
# serialize and round-trip (Spec 1 §4, review S4). Subclasses override the
# defaults; "instance" keeps existing SAM behavior unchanged.
input_layer: GpuInputLayer = "rgb"
supports_batching: bool = False
output_kind: GpuOutputKind = "instance"
# Post-inference cleanup: split a single instance label that spans spatially
# disconnected blobs into separate instances by connected components. A SAM
# mask (or a tile-merged objmap) can paint one label across distant regions;
# relabeling by connectivity gives each connected region its own id. Binary
# connected-components, so two *touching* distinct labels merge into one.
# ``instance`` output only — ``semantic`` already auto-labels by connectivity.
split_disconnected_labels: bool = True
# Connectivity for the relabel (1 = 4-neighbour, 2 = 8-neighbour). Structural,
# never tuned (TuneSpec(tunable=False) satisfies the annotation-coverage gate).
connectivity: Annotated[int, TuneSpec(tunable=False)] = 2
@abstractmethod
def _ensure_model_loaded(self) -> None:
"""Build/load the GPU model on first use (idempotent)."""
def _preprocess(self, array: np.ndarray) -> Any:
"""Turn a raw ``input_layer`` array into a model-ready ``uint8`` sample.
Single-channel 2D layers (``gray``/``detect_mat``) are stacked into an
``(H, W, 3)`` block so 3-channel models (SAM/DINO ViT) consume them
unchanged; ``rgb`` keeps its three channels. The result is then coerced
to ``uint8`` — float ``[0, 1]`` (``gray``/``detect_mat``) and ``uint16``
(16-bit ``rgb``) layers are max-normalized to ``0..255``, while an
already-``uint8`` ``rgb`` array passes through byte-identical — so every
layer reaches the model through this one shared conversion. Subclasses
rarely need to override this.
"""
if array.ndim == 2:
array = np.stack([array, array, array], axis=-1)
if array.dtype != np.uint8:
max_val = array.max()
if max_val > 0:
array = (array / max_val * 255).astype(np.uint8)
else:
array = np.zeros(array.shape, dtype=np.uint8)
return array
def _collate(self, samples: List[Any]) -> Any:
"""Merge per-sample ``_preprocess`` outputs into a batch.
Default returns the list unchanged (consumed by the looped
``_infer_batch``). Batchable subclasses override to stack into a tensor.
"""
return samples
def _infer_batch(self, batch: Any) -> List[np.ndarray]:
"""Run inference over a collated batch; return one result per sample.
Each result is a uint16 labeled map (``output_kind="instance"``) or a
boolean mask (``output_kind="semantic"``). The default loops
``_infer_one`` (correct for ``supports_batching=False``); batchable
subclasses override with a true ``(N, C, H, W)`` forward.
"""
self._ensure_model_loaded()
return [self._infer_one(sample) for sample in batch]
def _infer_one(self, sample: Any) -> np.ndarray:
"""Run the model on ONE preprocessed sample. Subclasses must implement.
Returns a uint16 labeled objmap (instance) or a boolean mask (semantic).
"""
raise NotImplementedError(
f"{type(self).__name__} must implement _infer_one()"
)
def _write_object_output(self, image: "Image", result: np.ndarray) -> None:
"""Write one ``infer_batch`` result onto the image per ``output_kind``.
- ``instance`` -> ``image.objmap[:]`` (detector-controlled labels).
When ``split_disconnected_labels`` is set, the written objmap is then
relabeled by connected components (``connectivity``) so one label
spanning spatially disconnected blobs becomes separate instances.
- ``semantic`` -> ``image.objmask[:]`` (auto-labels into the shared
``objmap`` backend, exactly like a threshold detector; see Spec 1 §8).
Already connectivity-labeled, so ``split_disconnected_labels`` is a
no-op here.
"""
if self.output_kind == "instance":
image.objmap[:] = result.astype(np.uint16)
if self.split_disconnected_labels:
image.objmap.relabel(connectivity=self.connectivity)
else: # semantic
image.objmask[:] = result.astype(bool)
def _operate(self, image: "Image") -> "Image":
"""Run GPU detection on one image (notebook / single-image path).
Reads the declared ``input_layer``, preprocesses, runs a one-element
batch through ``_collate`` + ``_infer_batch``, and writes the result via
``output_kind``. The batched CLI engine drives the same
``_preprocess``/``_collate``/``_infer_batch`` methods over many images.
"""
array = getattr(image, self.input_layer)[:]
sample = self._preprocess(array)
batch = self._collate([sample])
results = self._infer_batch(batch)
self._write_object_output(image, results[0])
return image
