You can deploy this model and start earning money today!

recursionpharma/
OpenPhenom

Name: OpenPhenom
Availability: InStock

Views

No views yet

Model Card for OpenPhenom-S/16

Channel-agnostic image encoding model CA-MAE with a ViT-S/16 encoder backbone designed for microscopy image featurization. The model uses a vision transformer backbone with channelwise cross-attention over patch tokens to create contextualized representations separately for each channel.

Model Details

Model Description

This model is a channel-agnostic masked autoencoder trained to reconstruct microscopy images over three datasets:

RxRx3
JUMP-CP overexpression
JUMP-CP gene-knockouts

Developed, funded, and shared by: Recursion
Model type: Vision transformer CA-MAE
Image modality: Optimized for microscopy images from the CellPainting assay
License: Non-Commercial End User License Agreement

Installation

Requires Python 3.10.4 or higher. From a clone of the repository:

cd /path/to/OpenPhenom
pip install -e .

Model Sources

Repository: https://github.com/recursionpharma/maes_microscopy
Paper: Masked Autoencoders for Microscopy are Scalable Learners of Cellular Biology

Uses

NOTE: model embeddings tend to extract features only after using standard batch correction post-processing techniques. We recommend, at a minimum, after inferencing the model over your images, to do the standard PCA-CenterScale pattern or better yet Typical Variation Normalization:

Fit a PCA kernel on all the control images (or all images if no controls) from across all experimental batches (e.g. the plates of wells from your assay),
Transform all the embeddings with that PCA kernel,
For each experimental batch, fit a separate StandardScaler on the transformed embeddings of the controls from step 2, then transform the rest of the embeddings from that batch with that StandardScaler.

Direct Use

Create biologically useful embeddings of microscopy images
Create contextualized embeddings of each channel of a microscopy image (set return_channelwise_embeddings=True)
Leverage the full MAE encoder + decoder to predict new channels / stains for images without all 6 CellPainting channels

Downstream Use

A determined ML expert could fine-tune the encoder for downstream tasks such as classification

Out-of-Scope Use

Unlikely to be especially performant on brightfield microscopy images
Out-of-domain medical images, such as H&E (maybe it would be a decent baseline though)

Bias, Risks, and Limitations

Primary limitation is that the embeddings tend to be more useful at scale. For example, if you only have 1 plate of microscopy images, the embeddings might underperform compared to a supervised bespoke model.

How to Get Started with the Model

You should be able to successfully run the below tests, which demonstrate how to use the model at inference time.

python

import pytest
import torch

from huggingface_mae import MAEModel

# huggingface_openphenom_model_dir = "."
huggingface_modelpath = "recursionpharma/OpenPhenom"


@pytest.fixture
def huggingface_model():
    # This step downloads the model to a local cache, takes a bit to run
    huggingface_model = MAEModel.from_pretrained(huggingface_modelpath)
    huggingface_model.eval()
    return huggingface_model


@pytest.mark.parametrize("C", [1, 4, 6, 11])
@pytest.mark.parametrize("return_channelwise_embeddings", [True, False])
def test_model_predict(huggingface_model, C, return_channelwise_embeddings):
    example_input_array = torch.randint(
        low=0,
        high=255,
        size=(2, C, 256, 256),
        dtype=torch.uint8,
        device=huggingface_model.device,
    )
    huggingface_model.return_channelwise_embeddings = return_channelwise_embeddings
    embeddings = huggingface_model.predict(example_input_array)
    expected_output_dim = 384 * C if return_channelwise_embeddings else 384
    assert embeddings.shape == (2, expected_output_dim)

We also provide a notebook for running inference on RxRx3-core.

Training, evaluation and testing details

See paper linked above for details on model training and evaluation. Primary hyperparameters are included in the repo linked above.

Environmental Impact

Hardware Type: Nvidia H100 Hopper nodes
Hours used: 400
Cloud Provider: private cloud
Carbon Emitted: 138.24 kg co2 (roughly the equivalent of one car driving from Toronto to Montreal)

BibTeX:

TeX

@inproceedings{kraus2024masked,
  title={Masked Autoencoders for Microscopy are Scalable Learners of Cellular Biology},
  author={Kraus, Oren and Kenyon-Dean, Kian and Saberian, Saber and Fallah, Maryam and McLean, Peter and Leung, Jess and Sharma, Vasudev and Khan, Ayla and Balakrishnan, Jia and Celik, Safiye and others},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={11757--11768},
  year={2024}
}

Model Card Contact

Kian Kenyon-Dean: kian.kd@recursion.com
Oren Kraus: oren.kraus@recursion.com
Or, email: info@rxrx.ai