SporaModelWrapper#

class spora_bench.wrapper.SporaModelWrapper(model_name)#

Parameters:: model_name (str)

compute_cell_tokens(dataset, tissue_id)#

Compute cell tokens for the given dataset and tissue ID.

Parameters:

dataset (MultiplexImagingDataset) – The input dataset containing the multiplexed images and segmentation masks.
tissue_id (str) – The ID of the tissue to compute cell tokens for.

Returns:

torch.Tensor: The computed cell token tensor. Shape: (N, D)

embed_tile(tissue)#

Embed a tile from the input image.

Parameters:: tissue (MultiplexTissue) – The input tissue containing the image and protein IDs.
Returns:: torch.Tensor: The embedded tile tensor. Shape: (D,)

embed_tissue(dataset, tissue_id, tissue_threshold=0.3)#

Embed the tissue from the input image.

Parameters:

dataset (MultiplexImagingDataset) – The input dataset containing the multiplexed images and segmentation masks.
tissue_id (str) – The ID of the tissue to embed.
tissue_threshold (float) – The threshold for considering a tile as containing tissue.

Returns:

torch.Tensor: A sequence-shaped embedding of the tissue. Shape: (N,D)

postprocess_tile_embeddings(tissue_embedding)#

Postprocess the tissue embedding if necessary (e.g., for dimensionality reduction for kronos patient level tasks).

Parameters:: tissue_embedding (torch.Tensor) – The raw tissue embedding tensor. Shape: (N, D)
Returns:: torch.Tensor: The postprocessed tissue token tensor. Shape: (N, D’)

predict_marker(tissue, target_channel_name, target_uniprot_id=None)#

Predict a target marker given multipelxed input image.

Parameters:

tissue (MultiplexTissue) – The input multiplexed tissue. Shape: (C, H, W)
target_channel_name (str) – The name of the target channel to be inpainted.
target_uniprot_id (str) – The uniprot ID of the target channel to be inpainted.

Returns:

torch.Tensor: The predicted single-marker image. Shape: (1, H, W)

Raises:

MarkerNotSupportedError – If the target uniprot ID is not supported by the model.

predict_instance_segmentation(tissue)#

Predicts the instance segmentation mask for the given tissue. Cell instances are enumerated with unique integers starting from 1.

Parameters:: tissue (MultiplexTissue) – The input multiplexed tissue. Shape: (C, H, W)
Returns:: torch.Tensor: The predicted instance segmentation mask. Shape: (H, W)

predict_cell_types(tissue, mask, exclude_classes=None)#

Predicts the cell types for the given tissue. Cell types are enumerated with unique integers starting from 1.

Parameters:

tissue (MultiplexTissue) – The input multiplexed tissue. Shape: (C, H, W)
mask (CellMask) – The input cell mask. Shape: (H, W)
exclude_classes (List[str]) – A list of cell types to exclude.

Returns:

torch.Tensor: The predicted cell type class ids for each cell instance in the mask. Shape: (num_cells,) where num_cells is the number of unique cells in the mask.

predict_markers_from_he(he_tile, target_markers=None)#

Predict marker maps from an H&E tile alone (no multiplex input required). Unlike predict_marker, this is for models that translate H&E directly into multiplex marker channels (e.g. ROSIE, HistoPlexer, GigaTIME), rather than inpainting a dropped multiplex channel from the remaining multiplex channels.

Implementations should canonicalize their native channel vocabulary into the shared marker tokens used across models/datasets (see spora_bench.utils.virtual_staining_utils.canon/build_marker_index), and are expected to expose the canonical markers they support via a supported_he_markers attribute (a set of canonical marker name strings) set in __init__, so callers can compute the scored marker set without running inference.

Parameters:

he_tile (torch.Tensor) – Imagenet-normalized H&E tile. Shape: (3, H, W)
target_markers (Optional[Iterable[str]]) – Canonical marker names to restrict/optimize prediction for. Implementations may ignore this and return every marker they support.

Returns:

Dict[str, torch.Tensor]: Canonical marker name -> predicted map. Shape of each value: (H, W)

predict_markers_from_he_batch(he_tiles, target_markers=None)#

Predict marker maps for a batch of H&E tiles. Default implementation loops over predict_markers_from_he one tile at a time; override for models that benefit from batching inference across tiles (e.g. patch-based models like ROSIE).

Parameters:

he_tiles (List[torch.Tensor]) – Imagenet-normalized H&E tiles, each of shape (3, H, W).
target_markers (Optional[Iterable[str]]) – See predict_markers_from_he.

Returns:

List[Dict[str, torch.Tensor]]: One canonical-marker-name -> predicted map dict per input tile.