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)

Model Description
UNeXt2 ConvNeXtV2-based encoder-decoder with 3D-to-2D stem and pixel shuffle head
FullyConvolutionalMAE Fully convolutional masked autoencoder for self-supervised pretraining
Unet2d Classic 2D UNet with configurable depth and residual connections
Unet25d 2.5D UNet with learned 3D-to-2D compression via skip interruption
Unet3d 3D U-Net (F-Net, Ounkomol et al. 2018) with recursive encoder-decoder. All spatial dims must be divisible by 2^depth.

Contrastive Encoders (viscy_models.contrastive)

Model Description
ContrastiveEncoder timm backbone (ConvNeXt/ResNet) with 3D-to-2D stem and projection MLP
ResNet3dEncoder MONAI ResNetFeatures for native 3D contrastive learning

Variational Autoencoders (viscy_models.vae)

Model Description
BetaVae25D 2.5D beta-VAE with timm encoder and custom decoder
BetaVaeMonai Beta-VAE wrapping MONAI's VarAutoEncoder

Heads (viscy_models.components.heads)

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

Class Description
BaseHead Abstract base — subclass to add custom heads. Defines forward, compute_loss, log_metrics
ClassificationHead Classification on backbone features. Uses MLP + optional CosineClassifier. Logs top-1 and top-k accuracy
MLP Configurable projection/classification MLP with BN/LN and dropout
CosineClassifier L2-normalised linear head with learnable temperature — recommended for large class counts

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.

Project links

Data verified by PyPI on Aug 20, 2026
Data provided by the project maintainers, verified at the time the release was uploaded to PyPI.

Key dates

PyPI data
Data sourced directly from PyPI's database.
  • Released:
Latest release

1 maintainer

PyPI data
Data sourced directly from PyPI's database.
Avatar for MehtaLab from gravatar.com MehtaLab

Credits

Author: Biohub

GitHub Statistics

Data verified by PyPI on Aug 20, 2026
The GitHub source repository was provided by the project maintainers and verified by PyPI at the time of upload. Stars, forks, and open issues/PRs are derived from that repository and have not been independently verified.

License expression

BSD-3-Clause
View SPDX License List

Requires

Python >=3.12

Provides Extra

celldiff

Tags

deep learning microscopy neural networks pytorch representation learning virtual staining

Classifiers

Report project as malware

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

viscy_models-0.1.0.tar.gz (76.7 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

viscy_models-0.1.0-py3-none-any.whl (85.2 kB view details)

Uploaded Python 3

File details

Details for the file viscy_models-0.1.0.tar.gz.

File metadata

  • Download URL: viscy_models-0.1.0.tar.gz
  • Upload date:
  • Size: 76.7 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for viscy_models-0.1.0.tar.gz
Algorithm Hash digest
SHA256 4c97f8a9f63cd80db7b59fadaf1afedc555aa04ab4d91636a059f148f5e31238
MD5 ba35275b3e3d288327a2f339c9fbb6fd
BLAKE2b-256 5d4a5a0bdaa43e1615d2febeb26765fa73ebb3ae319bcc548a661d479116ccc9

See more details on using hashes here.

Provenance

The following attestation bundles were made for viscy_models-0.1.0.tar.gz:

Publisher: release.yml on mehta-lab/VisCy

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file viscy_models-0.1.0-py3-none-any.whl.

File metadata

  • Download URL: viscy_models-0.1.0-py3-none-any.whl
  • Upload date:
  • Size: 85.2 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for viscy_models-0.1.0-py3-none-any.whl
Algorithm Hash digest
SHA256 71ca20a41f87619b5bf460834e9d7dde1f8dbae4e775e40738d34e9fd3a1f742
MD5 a265e5733f52d65a949232c6faa51f95
BLAKE2b-256 d071a2fb1f19a42aa00af838c22dc3ca182367d02e1b1227d4ab9b831d03611e

See more details on using hashes here.

Provenance

The following attestation bundles were made for viscy_models-0.1.0-py3-none-any.whl:

Publisher: release.yml on mehta-lab/VisCy

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Release history Release notifications | RSS feed

This release

0.1.0 This release

2 files

Pre-release

0.1.0a0

2 files

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page