viscy-models

Neural network architectures for virtual staining microscopy.

Part of the VisCy project.

Installation

From PyPI (when published)

pip install viscy-models

For development (from monorepo root)

# Using uv (recommended)
uv pip install -e packages/viscy-models

# Or via workspace sync
uv sync --package viscy-models

Usage

from viscy_models import UNeXt2, FullyConvolutionalMAE

# All models are pure nn.Modules — no Lightning or Hydra coupling
model = UNeXt2(
    in_channels=1,
    out_channels=2,
    in_stack_depth=5,
    backbone="convnextv2_tiny",
)

# Contrastive encoders
from viscy_models import ContrastiveEncoder, ResNet3dEncoder

encoder = ContrastiveEncoder(
    backbone="convnext_tiny",
    in_channels=2,
    in_stack_depth=15,
    embedding_dim=768,
    projection_dim=128,
)

# Variational autoencoders
from viscy_models import BetaVae25D, BetaVaeMonai

Models

UNet Family (viscy_models.unet)

Contrastive Encoders (viscy_models.contrastive)

Variational Autoencoders (viscy_models.vae)

Heads (viscy_models.components.heads)

Pluggable task heads for multi-task learning. Attach to ContrastiveModule via auxiliary_heads.

Label routing via SampleMeta

Auxiliary heads consume labels from anchor_meta["labels"] in the batch — a dict[str, int] populated by the dataset. The batch_key on each head selects which label to use:

# Dataset populates anchor_meta with integer labels
anchor_meta = [{"labels": {"condition": 0, "gene_ko": 42}}]

# Head config — batch_key must match a key in anchor_meta["labels"]
ClassificationHead(
    head_name="gene_ko",      # used for logging
    batch_key="gene_ko",      # key in anchor_meta["labels"]
    in_dims=768,
    hidden_dims=512,
    num_classes=1001,
    loss_weight=0.5,
)

Features

• Pure nn.Module architectures — no Lightning or Hydra dependencies
• Shared components in components/ (stems, heads, decoder blocks, ConvBlocks)
• Pluggable auxiliary heads via BaseHead — extend for custom losses and metrics
• State dict key compatibility with original VisCy checkpoints
• Immutable defaults for all model constructors

Dependencies

• torch>=2.10
• timm>=1.0.15
• monai>=1.5.2
• numpy>=2.4.1

Documentation

In the works!

License

BSD-3-Clause - see LICENSE in repository root.