core/cellpose_engine.py

"""Real Cellpose backend — generalist deep cell/nucleus segmentation.

Returns the same result dict shape as `fast_seg.analyze` so the viz + scoring are
shared. Needs the `cellpose` package + model weights (Dockerfile.cellpose installs
them; the model downloads on first use). See https://github.com/MouseLand/cellpose
"""
from __future__ import annotations

import numpy as np

from . import fast_seg

_MODEL = None


def available() -> bool:
    try:
        import cellpose  # noqa: F401
        return True
    except Exception:  # noqa: BLE001
        return False


def _get_model():
    global _MODEL
    if _MODEL is None:
        from cellpose import models
        # Cellpose-SAM (v4): one generalist model; weights hosted on HuggingFace.
        _MODEL = models.CellposeModel(gpu=False)
    return _MODEL


def analyze(img: np.ndarray, diameter: float = 0.0, gt: np.ndarray | None = None) -> dict:
    if not available():
        raise RuntimeError(
            "cellpose engine needs the `cellpose` package. Build with Dockerfile.cellpose. "
            "Use engine='fast' for the always-available classic segmentation."
        )
    out = _get_model().eval(img, diameter=(diameter or None))    # v4: (masks, flows, styles)
    lab = np.asarray(out[0], np.int32)
    import cellpose
    from skimage.measure import regionprops
    areas = [r.area for r in regionprops(lab)]
    report = {
        "engine": f"cellpose-SAM (v{cellpose.version})",
        "dims": list(img.shape),
        "n_cells": int(lab.max()),
        "mean_area_px": round(float(np.mean(areas)) if areas else 0.0, 1),
    }
    if gt is not None:
        report.update(fast_seg.score(lab, gt))
    return {"labels": lab, "image": img, "report": report}