Model networks and factories#

Create a PyTorch network for training. Start with a factory function such as cellpose_nuclei; use a network class when you need custom outputs. Set enc_pretrain=False to avoid downloading encoder weights.

import torch

from cellseg_models_pytorch.models.cellpose.cellpose_unet import cellpose_nuclei

network = cellpose_nuclei(3, enc_name="resnet18", enc_pretrain=False)
network.eval()
with torch.inference_mode():
    output = network(torch.rand(1, 3, 64, 64))["nuc"]

This example uses an untrained network. output.type_map contains class scores (logits) with shape (1, 3, 64, 64). output.aux_map contains two flow maps with shape (1, 2, 64, 64). For image prediction and reconstruction, use a model class.

Factory	Returned network	Maps used to separate objects	Separate foreground output
`cellpose_nuclei`	`CellPoseUnet`	2 flows, vertical then horizontal	No
`omnipose_nuclei`	`CellPoseUnet`	2 Omnipose flows, vertical then horizontal	No
`hovernet_nuclei`	`HoverNetUnet`	2 position maps, vertical then horizontal targets	No
`cellvit_nuclei`	`CellVitSamUnet`	2 position maps, vertical then horizontal targets	No
`stardist_nuclei`	`StarDistUnet`	`n_rays` radial distances	1 channel
`instanseg_nuclei`	`InstanSegUnet`	4 reconstruction channels	No

Outputs use SoftInstanceOutput. See inference for the prediction fields. Set return_pred_only=False if you also need encoder and decoder features.

For encoder freezing and a full training loop, see the training guide. Use standard PyTorch module methods to change training mode, move the network between devices and access its parameters.

Configure a custom network#

decoders names the decoder branches. heads names the outputs from each branch and sets their channel counts. Use nuc_, cyto_ or tissue_ prefixes and the reconstruction maps required by your model family. Stage settings such as out_channels must match the network depth. Changing the architecture requires compatible weights.

For CNN encoders such as resnet18, use image dimensions divisible by 32. CellVit supports samvit_* encoders, which have their own image-size requirements. For InstanSeg, use the factory or model class: these also create the pixel classifier needed for reconstruction.

CellPose family#

from cellseg_models_pytorch.models.cellpose.cellpose_unet import (
    CellPoseUnet,
    cellpose_nuclei,
    omnipose_nuclei,
)

cellpose_nuclei(n_nuc_classes, **kwargs)#

Create a CellPose network with standard nucleus heads.

https://www.nature.com/articles/s41592-020-01018-x

Parameters:

n_nuc_classes (int) – Number of cell-type output channels, including any background.
**kwargs – Network options accepted by CellPoseUnet, such as enc_name and enc_pretrain=False for offline construction. Decoder branches and head definitions are set by this factory.

Returns:

A CellPoseUnet PyTorch module. Call it with a normalized RGB batch to obtain raw output maps; it does not reconstruct individual instances.

Return type:

Module

omnipose_nuclei(n_nuc_classes, **kwargs)#

Create an Omnipose network with standard nucleus heads.

https://www.biorxiv.org/content/10.1101/2021.11.03.467199v2

Parameters:

n_nuc_classes (int) – Number of cell-type output channels, including any background.
**kwargs – Network options accepted by CellPoseUnet, such as enc_name and enc_pretrain=False for offline construction. Decoder branches and head definitions are set by this factory.

Returns:

A CellPoseUnet PyTorch module. Call it with a normalized RGB batch to obtain raw output maps; it does not reconstruct individual instances.

Return type:

Module

class CellPoseUnet(decoders, heads, depth=4, out_channels=(256, 128, 64, 32), layer_depths=(4, 4, 4, 4), style_channels=256, enc_name='efficientnet_b5', enc_pretrain=True, enc_freeze=False, enc_out_indices=None, upsampling='bilinear', long_skip='unet', merge_policy='sum', short_skip='basic', normalization='bn', activation='relu', convolution='conv', preactivate=True, attention=None, preattend=False, out_size=None, encoder_kws=None, skip_kws=None, stem_skip_kws=None, inst_key='type', **kwargs)#

A CellPose family network with configurable output branches.

Use the family factory for standard nucleus heads. Direct construction is for choosing your own branches and head channel counts. enc_name names the encoder stored at network[network.enc_name].

Configure the CellPoseUnet network and its output heads.

Based on CellPose and Omnipose. This version uses a timm encoder and separate output heads.

Parameters:

decoders (Tuple[str, ...]) – Names of decoder branches. Each name must have an entry in heads.
heads (Dict[str, Dict[str, int]]) – Output names and channel counts grouped by decoder. For example, {“type”: {“nuc_cellpose”: 2, “nuc_type”: 3}} for CellPose. Head names must match the chosen family and output type.
depth (int) – Number of decoder stages; align stage settings and encoder features.
out_channels (Tuple[int, ...]) – Channel counts at each decoder stage, from deepest to shallowest.
layer_depths (Tuple[int, ...]) – Number of convolution blocks at each decoder stage.
style_channels (int) – Size of the vector describing image appearance. None disables its use.
enc_name (str) – timm encoder name. Access it through network[network.enc_name].
enc_pretrain (bool) – Load pretrained encoder weights; may download on first use.
enc_freeze (bool) – Disable encoder weight gradients; does not keep it in evaluation mode.
enc_out_indices (Tuple[int, ...]) – Encoder feature indices. None selects range(depth).
upsampling (str) – Method for increasing feature resolution between decoder stages.
long_skip (str) – Connections from encoder features to decoder stages, such as “unet”.
merge_policy (str) – Combine connected features by “sum” or “cat” (concatenation).
short_skip (str) – Connections within decoder blocks: “basic”, “residual” or “dense”.
normalization (str) – Feature normalization, such as “bn” (batch normalization), or None.
activation (str) – Activation function used within decoder blocks, such as “relu”.
convolution (str) – Convolution implementation: “conv”, “wsconv” or “scaled_wsconv”.
preactivate (bool) – Apply normalization and activation before the convolution.
attention (str) – Channel attention method, or None to disable it.
preattend (bool) – Apply configured attention before the convolution block.
out_size (int) – Optional square output size in pixels. None keeps the decoder output size.
encoder_kws (Dict[str, Any]) – Additional keyword arguments passed to the timm encoder.
skip_kws (Dict[str, Any]) – Additional settings for the selected skip-connection implementation.
stem_skip_kws (Dict[str, Any]) – Additional settings for the connection from the input image to the decoder.
inst_key (str) – Output kind used to identify the instance foreground head.
**kwargs – Additional keyword arguments; do not use this to override explicit arguments.

forward(x, return_pred_only=True)#

Compute raw CellPose prediction maps for training or inference.

Parameters:

x (Tensor) – Floating-point RGB batch with shape (B, 3, H, W), on the network’s device and dtype. B is batch size, H height and W width. Apply the normalization used during training before calling the network.
return_pred_only (bool) – False also includes encoder/decoder features and encoder output under enc_feats, dec_feats and enc_out.

Returns:

Dictionary with nuc, cyto and tissue entries; unused outputs are None. Instance outputs contain raw class scores (logits) in type_map, maps for separating objects in aux_map, and optionally foreground values in binary_map. Map shape is (B, channels, H, W), using the configured output size. Maps keep the network device, dtype and gradients for training. Features are included only when requested.

Return type:

Dict[str, Any]

cellseg_models_pytorch.models.cellpose.cellpose_unet.CellPoseUnet.enc_name#: Encoder name used for access through network[network.enc_name].

StarDist family#

from cellseg_models_pytorch.models.stardist.stardist_unet import (
    StarDistUnet,
    stardist_nuclei,
)

stardist_nuclei(n_rays, n_nuc_classes, **kwargs)#

Create a StarDist network with standard nucleus heads.

https://arxiv.org/abs/1806.03535

Parameters:

n_rays (int) – Number of radial distances per pixel; match the training targets.
n_nuc_classes (int) – Number of cell-type output channels, including any background.
**kwargs – Network options accepted by StarDistUnet, such as enc_name and enc_pretrain=False for offline construction. Decoder branches and head definitions are set by this factory.

Returns:

A StarDistUnet PyTorch module. Call it with a normalized RGB batch to obtain raw output maps; it does not reconstruct individual instances.

Return type:

Module

class StarDistUnet(decoders, heads, depth=4, out_channels=(256, 128, 64, 32), style_channels=None, enc_name='efficientnet_b5', enc_pretrain=True, enc_freeze=False, enc_out_indices=None, upsampling='fixed-unpool', long_skip='unet', merge_policy='cat', short_skip='basic', block_type='basic', normalization=None, activation='relu', convolution='conv', preactivate=False, attention=None, preattend=False, out_size=None, encoder_kws=None, skip_kws=None, stem_skip_kws=None, inst_key='dist', **kwargs)#

A StarDist family network with configurable output branches.

Use the family factory for standard nucleus heads. Direct construction is for choosing your own branches and head channel counts. enc_name names the encoder stored at network[network.enc_name].

Configure the StarDistUnet network and its output heads.

Stardist:

https://arxiv.org/abs/1806.03535

Note

Minor differences from the original implementation. - long skip concatenation/sum applied before each conv layer rather than after.

Parameters:

decoders (Tuple[str, ...]) – Names of decoder branches. Each name must have an entry in heads.
heads (Dict[str, Dict[str, int]]) – Output names and channel counts grouped by decoder. Use nuc_stardist for the auxiliary head and nuc_type for cell-type logits; ray-based reconstruction also needs nuc_binary.
depth (int) – Number of decoder stages; align stage settings and encoder features.
out_channels (Tuple[int, ...]) – Channel counts at each decoder stage, from deepest to shallowest.
style_channels (int) – Size of the vector describing image appearance. None disables its use.
enc_name (str) – timm encoder name. Access it through network[network.enc_name].
enc_pretrain (bool) – Load pretrained encoder weights; may download on first use.
enc_freeze (bool) – Disable encoder weight gradients; does not keep it in evaluation mode.
enc_out_indices (Tuple[int, ...]) – Encoder feature indices. None selects range(depth).
upsampling (str) – Method for increasing feature resolution between decoder stages.
long_skip (str) – Connections from encoder features to decoder stages, such as “unet”.
merge_policy (str) – Combine connected features by “sum” or “cat” (concatenation).
short_skip (str) – Connections within decoder blocks: “basic”, “residual” or “dense”.
block_type (str) – Convolution block type used within each decoder stage.
normalization (str) – Feature normalization, such as “bn” (batch normalization), or None.
activation (str) – Activation function used within decoder blocks, such as “relu”.
convolution (str) – Convolution implementation: “conv”, “wsconv” or “scaled_wsconv”.
preactivate (bool) – Apply normalization and activation before the convolution.
attention (str) – Channel attention method, or None to disable it.
preattend (bool) – Apply configured attention before the convolution block.
out_size (int) – Optional square output size in pixels. None keeps the decoder output size.
encoder_kws (Dict[str, Any]) – Additional keyword arguments passed to the timm encoder.
skip_kws (Dict[str, Any]) – Additional settings for the selected skip-connection implementation.
stem_skip_kws (Dict[str, Any]) – Additional settings for the connection from the input image to the decoder.
inst_key (str) – Output kind used to identify the instance foreground head.
**kwargs – Additional keyword arguments are accepted but unused.

forward(x, return_pred_only=True)#

Compute raw StarDist prediction maps for training or inference.

Parameters:

x (Tensor) – Floating-point RGB batch with shape (B, 3, H, W), on the network’s device and dtype. B is batch size, H height and W width. Apply the normalization used during training before calling the network.
return_pred_only (bool) – False also includes encoder/decoder features and encoder output under enc_feats, dec_feats and enc_out.

Returns:

Return type:

Dict[str, Any]

cellseg_models_pytorch.models.stardist.stardist_unet.StarDistUnet.enc_name#: Encoder name used for access through network[network.enc_name].

HoverNet family#

from cellseg_models_pytorch.models.hovernet.hovernet_unet import (
    HoverNetUnet,
    hovernet_nuclei,
)

hovernet_nuclei(n_nuc_classes, **kwargs)#

Create a HoverNet network with standard nucleus heads.

https://www.sciencedirect.com/science/article/pii/S1361841519301045?via%3Dihub

Parameters:

n_nuc_classes (int) – Number of cell-type output channels, including any background.
**kwargs – Network options accepted by HoverNetUnet, such as enc_name and enc_pretrain=False for offline construction. Decoder branches and head definitions are set by this factory.

Returns:

A HoverNetUnet PyTorch module. Call it with a normalized RGB batch to obtain raw output maps; it does not reconstruct individual instances.

Return type:

Module

class HoverNetUnet(decoders, heads, depth=4, out_channels=(512, 256, 64, 64), style_channels=None, enc_name='efficientnet_b5', enc_pretrain=True, enc_freeze=False, enc_out_indices=None, upsampling='fixed-unpool', long_skip='unet', merge_policy='sum', n_dense=(8, 4), normalization='bn', activation='relu', convolution='conv', preactivate=True, attention=None, preattend=False, out_size=None, encoder_kws=None, skip_kws=None, stem_skip_kws=None, inst_key='inst', **kwargs)#

A HoverNet family network with configurable output branches.

Use the family factory for standard nucleus heads. Direct construction is for choosing your own branches and head channel counts. enc_name names the encoder stored at network[network.enc_name].

Configure the HoverNetUnet network and its output heads.

Based on HoVer-Net. This version supports timm encoders and uses transition convolutions between dense blocks.

Parameters:

decoders (Tuple[str, ...]) – Names of decoder branches. Each name must have an entry in heads.
heads (Dict[str, Dict[str, int]]) – Output names and channel counts grouped by decoder. Use nuc_hovernet for the auxiliary head and nuc_type for cell-type logits.
depth (int) – Number of decoder stages; align stage settings and encoder features.
out_channels (Tuple[int, ...]) – Channel counts at each decoder stage, from deepest to shallowest.
style_channels (int) – Size of the vector describing image appearance. None disables its use.
enc_name (str) – timm encoder name. Access it through network[network.enc_name].
enc_pretrain (bool) – Load pretrained encoder weights; may download on first use.
enc_freeze (bool) – Disable encoder weight gradients; does not keep it in evaluation mode.
enc_out_indices (Tuple[int, ...]) – Encoder feature indices. None selects range(depth).
upsampling (str) – Method for increasing feature resolution between decoder stages.
long_skip (str) – Connections from encoder features to decoder stages, such as “unet”.
merge_policy (str) – Combine connected features by “sum” or “cat” (concatenation).
n_dense (Tuple[int, int]) – Number of densely connected blocks at each decoder stage.
normalization (str) – Feature normalization, such as “bn” (batch normalization), or None.
activation (str) – Activation function used within decoder blocks, such as “relu”.
convolution (str) – Convolution implementation: “conv”, “wsconv” or “scaled_wsconv”.
preactivate (bool) – Apply normalization and activation before the convolution.
attention (str) – Channel attention method, or None to disable it.
preattend (bool) – Apply configured attention before the convolution block.
out_size (int) – Optional square output size in pixels. None keeps the decoder output size.
encoder_kws (Dict[str, Any]) – Additional keyword arguments passed to the timm encoder.
skip_kws (Dict[str, Any]) – Additional settings for the selected skip-connection implementation.
stem_skip_kws (Dict[str, Any]) – Additional settings for the connection from the input image to the decoder.
inst_key (str) – Output kind used to identify the instance foreground head.
**kwargs – Additional keyword arguments are accepted but unused.

forward(x, return_pred_only=True)#

Compute raw HoverNet prediction maps for training or inference.

Parameters:

x (Tensor) – Floating-point RGB batch with shape (B, 3, H, W), on the network’s device and dtype. B is batch size, H height and W width. Apply the normalization used during training before calling the network.
return_pred_only (bool) – False also includes encoder/decoder features and encoder output under enc_feats, dec_feats and enc_out.

Returns:

Return type:

Dict[str, Any]

cellseg_models_pytorch.models.hovernet.hovernet_unet.HoverNetUnet.enc_name#: Encoder name used for access through network[network.enc_name].

CellVit family#

from cellseg_models_pytorch.models.cellvit.cellvit_unet import (
    CellVitSamUnet,
    cellvit_nuclei,
)

cellvit_nuclei(enc_name, n_nuc_classes, **kwargs)#

Create a CellVit network with standard nucleus heads.

https://arxiv.org/abs/2306.15350

Parameters:

enc_name (str) – Name of a supported timm SAM encoder.
n_nuc_classes (int) – Number of cell-type output channels, including any background.
**kwargs – Network options accepted by CellVitSamUnet, such as enc_name and enc_pretrain=False for offline construction. Decoder branches and head definitions are set by this factory.

Returns:

A CellVitSamUnet PyTorch module. Call it with a normalized RGB batch to obtain raw output maps; it does not reconstruct individual instances.

Return type:

Module

class CellVitSamUnet(decoders, heads, out_channels=(512, 256, 128, 64), layer_depths=(3, 2, 2, 2), style_channels=None, enc_name='samvit_base_patch16', enc_pretrain=True, enc_freeze=False, enc_out_channels=None, enc_out_indices=None, long_skip='unet', merge_policy='cat', short_skip='basic', normalization='bn', activation='relu', convolution='conv', preactivate=True, attention=None, preattend=False, out_size=None, encoder_kws=None, skip_kws=None, stem_skip_kws=None, inst_key='inst', **kwargs)#

A CellVit family network with configurable output branches.

Use the family factory for standard nucleus heads. Direct construction is for choosing your own branches and head channel counts. enc_name names the encoder stored at network[network.enc_name].

Configure the CellVitSamUnet network and its output heads.

CellVit:

https://arxiv.org/abs/2306.15350

Parameters:

decoders (Tuple[str, ...]) – Names of decoder branches. Each name must have an entry in heads.
heads (Dict[str, Dict[str, int]]) – Output names and channel counts grouped by decoder. Use nuc_hovernet for the auxiliary head and nuc_type for cell-type logits.
out_channels (Tuple[int, ...]) – Channel counts at each decoder stage, from deepest to shallowest.
layer_depths (Tuple[int, ...]) – Number of convolution blocks at each decoder stage.
style_channels (int) – Size of the vector describing image appearance. None disables its use.
enc_name (str) – One of “samvit_base_patch16”, “samvit_base_patch16_224”, “samvit_large_patch16” or “samvit_huge_patch16”. Access it through network[network.enc_name].
enc_pretrain (bool) – Load pretrained encoder weights; may download on first use.
enc_freeze (bool) – Disable encoder weight gradients; does not keep it in evaluation mode.
enc_out_channels (Tuple[int, ...]) – Retained compatibility option; currently does not affect feature upsampling.
enc_out_indices (Tuple[int, ...]) – Encoder feature indices. None selects one per decoder stage.
long_skip (str) – Connections from encoder features to decoder stages, such as “unet”.
merge_policy (str) – Combine connected features by “sum” or “cat” (concatenation).
short_skip (str) – Connections within decoder blocks: “basic”, “residual” or “dense”.
normalization (str) – Feature normalization, such as “bn” (batch normalization), or None.
activation (str) – Activation function used within decoder blocks, such as “relu”.
convolution (str) – Convolution implementation: “conv”, “wsconv” or “scaled_wsconv”.
preactivate (bool) – Apply normalization and activation before the convolution.
attention (str) – Channel attention method, or None to disable it.
preattend (bool) – Apply configured attention before the convolution block.
out_size (int) – Optional square output size in pixels. None keeps the decoder output size.
encoder_kws (Dict[str, Any]) – Additional keyword arguments passed to the timm encoder.
skip_kws (Dict[str, Any]) – Additional settings for the selected skip-connection implementation.
stem_skip_kws (Dict[str, Any]) – Additional settings for the connection from the input image to the decoder.
inst_key (str) – Output kind used to identify the instance foreground head.
**kwargs – Additional keyword arguments are accepted but unused.

forward(x, return_pred_only=True)#

Compute raw CellVit prediction maps for training or inference.

Parameters:

x (Tensor) – Floating-point RGB batch with shape (B, 3, H, W), on the network’s device and dtype. B is batch size, H height and W width. Apply the normalization used during training before calling the network.
return_pred_only (bool) – False also includes encoder/decoder features and encoder output under enc_feats, dec_feats and enc_out.

Returns:

Return type:

Dict[str, Any]

cellseg_models_pytorch.models.cellvit.cellvit_unet.CellVitSamUnet.enc_name#: Encoder name used for access through network[network.enc_name].

InstanSeg family#

from cellseg_models_pytorch.models.instanseg import (
    InstanSegUnet,
    instanseg_nuclei,
)

instanseg_nuclei(n_nuc_classes, **kwargs)#

Create an InstanSeg network with standard nucleus heads.

Includes the pixel classifier used by InstanSeg reconstruction.

Parameters:

n_nuc_classes (int) – Number of cell-type output channels, including any background.
**kwargs – Network options accepted by InstanSegUnet, such as enc_name and enc_pretrain=False for offline construction. Decoder branches and head definitions are set by this factory.

Returns:

A InstanSegUnet PyTorch module. Call it with a normalized RGB batch to obtain raw output maps; it does not reconstruct individual instances.

Return type:

Module

class InstanSegUnet(decoders, heads, depth=4, out_channels=(256, 128, 64, 32), style_channels=None, enc_name='efficientnet_b5', enc_pretrain=True, enc_freeze=False, enc_out_indices=None, upsampling='fixed-unpool', long_skip='unet', merge_policy='cat', short_skip='basic', block_type='basic', normalization=None, activation='relu', convolution='conv', preactivate=False, attention=None, preattend=False, out_size=None, encoder_kws=None, skip_kws=None, stem_skip_kws=None, inst_key='type', **kwargs)#

A InstanSeg family network with configurable output branches.

Use the family factory for standard nucleus heads. Direct construction is for choosing your own branches and head channel counts. enc_name names the encoder stored at network[network.enc_name].

Configure the InstanSeg network for training or prediction.

Parameters:

decoders (Tuple[str, ...]) – Names of decoder branches. Each name must have an entry in heads.
heads (Dict[str, Dict[str, int]]) – Output names and channel counts grouped by decoder. Use nuc_instanseg for the auxiliary head and nuc_type for cell-type logits; ray-based reconstruction also needs nuc_binary.
depth (int) – Number of decoder stages; align stage settings and encoder features.
out_channels (Tuple[int, ...]) – Channel counts at each decoder stage, from deepest to shallowest.
style_channels (int) – Size of the vector describing image appearance. None disables its use.
enc_name (str) – timm encoder name. Access it through network[network.enc_name].
enc_pretrain (bool) – Load pretrained encoder weights; may download on first use.
enc_freeze (bool) – Disable encoder weight gradients; does not keep it in evaluation mode.
enc_out_indices (Tuple[int, ...]) – Encoder feature indices. None selects range(depth).
upsampling (str) – Method for increasing feature resolution between decoder stages.
long_skip (str) – Connections from encoder features to decoder stages, such as “unet”.
merge_policy (str) – Combine connected features by “sum” or “cat” (concatenation).
short_skip (str) – Connections within decoder blocks: “basic”, “residual” or “dense”.
block_type (str) – Convolution block type used within each decoder stage.
normalization (str) – Feature normalization, such as “bn” (batch normalization), or None.
activation (str) – Activation function used within decoder blocks, such as “relu”.
convolution (str) – Convolution implementation: “conv”, “wsconv” or “scaled_wsconv”.
preactivate (bool) – Apply normalization and activation before the convolution.
attention (str) – Channel attention method, or None to disable it.
preattend (bool) – Apply configured attention before the convolution block.
out_size (int) – Optional square output size in pixels. None keeps the decoder output size.
encoder_kws (Dict[str, Any]) – Additional keyword arguments passed to the timm encoder.
skip_kws (Dict[str, Any]) – Additional settings for the selected skip-connection implementation.
stem_skip_kws (Dict[str, Any]) – Additional settings for the connection from the input image to the decoder.
inst_key (str) – Output kind used to identify the instance foreground head.
**kwargs – Additional keyword arguments are accepted but unused.

forward(x, return_pred_only=True)#

Compute raw InstanSeg prediction maps for training or inference.

Parameters:

x (Tensor) – Floating-point RGB batch with shape (B, 3, H, W), on the network’s device and dtype. B is batch size, H height and W width. Apply the normalization used during training before calling the network.
return_pred_only (bool) – False also includes encoder/decoder features and encoder output under enc_feats, dec_feats and enc_out.

Returns:

Return type:

Dict[str, Any]

cellseg_models_pytorch.models.instanseg.InstanSegUnet.enc_name#: Encoder name used for access through network[network.enc_name].