cellmil.models.segmentation.cellposeV2

Classes

CellposeSAM([pretrained_model, device, ...])

Simplified PyTorch wrapper around Cellpose that operates directly on torch tensors with true batching for improved performance.

class cellmil.models.segmentation.cellposeV2.CellposeSAM(pretrained_model: str = 'cpsam', device: Optional[device] = None, use_bfloat16: bool = True)

Bases: Module

Simplified PyTorch wrapper around Cellpose that operates directly on torch tensors with true batching for improved performance. Designed for patch-based processing.

__init__(pretrained_model: str = 'cpsam', device: Optional[device] = None, use_bfloat16: bool = True)

Initialize CellposeSAMV2 wrapper.

Parameters:

• pretrained_model – Path to pretrained cellpose model or model name

• device – Device to run the model on

• use_bfloat16 – Use bfloat16 precision for model weights

_init_cellpose_network()

Initialize the Cellpose network directly without the wrapper.

forward(x: Tensor, normalize: bool = True, resample: bool = True, niter: int = 200, flow_threshold: float = 0.4, cellprob_threshold: float = 0.0, min_size: int = 15, max_size_fraction: float = 0.4) → Dict[str, Tensor]

Forward pass through Cellpose model with true batched processing.

Parameters:

• x – Input tensor of shape (B, C, H, W) where C=3 (RGB channels)

• normalize – Whether to normalize input

• resample – Whether to resize flows and cellprob back to original image size

• niter – Number of iterations for mask refinement

• flow_threshold – Threshold for flow field

• cellprob_threshold – Threshold for cell probability map

• min_size – Minimum size of masks to keep

• max_size_fraction – Maximum size fraction of masks to keep

Returns:

• masks: Instance segmentation masks (B, H, W)

• flows: Flow fields (B, H, W, 2)

• cellprob: Cell probability maps (B, H, W)

• styles: Style vectors (B, style_dim)

Return type:

Dictionary containing

_compute_masks(flows: ndarray[Any, Any], cellprob: ndarray[Any, Any], flow_threshold: float = 0.4, cellprob_threshold: float = 0.0, min_size: int = 15, max_size_fraction: float = 0.4, niter: int = 200) → ndarray[Any, Any]

Compute masks from flows and cell probabilities using cellpose dynamics.

_normalize_batch(x: Tensor) → Tensor

Normalize batch of images using torch operations.

_resize_flows_batch(flows: Tensor, Ly: int, Lx: int) → Tensor

Resize flow fields to target dimensions.

_resize_cellprob_batch(cellprob: Tensor, Ly: int, Lx: int) → Tensor

Resize cell probability maps to target dimensions.

eval()

Set model to evaluation mode.

train(mode: bool = True)

Set model to training mode (cellpose doesn’t support training).

to(device: Union[device, str])

Move model to specified device.

cuda(device: Optional[Union[int, device]] = None)

Move model to CUDA device.

cpu()

Move model to CPU.

calculate_instance_map(predictions: Dict[str, Tensor], magnification: float = 40.0) → Tuple[Tensor, list[Dict[int, Dict[str, Any]]]]

Calculate instance map and extract cell information from predictions.

Parameters:

• predictions – Dictionary containing model outputs

• magnification – Magnification level

Returns:

Tuple containing instance map and cell information

_extract_contour(mask: ndarray[Any, Any]) → ndarray[Any, Any]

Extract contour from binary mask.