Models#

Load weights and predict individual cells with a model class. For training, use its .model PyTorch network. Start with the quick start or image inference guide.

Choose a family and checkpoint#

Family

HGSC checkpoint name

Cervix (CIN2)

CellPose

hgsc_v1_efficientnet_b5

cin2_v1_efficientnet_b5

StarDist

hgsc_v1_efficientnet_b5

cin2_v1_efficientnet_b5

HoverNet

hgsc_v1_efficientnet_b5

cin2_v1_efficientnet_b5

CellVit

hgsc_v1_samvit_base_patch16

cin2_v1_samvit_base_patch16

InstanSeg

Use local weights

—

Load pretrained weights by checkpoint name, or pass a local path to load your own. Named checkpoints may download weights from Hugging Face. Use the normalization and class labels supplied with the weights.

CellVit’s former name, hgsc_v1_efficientnet_b5, still loads the same weights.

CIN2 repositories currently contain model cards only; weights cannot yet be loaded by name. Load any compatible local CIN2 weights by path.

Predict images#

Call set_inference_mode() before prediction. For CPU prediction, pass mixed_precision=False. CellPose also needs CPU reconstruction configured as shown in the image inference guide.

predict produces maps; post_process turns them into instance and class masks. See inference for input formats and output fields.

Train a model#

Use .model with a normal PyTorch training loop. After training, call set_inference_mode() again before predicting images. The training guide includes encoder freezing; the network reference covers architecture options.

CellPose#

from cellseg_models_pytorch.models.cellpose import CellPose
class CellPose(n_nuc_classes, enc_name='efficientnet_b5', enc_pretrain=True, enc_freeze=False, device=device(type='cuda'), model_kwargs={})#

Find nuclei and predict their cell types with a CellPose network.

Use model to access the PyTorch network for training. For image prediction, call set_inference_mode before predict and post_process. Loading weights alone does not enable inference.

The architecture follows CellPose.

Configure the network and move it to the requested device.

Parameters:
  • n_nuc_classes (int) – Number of cell-type output channels, including any background class in your label mapping.

  • enc_name (str) – Encoder name supported by timm (PyTorch Image Models).

  • enc_pretrain (bool) – Load pretrained encoder weights. This may download weights; it does not load trained segmentation layers.

  • enc_freeze (bool) – Disable gradient updates to the encoder weights. Training mode still controls its normalization layers.

  • device (device) – Device on which to create the network. Pass torch.device("cpu") for CPU use; the default requires CUDA.

  • model_kwargs (Dict[str, Any]) – Additional options for CellPoseUnet, such as decoder settings. These must match the architecture when loading an existing checkpoint.

Example

Create an untrained network without downloading encoder weights:

>>> import torch
>>> from cellseg_models_pytorch.models.cellpose import CellPose
>>> model = CellPose(
...     3,
...     enc_name="resnet18",
...     enc_pretrain=False,
...     device=torch.device("cpu"),
... )
>>> model.inference_mode
False
set_inference_mode(mixed_precision=True)#

Prepare image prediction and instance reconstruction.

Sets the network to evaluation mode and replaces its predictor and post-processor. This affects later training; call model.train() on the underlying network when returning to training.

CellPose reconstruction initially uses use_gpu=True. For CPU-only reconstruction, replace post_processor with PostProcessor("cellpose", {"use_gpu": False}) after this call.

Parameters:

mixed_precision (bool) – Use float16 autocasting during prediction. Choose False for ordinary float32 inference, especially on CPU. Reconstruction has its own device settings.

Return type:

None

classmethod CellPose.from_pretrained(weights, device=device(type='cuda'), model_kwargs={})#

Create a model from saved segmentation weights.

The checkpoint determines the encoder name and number of cell-type output channels, plus the ray count for StarDist. Use the same preprocessing and class mapping as the training run. Call set_inference_mode before predicting images.

Parameters:
  • weights (str | Path) – Path to a local checkpoint, or a registered checkpoint name for this model family. Registered names download weights from Hugging Face when they are not already cached. Safetensors files and PyTorch state dictionaries must match this library’s network.

  • device (device) – Device on which to load the weights and create the network. Pass torch.device("cpu") for CPU use; the default requires CUDA.

  • model_kwargs (Dict[str, Any]) – Additional network options needed to reproduce the checkpoint’s architecture. Encoder name, class count and StarDist ray count are read from the weights.

Returns:

An instance of the model class used to call this method, with loaded weights. Inference is not yet enabled; its PyTorch network is available as model for training.

Raises:
  • ValueError – The path is not a file or a registered checkpoint name.

  • ImportError – Loading a safetensors file requires safetensors.

  • RuntimeError – The checkpoint weights do not match the network.

Return type:

BaseModelInst

Example

Load your checkpoint for image inference on the CPU:

from pathlib import Path

import torch

from cellseg_models_pytorch.models.cellpose import CellPose

model = CellPose.from_pretrained(
    Path("checkpoints/cellpose.safetensors"),
    device=torch.device("cpu"),
)

Then follow the quick start to configure prediction and reconstruction.

CellPose.predict(x, *, use_sliding_win=False, window_size=None, stride=None)#

Predict class labels and reconstruction maps for an image or batch.

Normalize inputs as required by the checkpoint. This method does not reconstruct individual objects; pass its result to post_process.

Parameters:
  • x (Tensor | ndarray | Image) – RGB NumPy/PIL image with height, width and channels, or a tensor with channels first: (C, H, W) or (B, C, H, W). B is batch size, C channels, H height and W width. NumPy/PIL uint8 images are divided by 255; floating-point images keep their values. Tensor inputs must already have the network’s device, floating dtype and expected value range.

  • use_sliding_win (bool) – Predict overlapping windows and combine them into full-image maps. Each window must meet the network’s size constraints.

  • window_size (Tuple[int, int]) – Window (height, width), required for sliding-window inference.

  • stride (int) – Window step in pixels, required with sliding windows. Must be positive and no larger than either window dimension.

Returns:

Dictionary with nuc, cyto and tissue entries. Unused outputs are None. Instance entries contain type_map class labels of shape (B, H, W), and floating aux_map reconstruction values of shape (B, channels, H, W). These remain on the network device. A tissue entry contains class labels of shape (B, H, W). These are not probability maps. Class meanings and reconstruction maps depend on the checkpoint.

Raises:

ValueError – Inference has not been enabled, or the window settings or input dimensions are invalid.

Return type:

dict[str, SoftSemanticOutput | SoftInstanceOutput | None]

CellPose.post_process(x, *, use_async_postproc=True, start_method='threading', n_jobs=4, save_paths_nuc=None, save_paths_cyto=None, coords=None, class_dict_nuc=None, class_dict_cyto=None)#

Turn predicted maps into individual nuclei and cell-type masks.

Parameters:
  • x (Dict[str, SoftSemanticOutput | SoftInstanceOutput]) – Dictionary returned by predict, including a nucleus entry.

  • use_async_postproc (bool) – Use asynchronous worker tasks for batches larger than one. This call waits for their results before returning. False uses synchronous workers; a single image always runs serially.

  • start_method (str) – Worker mode: threading, fork or spawn. Available process modes depend on the operating system.

  • n_jobs (int) – Number of reconstruction workers for batches larger than one.

  • save_paths_nuc (List[Path | str]) – One vector-output path per image for nuclei. Without paths, return masks in memory. Saving returns None for that image.

  • save_paths_cyto (List[Path | str]) – Equivalent paths for configured cytoplasm outputs.

  • coords (List[Tuple[int, int, int, int]]) – One (x, y, width, height) tuple per image. Vector coordinates are offset by x and y when saving.

  • class_dict_nuc (Dict[int, str]) – Map nucleus class numbers to names in saved vector data.

  • class_dict_cyto (Dict[int, str]) – Map cytoplasm class numbers to names in saved vector data.

Returns:

Dictionary with nuc, cyto and tissue entries. Nucleus and cytoplasm results are lists of (instances, classes) NumPy mask pairs, one per image. Both masks are int32 (H, W). Instance 0 is background; positive IDs identify objects. Each object gets the most common predicted class within its mask. Unused outputs and saved-image results are None. Tissue results are class masks. Tensor maps are copied to CPU.

Raises:

ValueError – Inference has not been enabled.

Return type:

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

Note

Use PostProcessor directly for tissue saving or custom contour smoothing. The selected model family controls the reconstruction method.

cellseg_models_pytorch.models.cellpose.CellPose.model#

PyTorch network used for training.

cellseg_models_pytorch.models.cellpose.CellPose.device#

Device used by the model.

cellseg_models_pytorch.models.cellpose.CellPose.inference_mode#

Whether inference has been set up. Calling .model.train() does not reset this flag.

cellseg_models_pytorch.models.cellpose.CellPose.post_processor#

Post-processor created by set_inference_mode. Replace it to change reconstruction settings.

ONNX export: export_onnx() in the export reference.

StarDist#

from_pretrained reads the ray count from the saved weights.

from cellseg_models_pytorch.models.stardist import StarDist
class StarDist(n_nuc_classes, n_rays=32, enc_name='efficientnet_b5', enc_pretrain=True, enc_freeze=False, device=device(type='cuda'), model_kwargs={})#

Find nuclei and predict their cell types with a StarDist network.

Use model for training. Call set_inference_mode before predict and post_process to enable image inference.

Configure the network and move it to the requested device.

Parameters:
  • n_nuc_classes (int) – Number of cell-type output channels, including any background class in the training label mapping.

  • n_rays (int) – Number of radial distances per pixel. Match the training targets. from_pretrained reads this count from the checkpoint.

  • enc_name (str) – Encoder name supported by timm (PyTorch Image Models).

  • enc_pretrain (bool) – Load pretrained encoder weights, downloading if needed. Segmentation layers still need training or a segmentation checkpoint.

  • enc_freeze (bool) – Disable encoder weight updates. Training mode still controls its normalization layers.

  • device (device) – Network device. Pass torch.device("cpu") for CPU use; the default requires CUDA.

  • model_kwargs (Dict[str, Any]) – Additional StarDistUnet options. Use the same architecture settings as the checkpoint when loading existing weights.

set_inference_mode(mixed_precision=True)#

Prepare image prediction and instance reconstruction.

Sets the network to evaluation mode and replaces its predictor and post-processor. This affects later training; call model.train() on the underlying network when returning to training.

Configures StarDist reconstruction with bounding-box trimming and normalization enabled.

Parameters:

mixed_precision (bool) – Use float16 autocasting during prediction. Choose False for ordinary float32 inference, especially on CPU. Reconstruction has its own device settings.

Return type:

None

classmethod StarDist.from_pretrained(weights, device=device(type='cuda'), model_kwargs={})#

Create a model from saved segmentation weights.

The checkpoint determines the encoder name and number of cell-type output channels, plus the ray count for StarDist. Use the same preprocessing and class mapping as the training run. Call set_inference_mode before predicting images.

Parameters:
  • weights (str | Path) – Path to a local checkpoint, or a registered checkpoint name for this model family. Registered names download weights from Hugging Face when they are not already cached. Safetensors files and PyTorch state dictionaries must match this library’s network.

  • device (device) – Device on which to load the weights and create the network. Pass torch.device("cpu") for CPU use; the default requires CUDA.

  • model_kwargs (Dict[str, Any]) – Additional network options needed to reproduce the checkpoint’s architecture. Encoder name, class count and StarDist ray count are read from the weights.

Returns:

An instance of the model class used to call this method, with loaded weights. Inference is not yet enabled; its PyTorch network is available as model for training.

Raises:
  • ValueError – The path is not a file or a registered checkpoint name.

  • ImportError – Loading a safetensors file requires safetensors.

  • RuntimeError – The checkpoint weights do not match the network.

Return type:

BaseModelInst

Example

Load your checkpoint for image inference on the CPU:

from pathlib import Path

import torch

from cellseg_models_pytorch.models.cellpose import CellPose

model = CellPose.from_pretrained(
    Path("checkpoints/cellpose.safetensors"),
    device=torch.device("cpu"),
)

Then follow the quick start to configure prediction and reconstruction.

StarDist.predict(x, *, use_sliding_win=False, window_size=None, stride=None)#

Predict class labels and reconstruction maps for an image or batch.

Normalize inputs as required by the checkpoint. This method does not reconstruct individual objects; pass its result to post_process.

Parameters:
  • x (Tensor | ndarray | Image) – RGB NumPy/PIL image with height, width and channels, or a tensor with channels first: (C, H, W) or (B, C, H, W). B is batch size, C channels, H height and W width. NumPy/PIL uint8 images are divided by 255; floating-point images keep their values. Tensor inputs must already have the network’s device, floating dtype and expected value range.

  • use_sliding_win (bool) – Predict overlapping windows and combine them into full-image maps. Each window must meet the network’s size constraints.

  • window_size (Tuple[int, int]) – Window (height, width), required for sliding-window inference.

  • stride (int) – Window step in pixels, required with sliding windows. Must be positive and no larger than either window dimension.

Returns:

Dictionary with nuc, cyto and tissue entries. Unused outputs are None. Instance entries contain type_map class labels of shape (B, H, W), and floating aux_map reconstruction values of shape (B, channels, H, W). These remain on the network device. A tissue entry contains class labels of shape (B, H, W). These are not probability maps. Class meanings and reconstruction maps depend on the checkpoint.

Raises:

ValueError – Inference has not been enabled, or the window settings or input dimensions are invalid.

Return type:

dict[str, SoftSemanticOutput | SoftInstanceOutput | None]

StarDist.post_process(x, *, use_async_postproc=True, start_method='threading', n_jobs=4, save_paths_nuc=None, save_paths_cyto=None, coords=None, class_dict_nuc=None, class_dict_cyto=None)#

Turn predicted maps into individual nuclei and cell-type masks.

Parameters:
  • x (Dict[str, SoftSemanticOutput | SoftInstanceOutput]) – Dictionary returned by predict, including a nucleus entry.

  • use_async_postproc (bool) – Use asynchronous worker tasks for batches larger than one. This call waits for their results before returning. False uses synchronous workers; a single image always runs serially.

  • start_method (str) – Worker mode: threading, fork or spawn. Available process modes depend on the operating system.

  • n_jobs (int) – Number of reconstruction workers for batches larger than one.

  • save_paths_nuc (List[Path | str]) – One vector-output path per image for nuclei. Without paths, return masks in memory. Saving returns None for that image.

  • save_paths_cyto (List[Path | str]) – Equivalent paths for configured cytoplasm outputs.

  • coords (List[Tuple[int, int, int, int]]) – One (x, y, width, height) tuple per image. Vector coordinates are offset by x and y when saving.

  • class_dict_nuc (Dict[int, str]) – Map nucleus class numbers to names in saved vector data.

  • class_dict_cyto (Dict[int, str]) – Map cytoplasm class numbers to names in saved vector data.

Returns:

Dictionary with nuc, cyto and tissue entries. Nucleus and cytoplasm results are lists of (instances, classes) NumPy mask pairs, one per image. Both masks are int32 (H, W). Instance 0 is background; positive IDs identify objects. Each object gets the most common predicted class within its mask. Unused outputs and saved-image results are None. Tissue results are class masks. Tensor maps are copied to CPU.

Raises:

ValueError – Inference has not been enabled.

Return type:

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

Note

Use PostProcessor directly for tissue saving or custom contour smoothing. The selected model family controls the reconstruction method.

cellseg_models_pytorch.models.stardist.StarDist.model#

PyTorch network used for training.

cellseg_models_pytorch.models.stardist.StarDist.device#

Device used by the model.

cellseg_models_pytorch.models.stardist.StarDist.inference_mode#

Whether inference has been set up. Calling .model.train() does not reset this flag.

cellseg_models_pytorch.models.stardist.StarDist.post_processor#

Post-processor created by set_inference_mode. Replace it to change reconstruction settings.

ONNX export: export_onnx() in the export reference.

HoverNet#

from cellseg_models_pytorch.models.hovernet import HoverNet
class HoverNet(n_nuc_classes, enc_name='efficientnet_b5', enc_pretrain=True, enc_freeze=False, device=device(type='cuda'), model_kwargs={})#

Find nuclei and predict their cell types with a HoverNet network.

Use model for training. Call set_inference_mode before predict and post_process to enable image inference.

Configure the network and move it to the requested device.

Parameters:
  • n_nuc_classes (int) – Number of cell-type output channels, including any background class in the training label mapping.

  • enc_name (str) – Encoder name supported by timm (PyTorch Image Models).

  • enc_pretrain (bool) – Load pretrained encoder weights, downloading if needed. Segmentation layers still need training or a segmentation checkpoint.

  • enc_freeze (bool) – Disable encoder weight updates. Training mode still controls its normalization layers.

  • device (device) – Network device. Pass torch.device("cpu") for CPU use; the default requires CUDA.

  • model_kwargs (Dict[str, Any]) – Additional HoverNetUnet options. Use the same architecture settings as the checkpoint when loading existing weights.

set_inference_mode(mixed_precision=True)#

Prepare image prediction and instance reconstruction.

Sets the network to evaluation mode and replaces its predictor and post-processor. This affects later training; call model.train() on the underlying network when returning to training.

Configures HoVer-Net reconstruction for the horizontal/vertical maps.

Parameters:

mixed_precision (bool) – Use float16 autocasting during prediction. Choose False for ordinary float32 inference, especially on CPU. Reconstruction has its own device settings.

Return type:

None

classmethod HoverNet.from_pretrained(weights, device=device(type='cuda'), model_kwargs={})#

Create a model from saved segmentation weights.

The checkpoint determines the encoder name and number of cell-type output channels, plus the ray count for StarDist. Use the same preprocessing and class mapping as the training run. Call set_inference_mode before predicting images.

Parameters:
  • weights (str | Path) – Path to a local checkpoint, or a registered checkpoint name for this model family. Registered names download weights from Hugging Face when they are not already cached. Safetensors files and PyTorch state dictionaries must match this library’s network.

  • device (device) – Device on which to load the weights and create the network. Pass torch.device("cpu") for CPU use; the default requires CUDA.

  • model_kwargs (Dict[str, Any]) – Additional network options needed to reproduce the checkpoint’s architecture. Encoder name, class count and StarDist ray count are read from the weights.

Returns:

An instance of the model class used to call this method, with loaded weights. Inference is not yet enabled; its PyTorch network is available as model for training.

Raises:
  • ValueError – The path is not a file or a registered checkpoint name.

  • ImportError – Loading a safetensors file requires safetensors.

  • RuntimeError – The checkpoint weights do not match the network.

Return type:

BaseModelInst

Example

Load your checkpoint for image inference on the CPU:

from pathlib import Path

import torch

from cellseg_models_pytorch.models.cellpose import CellPose

model = CellPose.from_pretrained(
    Path("checkpoints/cellpose.safetensors"),
    device=torch.device("cpu"),
)

Then follow the quick start to configure prediction and reconstruction.

HoverNet.predict(x, *, use_sliding_win=False, window_size=None, stride=None)#

Predict class labels and reconstruction maps for an image or batch.

Normalize inputs as required by the checkpoint. This method does not reconstruct individual objects; pass its result to post_process.

Parameters:
  • x (Tensor | ndarray | Image) – RGB NumPy/PIL image with height, width and channels, or a tensor with channels first: (C, H, W) or (B, C, H, W). B is batch size, C channels, H height and W width. NumPy/PIL uint8 images are divided by 255; floating-point images keep their values. Tensor inputs must already have the network’s device, floating dtype and expected value range.

  • use_sliding_win (bool) – Predict overlapping windows and combine them into full-image maps. Each window must meet the network’s size constraints.

  • window_size (Tuple[int, int]) – Window (height, width), required for sliding-window inference.

  • stride (int) – Window step in pixels, required with sliding windows. Must be positive and no larger than either window dimension.

Returns:

Dictionary with nuc, cyto and tissue entries. Unused outputs are None. Instance entries contain type_map class labels of shape (B, H, W), and floating aux_map reconstruction values of shape (B, channels, H, W). These remain on the network device. A tissue entry contains class labels of shape (B, H, W). These are not probability maps. Class meanings and reconstruction maps depend on the checkpoint.

Raises:

ValueError – Inference has not been enabled, or the window settings or input dimensions are invalid.

Return type:

dict[str, SoftSemanticOutput | SoftInstanceOutput | None]

HoverNet.post_process(x, *, use_async_postproc=True, start_method='threading', n_jobs=4, save_paths_nuc=None, save_paths_cyto=None, coords=None, class_dict_nuc=None, class_dict_cyto=None)#

Turn predicted maps into individual nuclei and cell-type masks.

Parameters:
  • x (Dict[str, SoftSemanticOutput | SoftInstanceOutput]) – Dictionary returned by predict, including a nucleus entry.

  • use_async_postproc (bool) – Use asynchronous worker tasks for batches larger than one. This call waits for their results before returning. False uses synchronous workers; a single image always runs serially.

  • start_method (str) – Worker mode: threading, fork or spawn. Available process modes depend on the operating system.

  • n_jobs (int) – Number of reconstruction workers for batches larger than one.

  • save_paths_nuc (List[Path | str]) – One vector-output path per image for nuclei. Without paths, return masks in memory. Saving returns None for that image.

  • save_paths_cyto (List[Path | str]) – Equivalent paths for configured cytoplasm outputs.

  • coords (List[Tuple[int, int, int, int]]) – One (x, y, width, height) tuple per image. Vector coordinates are offset by x and y when saving.

  • class_dict_nuc (Dict[int, str]) – Map nucleus class numbers to names in saved vector data.

  • class_dict_cyto (Dict[int, str]) – Map cytoplasm class numbers to names in saved vector data.

Returns:

Dictionary with nuc, cyto and tissue entries. Nucleus and cytoplasm results are lists of (instances, classes) NumPy mask pairs, one per image. Both masks are int32 (H, W). Instance 0 is background; positive IDs identify objects. Each object gets the most common predicted class within its mask. Unused outputs and saved-image results are None. Tissue results are class masks. Tensor maps are copied to CPU.

Raises:

ValueError – Inference has not been enabled.

Return type:

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

Note

Use PostProcessor directly for tissue saving or custom contour smoothing. The selected model family controls the reconstruction method.

cellseg_models_pytorch.models.hovernet.HoverNet.model#

PyTorch network used for training.

cellseg_models_pytorch.models.hovernet.HoverNet.device#

Device used by the model.

cellseg_models_pytorch.models.hovernet.HoverNet.inference_mode#

Whether inference has been set up. Calling .model.train() does not reset this flag.

cellseg_models_pytorch.models.hovernet.HoverNet.post_processor#

Post-processor created by set_inference_mode. Replace it to change reconstruction settings.

ONNX export: export_onnx() in the export reference.

CellVit#

ONNX export: export_onnx() in the export reference.

Use a supported samvit_* encoder. Loading a checkpoint selects the encoder stored in its weights, which can differ from the constructor default.

from cellseg_models_pytorch.models.cellvit import CellVit
class CellVit(n_nuc_classes, enc_name='samvit_base_patch16', enc_pretrain=True, enc_freeze=False, device=device(type='cuda'), model_kwargs={})#

Find nuclei and predict their cell types with a CellVit network.

Use model for training. Call set_inference_mode before predict and post_process to enable image inference.

Configure the network and move it to the requested device.

Parameters:
  • n_nuc_classes (int) – Number of cell-type output channels, including any background class in the training label mapping.

  • enc_name (str) – Supported timm SAM encoder, such as samvit_base_patch16_224.

  • enc_pretrain (bool) – Load pretrained encoder weights, downloading if needed. Segmentation layers still need training or a segmentation checkpoint.

  • enc_freeze (bool) – Disable encoder weight updates. Training mode still controls its normalization layers.

  • device (device) – Network device. Pass torch.device("cpu") for CPU use; the default requires CUDA.

  • model_kwargs (Dict[str, Any]) – Additional CellVitSamUnet options. Use the same architecture settings as the checkpoint when loading existing weights.

set_inference_mode(mixed_precision=True)#

Prepare image prediction and instance reconstruction.

Sets the network to evaluation mode and replaces its predictor and post-processor. This affects later training; call model.train() on the underlying network when returning to training.

Configures HoVer-Net reconstruction for the horizontal/vertical maps.

Parameters:

mixed_precision (bool) – Use float16 autocasting during prediction. Choose False for ordinary float32 inference, especially on CPU. Reconstruction has its own device settings.

Return type:

None

classmethod CellVit.from_pretrained(weights, device=device(type='cuda'), model_kwargs={})#

Create a model from saved segmentation weights.

The checkpoint determines the encoder name and number of cell-type output channels, plus the ray count for StarDist. Use the same preprocessing and class mapping as the training run. Call set_inference_mode before predicting images.

Parameters:
  • weights (str | Path) – Path to a local checkpoint, or a registered checkpoint name for this model family. Registered names download weights from Hugging Face when they are not already cached. Safetensors files and PyTorch state dictionaries must match this library’s network.

  • device (device) – Device on which to load the weights and create the network. Pass torch.device("cpu") for CPU use; the default requires CUDA.

  • model_kwargs (Dict[str, Any]) – Additional network options needed to reproduce the checkpoint’s architecture. Encoder name, class count and StarDist ray count are read from the weights.

Returns:

An instance of the model class used to call this method, with loaded weights. Inference is not yet enabled; its PyTorch network is available as model for training.

Raises:
  • ValueError – The path is not a file or a registered checkpoint name.

  • ImportError – Loading a safetensors file requires safetensors.

  • RuntimeError – The checkpoint weights do not match the network.

Return type:

BaseModelInst

Example

Load your checkpoint for image inference on the CPU:

from pathlib import Path

import torch

from cellseg_models_pytorch.models.cellpose import CellPose

model = CellPose.from_pretrained(
    Path("checkpoints/cellpose.safetensors"),
    device=torch.device("cpu"),
)

Then follow the quick start to configure prediction and reconstruction.

CellVit.predict(x, *, use_sliding_win=False, window_size=None, stride=None)#

Predict class labels and reconstruction maps for an image or batch.

Normalize inputs as required by the checkpoint. This method does not reconstruct individual objects; pass its result to post_process.

Parameters:
  • x (Tensor | ndarray | Image) – RGB NumPy/PIL image with height, width and channels, or a tensor with channels first: (C, H, W) or (B, C, H, W). B is batch size, C channels, H height and W width. NumPy/PIL uint8 images are divided by 255; floating-point images keep their values. Tensor inputs must already have the network’s device, floating dtype and expected value range.

  • use_sliding_win (bool) – Predict overlapping windows and combine them into full-image maps. Each window must meet the network’s size constraints.

  • window_size (Tuple[int, int]) – Window (height, width), required for sliding-window inference.

  • stride (int) – Window step in pixels, required with sliding windows. Must be positive and no larger than either window dimension.

Returns:

Dictionary with nuc, cyto and tissue entries. Unused outputs are None. Instance entries contain type_map class labels of shape (B, H, W), and floating aux_map reconstruction values of shape (B, channels, H, W). These remain on the network device. A tissue entry contains class labels of shape (B, H, W). These are not probability maps. Class meanings and reconstruction maps depend on the checkpoint.

Raises:

ValueError – Inference has not been enabled, or the window settings or input dimensions are invalid.

Return type:

dict[str, SoftSemanticOutput | SoftInstanceOutput | None]

CellVit.post_process(x, *, use_async_postproc=True, start_method='threading', n_jobs=4, save_paths_nuc=None, save_paths_cyto=None, coords=None, class_dict_nuc=None, class_dict_cyto=None)#

Turn predicted maps into individual nuclei and cell-type masks.

Parameters:
  • x (Dict[str, SoftSemanticOutput | SoftInstanceOutput]) – Dictionary returned by predict, including a nucleus entry.

  • use_async_postproc (bool) – Use asynchronous worker tasks for batches larger than one. This call waits for their results before returning. False uses synchronous workers; a single image always runs serially.

  • start_method (str) – Worker mode: threading, fork or spawn. Available process modes depend on the operating system.

  • n_jobs (int) – Number of reconstruction workers for batches larger than one.

  • save_paths_nuc (List[Path | str]) – One vector-output path per image for nuclei. Without paths, return masks in memory. Saving returns None for that image.

  • save_paths_cyto (List[Path | str]) – Equivalent paths for configured cytoplasm outputs.

  • coords (List[Tuple[int, int, int, int]]) – One (x, y, width, height) tuple per image. Vector coordinates are offset by x and y when saving.

  • class_dict_nuc (Dict[int, str]) – Map nucleus class numbers to names in saved vector data.

  • class_dict_cyto (Dict[int, str]) – Map cytoplasm class numbers to names in saved vector data.

Returns:

Dictionary with nuc, cyto and tissue entries. Nucleus and cytoplasm results are lists of (instances, classes) NumPy mask pairs, one per image. Both masks are int32 (H, W). Instance 0 is background; positive IDs identify objects. Each object gets the most common predicted class within its mask. Unused outputs and saved-image results are None. Tissue results are class masks. Tensor maps are copied to CPU.

Raises:

ValueError – Inference has not been enabled.

Return type:

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

Note

Use PostProcessor directly for tissue saving or custom contour smoothing. The selected model family controls the reconstruction method.

cellseg_models_pytorch.models.cellvit.CellVit.model#

PyTorch network used for training.

cellseg_models_pytorch.models.cellvit.CellVit.device#

Device used by the model.

cellseg_models_pytorch.models.cellvit.CellVit.inference_mode#

Whether inference has been set up. Calling .model.train() does not reset this flag.

cellseg_models_pytorch.models.cellvit.CellVit.post_processor#

Post-processor created by set_inference_mode. Replace it to change reconstruction settings.

InstanSeg#

from cellseg_models_pytorch.models.instanseg import InstanSeg
class InstanSeg(n_nuc_classes, enc_name='efficientnet_b5', enc_pretrain=True, enc_freeze=False, device=device(type='cuda'), model_kwargs={})#

Find nuclei and predict their cell types with an InstanSeg network.

Use local trained weights; this family has no registered pretrained names.

Use model for training. Call set_inference_mode before predict and post_process to enable image inference.

Configure the network and move it to the requested device.

Parameters:
  • n_nuc_classes (int) – Number of cell-type output channels, including any background class in the training label mapping.

  • enc_name (str) – Encoder name supported by timm (PyTorch Image Models).

  • enc_pretrain (bool) – Load pretrained encoder weights, downloading if needed. Segmentation layers still need training or a segmentation checkpoint.

  • enc_freeze (bool) – Disable encoder weight updates. Training mode still controls its normalization layers.

  • device (device) – Network device. Pass torch.device("cpu") for CPU use; the default requires CUDA.

  • model_kwargs (Dict[str, Any]) – Additional InstanSegUnet options. Use the same architecture settings as the checkpoint when loading existing weights.

set_inference_mode(mixed_precision=True)#

Prepare image prediction and instance reconstruction.

Sets the network to evaluation mode and replaces its predictor and post-processor. This affects later training; call model.train() on the underlying network when returning to training.

Uses the network’s attached pixel classifier for reconstruction. On CPU, mixed precision is disabled even when requested.

Parameters:

mixed_precision (bool) – Use float16 autocasting during prediction. Choose False for ordinary float32 inference, especially on CPU. Reconstruction has its own device settings.

Return type:

None

classmethod InstanSeg.from_pretrained(weights, device=device(type='cuda'), model_kwargs={})#

Create a model from saved segmentation weights.

The checkpoint determines the encoder name and number of cell-type output channels, plus the ray count for StarDist. Use the same preprocessing and class mapping as the training run. Call set_inference_mode before predicting images.

Parameters:
  • weights (str | Path) – Path to a local checkpoint, or a registered checkpoint name for this model family. Registered names download weights from Hugging Face when they are not already cached. Safetensors files and PyTorch state dictionaries must match this library’s network.

  • device (device) – Device on which to load the weights and create the network. Pass torch.device("cpu") for CPU use; the default requires CUDA.

  • model_kwargs (Dict[str, Any]) – Additional network options needed to reproduce the checkpoint’s architecture. Encoder name, class count and StarDist ray count are read from the weights.

Returns:

An instance of the model class used to call this method, with loaded weights. Inference is not yet enabled; its PyTorch network is available as model for training.

Raises:
  • ValueError – The path is not a file or a registered checkpoint name.

  • ImportError – Loading a safetensors file requires safetensors.

  • RuntimeError – The checkpoint weights do not match the network.

Return type:

BaseModelInst

Example

Load your checkpoint for image inference on the CPU:

from pathlib import Path

import torch

from cellseg_models_pytorch.models.cellpose import CellPose

model = CellPose.from_pretrained(
    Path("checkpoints/cellpose.safetensors"),
    device=torch.device("cpu"),
)

Then follow the quick start to configure prediction and reconstruction.

InstanSeg.predict(x, *, use_sliding_win=False, window_size=None, stride=None)#

Predict class labels and reconstruction maps for an image or batch.

Normalize inputs as required by the checkpoint. This method does not reconstruct individual objects; pass its result to post_process.

Parameters:
  • x (Tensor | ndarray | Image) – RGB NumPy/PIL image with height, width and channels, or a tensor with channels first: (C, H, W) or (B, C, H, W). B is batch size, C channels, H height and W width. NumPy/PIL uint8 images are divided by 255; floating-point images keep their values. Tensor inputs must already have the network’s device, floating dtype and expected value range.

  • use_sliding_win (bool) – Predict overlapping windows and combine them into full-image maps. Each window must meet the network’s size constraints.

  • window_size (Tuple[int, int]) – Window (height, width), required for sliding-window inference.

  • stride (int) – Window step in pixels, required with sliding windows. Must be positive and no larger than either window dimension.

Returns:

Dictionary with nuc, cyto and tissue entries. Unused outputs are None. Instance entries contain type_map class labels of shape (B, H, W), and floating aux_map reconstruction values of shape (B, channels, H, W). These remain on the network device. A tissue entry contains class labels of shape (B, H, W). These are not probability maps. Class meanings and reconstruction maps depend on the checkpoint.

Raises:

ValueError – Inference has not been enabled, or the window settings or input dimensions are invalid.

Return type:

dict[str, SoftSemanticOutput | SoftInstanceOutput | None]

InstanSeg.post_process(x, *, use_async_postproc=True, start_method='threading', n_jobs=4, save_paths_nuc=None, save_paths_cyto=None, coords=None, class_dict_nuc=None, class_dict_cyto=None)#

Turn predicted maps into individual nuclei and cell-type masks.

Parameters:
  • x (Dict[str, SoftSemanticOutput | SoftInstanceOutput]) – Dictionary returned by predict, including a nucleus entry.

  • use_async_postproc (bool) – Use asynchronous worker tasks for batches larger than one. This call waits for their results before returning. False uses synchronous workers; a single image always runs serially.

  • start_method (str) – Worker mode: threading, fork or spawn. Available process modes depend on the operating system.

  • n_jobs (int) – Number of reconstruction workers for batches larger than one.

  • save_paths_nuc (List[Path | str]) – One vector-output path per image for nuclei. Without paths, return masks in memory. Saving returns None for that image.

  • save_paths_cyto (List[Path | str]) – Equivalent paths for configured cytoplasm outputs.

  • coords (List[Tuple[int, int, int, int]]) – One (x, y, width, height) tuple per image. Vector coordinates are offset by x and y when saving.

  • class_dict_nuc (Dict[int, str]) – Map nucleus class numbers to names in saved vector data.

  • class_dict_cyto (Dict[int, str]) – Map cytoplasm class numbers to names in saved vector data.

Returns:

Dictionary with nuc, cyto and tissue entries. Nucleus and cytoplasm results are lists of (instances, classes) NumPy mask pairs, one per image. Both masks are int32 (H, W). Instance 0 is background; positive IDs identify objects. Each object gets the most common predicted class within its mask. Unused outputs and saved-image results are None. Tissue results are class masks. Tensor maps are copied to CPU.

Raises:

ValueError – Inference has not been enabled.

Return type:

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

Note

Use PostProcessor directly for tissue saving or custom contour smoothing. The selected model family controls the reconstruction method.

cellseg_models_pytorch.models.instanseg.InstanSeg.model#

PyTorch network used for training.

cellseg_models_pytorch.models.instanseg.InstanSeg.device#

Device used by the model.

cellseg_models_pytorch.models.instanseg.InstanSeg.inference_mode#

Whether inference has been set up. Calling .model.train() does not reset this flag.

cellseg_models_pytorch.models.instanseg.InstanSeg.post_processor#

Post-processor created by set_inference_mode. Replace it to change reconstruction settings.

ONNX export: export_onnx() in the export reference.