Image prediction and instance masks#

Use a model class for image prediction. Use Predictor to produce prediction maps and PostProcessor to turn those maps into individual objects. For whole slides, see slide inference or the WSI guide.

Read prediction results#

Predictor returns a dictionary with nuc, cyto and tissue entries. Unused entries are None. Each type_map contains the predicted class at each pixel.

Field	Direct network call	After Predictor
`type_map`	Floating class logits, `(B, classes, H, W)`	Integer class IDs, `(B, H, W)`
`aux_map`	Flows or distances used to separate objects	Same maps, combined across image windows if requested
`binary_map`	Optional foreground map	Meaning depends on the model
`parents`	Names of the output branches	Unchanged

B is batch size, H height and W width. Tensor maps stay on the network device. PostProcessor converts them to CPU NumPy arrays.

Class IDs follow the checkpoint’s labels. Class 0 can include untyped nuclei; do not assume it always means background.

Predictor#

NumPy and PIL images use height, width, channels. Tensors use channels, height, width, with an optional leading batch axis. Normalize images as required by the checkpoint. Tensor inputs must match the network’s device and floating dtype.

To keep raw class scores, call the network directly. save_intermediate=True returns the same processed maps twice; it does not keep the raw scores.

from cellseg_models_pytorch.inference.predictor import Predictor

class Predictor(model, mixed_precision=True)#

Run a configured network on images, including overlapping windows.

This is the lower-level alternative to a model class’s predict method. It returns class labels and reconstruction maps, not individual objects.

Choose the network and prediction precision.

Parameters:

model (Module) – A package network that returns the nuc, cyto and tissue map entries. Put it on the intended device and call eval() first; the predictor does not change its training mode.
mixed_precision (bool) – Use float16 autocasting on the network’s device. False keeps the network’s normal computation dtype; use it for ordinary float32 CPU prediction.

predict_sliding_win(x, window_size, stride, padding=20, apply_boundary_weight=True, save_intermediate=False)#

Predict overlapping windows and combine them at the original image size.

Parameters:

x (Tensor | ndarray | Image) – RGB NumPy/PIL image (H, W, C), or floating tensor (C, H, W) or (B, C, H, W). B is batch size, C channels, H height and W width. NumPy/PIL uint8 values are divided by 255. Other image values are unchanged. Tensors must already have the correct dtype/device. Apply the checkpoint’s normalization before prediction.
window_size (tuple[int, int]) – Window (height, width). Each dimension must be positive and compatible with the network’s encoder and decoder.
stride (int) – Positive step in pixels, no larger than either window dimension.
padding (int | None) – Extra coverage used to build the window grid. None uses 20. Reflection padding is used where possible, otherwise edge replication. Padding is removed from the returned maps.
apply_boundary_weight (bool) – Reduce the influence of window edges on maps used to separate objects and foreground maps. Class probabilities are averaged before selecting class labels.
save_intermediate (bool) – Return an additional reference to the result; it does not keep a separate copy of the raw class scores.

Returns:

A dictionary with nuc, cyto and tissue entries; unused outputs are None. Instance type_map values are integer class labels (B, H, W); maps for separating objects remain floating tensors (B, channels, H, W). Results stay on the network device. With save_intermediate=True, returns the same dictionary twice. Call the network directly for raw class scores.

Raises:

ValueError – Window sizes, stride, padding or image dimensions are invalid.

Return type:

Predictor.predict(x, apply_boundary_weight=False, save_intermediate=False)#

Predict a single image or a batch without windowing.

Parameters:

x (Tensor | ndarray | Image) – RGB NumPy/PIL image (H, W, C), or floating tensor (C, H, W) or (B, C, H, W). B is batch size, C channels, H height and W width. NumPy/PIL uint8 values are divided by 255. Other image values are unchanged. Tensors must already have the correct dtype/device. Apply the checkpoint’s normalization before prediction.
apply_boundary_weight (bool) – Give more weight to the image center and less to its edges in nucleus and cytoplasm reconstruction and foreground maps. Class labels and tissue maps are unchanged. Defaults to False.
save_intermediate (bool) – Return an additional reference to the result. The current implementation does not preserve a separate raw-score copy.

Returns:

Raises:

ValueError – The input is not a three- or four-dimensional tensor after image conversion.

Return type:

PostProcessor#

from cellseg_models_pytorch.inference.post_processor import PostProcessor

Choose the method matching your model: cellpose, omnipose, hovernet for HoverNet/CellVit, or stardist for StarDist.

Batch methods require a nuc entry. Each result contains two int32 (H, W) masks: instances identifies objects. classes gives the most common predicted class within each object. Instance 0 is background; positive IDs identify objects within that image.

import torch

from cellseg_models_pytorch.models.cellpose import CellPose
from cellseg_models_pytorch.transforms import minmax_normalize
from cellseg_models_pytorch.utils import FileHandler

model = CellPose.from_pretrained(
    "checkpoints/cellpose.safetensors", device=torch.device("cpu")
)
model.set_inference_mode(mixed_precision=False)
image = minmax_normalize(FileHandler.read_img("data/image.png"))
prediction = model.predict(image)
post_processor = PostProcessor("cellpose", {"use_gpu": False})
result = post_processor.postproc_serial(prediction)
instances, classes = result["nuc"][0]

Provide save paths to write polygons instead of returning masks. Paths may be overwritten. Polygon coordinates are in pixels, with optional x/y offsets. Smoothing changes saved polygon outlines only.

class PostProcessor(postproc_method, postproc_kwargs={})#

Reconstruct individual objects and assign a class to each object.

Choose the method that matches your model family. Batch methods copy tensor maps to CPU NumPy arrays to separate objects.

Choose the reconstruction method and its options.

Parameters:

postproc_method (str) – Method name, for example cellpose, omnipose, hovernet, stardist or instanseg. Must match the model output.
postproc_kwargs (dict) – Options passed to that method. For CellPose, use {"use_gpu": False} for CPU reconstruction. InstanSeg needs the trained network’s pixel_classifier.

Raises:

ValueError – The method name is unknown.

postproc_inst(inst_map, aux_map, type_map, save_path=None, coords=None, class_dict=None, smooth_func=<function gaussian_smooth>)#

Reconstruct one image’s objects and vote on their classes.

Parameters:

inst_map (ndarray) – HW foreground labels/scores expected by the reconstruction method. This is the foreground input, not already reconstructed instance IDs.
aux_map (ndarray) – Floating family-specific reconstruction map (channels, H, W).
type_map (ndarray) – Integer class labels (H, W) used for majority voting per object.
save_path (Path | str) – Optional vector file path (.parquet, .geojson or .feather). Existing-file behavior follows the selected writer.
coords (Tuple[int, int, int, int]) – Optional (x, y, width, height) tuple. Saved geometry is offset by x and y; masks returned in memory keep local pixel coordinates.
class_dict (dict[int, str] | None) – Map integer class numbers to names for vector output.
smooth_func (Callable) – Optional polygon-smoothing function used only when saving.

Returns:

(instances, classes) as int32 NumPy arrays of shape (H, W). Instance 0 is background; positive IDs identify objects. Returns None instead when saving to a file. Smoothing affects saved polygons, not returned masks.

Return type:

tuple[ndarray, ndarray] | None

postproc_parallel(soft_masks, n_jobs=4, start_method='threading', progress_bar=False, save_paths_nuc=None, save_paths_cyto=None, save_paths_tissue=None, coords=None, class_dict_nuc=None, class_dict_cyto=None, class_dict_tissue=None, nuc_smooth_func=<function gaussian_smooth>, cyto_smooth_func=<function gaussian_smooth>, tissue_smooth_func=None)#

Reconstruct a batch using workers.

Parameters:

soft_masks – Prediction dictionary with a nucleus entry and optional cyto/tissue entries. Class maps must already contain class labels, as returned by Predictor, not raw class logits.
n_jobs (int) – Number of workers; a batch of one uses one worker.
start_method (str) – Worker mode: “threading”, “fork” or “spawn”, subject to operating-system support.
progress_bar (bool) – Show worker progress.
save_paths_nuc (List[Path | str]) – Optional vector output paths, one per image, for nuclei.
save_paths_cyto (List[Path | str]) – Equivalent paths for cytoplasm.
save_paths_tissue (List[Path | str]) – Equivalent paths for tissue.
coords (List[Tuple[int, int, int, int]]) – Optional (x, y, width, height) per image for vector coordinate offsets.
class_dict_nuc (Dict[int, str]) – Map nucleus class numbers to output names.
class_dict_cyto (Dict[int, str]) – Map cytoplasm class numbers to output names.
class_dict_tissue (Dict[int, str]) – Map tissue class numbers to output names.
nuc_smooth_func (Callable) – Polygon smoothing for saved nuclei; None disables it.
cyto_smooth_func (Callable) – Polygon smoothing for saved cytoplasm; None disables it.
tissue_smooth_func (Callable) – Polygon smoothing for saved tissue; None disables it.

Returns:

Dictionary with nuc, cyto and tissue lists, in input order. Missing heads are None. Instance results are (instances, classes) int32 NumPy masks (H, W); tissue results are int32 class masks (H, W). Saving an image returns None at that list position. All worker results are resolved before returning; no future is returned to the caller.

Return type:

dict[str, list[tuple[ndarray, ndarray] | ndarray | None] | None]

postproc_parallel_async(soft_masks, n_jobs=4, start_method='threading', save_paths_nuc=None, save_paths_cyto=None, save_paths_tissue=None, coords=None, class_dict_nuc=None, class_dict_cyto=None, class_dict_tissue=None, nuc_smooth_func=<function gaussian_smooth>, cyto_smooth_func=<function gaussian_smooth>, tissue_smooth_func=None)#

Schedule batch reconstruction on workers and wait for all results.

Parameters:

soft_masks – Prediction dictionary with a nucleus entry and optional cyto/tissue entries. Class maps must already contain class labels, as returned by Predictor, not raw class logits.
n_jobs (int) – Number of workers; a batch of one uses one worker.
start_method (str) – Worker mode: “threading”, “fork” or “spawn”, subject to operating-system support.
save_paths_nuc (List[Path | str]) – Optional vector output paths, one per image, for nuclei.
save_paths_cyto (List[Path | str]) – Equivalent paths for cytoplasm.
save_paths_tissue (List[Path | str]) – Equivalent paths for tissue.
coords (List[Tuple[int, int, int, int]]) – Optional (x, y, width, height) per image for vector coordinate offsets.
class_dict_nuc (Dict[int, str]) – Map nucleus class numbers to output names.
class_dict_cyto (Dict[int, str]) – Map cytoplasm class numbers to output names.
class_dict_tissue (Dict[int, str]) – Map tissue class numbers to output names.
nuc_smooth_func (Callable) – Polygon smoothing for saved nuclei; None disables it.
cyto_smooth_func (Callable) – Polygon smoothing for saved cytoplasm; None disables it.
tissue_smooth_func (Callable) – Polygon smoothing for saved tissue; None disables it.

Returns:

Return type:

dict[str, list[tuple[ndarray, ndarray] | ndarray | None] | None]

postproc_serial(soft_masks, save_paths_nuc=None, save_paths_cyto=None, save_paths_tissue=None, coords=None, class_dict_nuc=None, class_dict_cyto=None, class_dict_tissue=None, nuc_smooth_func=<function gaussian_smooth>, cyto_smooth_func=<function gaussian_smooth>, tissue_smooth_func=None)#

Reconstruct a batch one image at a time.

Parameters:

soft_masks – Prediction dictionary with a nucleus entry and optional cyto/tissue entries. Class maps must already contain class labels, as returned by Predictor, not raw class logits.
save_paths_nuc (List[Path | str]) – Optional vector output paths, one per image, for nuclei.
save_paths_cyto (List[Path | str]) – Equivalent paths for cytoplasm.
save_paths_tissue (List[Path | str]) – Equivalent paths for tissue.
coords (List[Tuple[int, int, int, int]]) – Optional (x, y, width, height) per image for vector coordinate offsets.
class_dict_nuc (Dict[int, str]) – Map nucleus class numbers to output names.
class_dict_cyto (Dict[int, str]) – Map cytoplasm class numbers to output names.
class_dict_tissue (Dict[int, str]) – Map tissue class numbers to output names.
nuc_smooth_func (Callable) – Polygon smoothing for saved nuclei; None disables it.
cyto_smooth_func (Callable) – Polygon smoothing for saved cytoplasm; None disables it.
tissue_smooth_func (Callable) – Polygon smoothing for saved tissue; None disables it.

Returns:

Return type:

dict[str, list[tuple[ndarray, ndarray] | ndarray | None] | None]

postproc_tissuemap(tissue_map, save_path=None, coords=None, class_dict=None, smooth_func=None)#

Clean one tissue-class mask and optionally save polygons.

Parameters:

tissue_map (ndarray) – Integer class mask (H, W). Small regions below 5000 pixels are removed and holes are filled before output.
save_path (Path | str) – Optional vector output path; without it, return the mask.
coords (Tuple[int, int, int, int]) – Optional (x, y, width, height) tuple. Saved geometry is offset by x and y; masks returned in memory keep local pixel coordinates.
class_dict (dict[int, str] | None) – Map integer class numbers to names for vector output.
smooth_func (Callable) – Optional polygon-smoothing function used only when saving.

Returns:

Cleaned int32 class mask (H, W), or None when saving polygons.

Return type:

ndarray | None

Prediction map fields#

Each container holds the maps for one output branch. Despite the “Soft” name, Predictor replaces class scores with integer class labels.

class SoftInstanceOutput(type_map, aux_map, binary_map=None, parents=None)#

Prediction maps for one instance head, such as nuclei or cytoplasm.

Parameters:

type_map (Tensor) – Raw network class logits of shape (B, classes, H, W). After Predictor runs, this field contains integer class labels of shape (B, H, W), not probabilities.
aux_map (Tensor) – Floating reconstruction values of shape (B, channels, H, W). Channels depend on the model family: flows, horizontal/vertical distances, radial distances or InstanSeg values.
binary_map (Tensor | None) – Optional separate foreground map of shape (B, 1, H, W). Interpretation depends on the head; it is not universally a probability.
parents (dict[str, list[str] | None] | None) – Optional mapping from field names to their producing branch/head names. Useful for identifying which network outputs populated a field.

Maps remain on the network device and use its computation dtype, except class labels produced by prediction are integer tensors. B is batch size, H height, W width. Reading a field does not copy its tensor.

class SoftSemanticOutput(type_map, binary_map=None, parents=None)#

Prediction maps for a semantic head, such as tissue classes.

Parameters:

type_map (Tensor) – Raw class logits of shape (B, classes, H, W) from the network. Predictor smooths the tissue probabilities and then replaces this field with integer class labels of shape (B, H, W).
binary_map (Tensor | None) – Optional separate binary output map, otherwise None.
parents (dict[str, list[str] | None] | None) – Optional mapping from field names to producing branch/head names.

B is batch size, H height and W width. Tensors remain on the network device; raw maps use its computation dtype and predicted class labels are integers.