deepcell_types.predict

deepcell_types.predict(raw, mask, channel_names, mpp, *, model_name, device=None, batch_size=256, num_workers=0, zarr_path=None, return_probabilities=False, ct_abstention_k=0, preprocess=None) → list[str] | PredictionResult

Run the cell-type prediction pipeline.

Given a spatial proteomics image raw, a corresponding segmentation mask, and a list of markers (channel_names) corresponding to the channels of raw, predict the cell type associated with each index in mask.

Parameters:

rawA spatial proteomic image as an numpy.ndarray with shape (C, H, W).

A 2D multiplexed image in channel-first format. The image will be converted internally to dtype=np.float32.

mask2D label image

Segmentation mask of raw as a 2D label image with shape (H, W).

channel_nameslist of str

A list of channel markers. Must have the same length as the number of channels in raw and be given in the same order as the channels in raw.

mppfloat

The image resolution in microns-per-pixel. Improves prediction performance by removing scale variability.

model_namestr

Which pre-trained model to use. Accepts a registry version string (see list_model_versions(), e.g. "2026-06-15") or the sentinel "latest", either of which is downloaded and cached automatically. A bare name is searched for at Path.home() / ".deepcell/models", and a filesystem path to a .pt file is also accepted.

devicetorch.device or str

Which device to run inference on, e.g. "cpu", "cuda", or "cuda:0" to select a specific GPU. All arguments after mpp are keyword-only.

batch_sizeint, default=256

Batch size to be used for inference. Larger batch_size will increase performance by increasing VRAM usage. Default value of 256 is conservative and should be appropriate for systems with <16GB VRAM.

num_workersint, default=0

Number of DataLoader worker processes. Default 0 runs the patch generator in-process (safe on all machines). PatchDataset is an IterableDataset that holds the full FOV in memory, so each worker is an extra copy AND re-runs the per-FOV preprocessing — only raise this on machines with abundant RAM and CPU.

zarr_pathstr or pathlib.Path, optional, default=None

Optional path to a TissueNet zarr archive to read the marker / cell-type registry from (or set DEEPCELL_TYPES_ZARR_PATH). When omitted, the registry is read from the vocabulary snapshot shipped with the package, so inference does not require the archive.

return_probabilitiesbool, default=False

If False (default, back-compat), returns a list of cell-type names. If True, returns a PredictionResult with the full per-cell softmax probability matrix and the cell indices.

ct_abstention_kfloat, default=0

IQR-fence post-hoc abstention multiplier. Abstention is opt-in: the default 0 (any k <= 0) returns the raw argmax cell-type label for every cell and never relabels to "Unknown". Pass a positive float to enable it: the fence is Q1 - k*IQR on the per-FOV cell-wise max-softmax distribution and cells below it are relabelled to "Unknown" (k=0.2 is a historical opt-in ablation; the paper headline is full-coverage with no abstention). Pass return_probabilities=True to recover the pre-abstention labels and the abstained mask. Has no effect on FOVs with fewer than 4 cells (the IQR is undefined).

preprocesscallable, optional, default=None

Custom per-FOV preprocessing hook. Called as preprocess(raw, channel_names) -> raw where raw is a (C, H, W) float32 array already resampled to the model’s target MPP and restricted to in-vocabulary channels, and channel_names are the resolved standard marker names aligned to raw. Must return a (C, H, W) array in [0, 1]. When None (default), the built-in per-channel p99.9 clip + min-max normalization is used. Build one declaratively with deepcell_types.make_preprocessor(). Note: deepcell_types.preprocess_fov() is NOT a valid hook — it has a different signature (takes mask / keyword-only native_mpp) and returns a PreprocessedFov, not a bare (C, H, W) array.

Returns:

list of str

(default) Predicted cell-type name for each unique cell index in mask, ordered by ascending cell index. Cells flagged by the IQR-fence abstention carry the sentinel "Unknown" — but only when ct_abstention_k > 0; with the default 0 every cell carries its raw argmax label.

PredictionResult

(when return_probabilities=True) Full per-cell probabilities, cell indices, predicted names, and an abstained boolean mask. See PredictionResult.