Add package code and model weights

.gitattributes

@@ -19,6 +19,7 @@
 *.pb filter=lfs diff=lfs merge=lfs -text
 *.pickle filter=lfs diff=lfs merge=lfs -text
 *.pkl filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,12 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+
+.vscode/

.python-version

@@ -0,0 +1 @@
+3.13

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) Microsoft Corporation.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE

README.md

@@ -1,5 +1,180 @@
 ---
 license: mit
+language:
+- en
+pipeline_tag: zero-shot-image-classification
 ---
 
-This repository will be populated very soon.
+# Model card for COLIPRI
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+COLIPRI is a 3D vision&ndash;language transformer model trained to encode chest CT scans and reports.
+
+## Model description
+
+<!-- Provide a longer summary of what this model is. -->
+
+COLIPRI was trained using tens of thousands of chest CT scans and reports, without any annotations, using multiple objectives to learn strong joint representations of 3D images and text.
+The procedure is described in detail in our manuscript, [_Comprehensive language-image pre-training for 3D medical image understanding_](https://arxiv.org/abs/2510.15042) (Wald et al. 2026).
+
+The weights shared here correspond to our best-performing model, COLIPRI-CRM.
+
+- **Developed by:** Microsoft Health Futures
+- **Model type:** 3D vision&ndash;language encoder
+- **License:** [MIT](./LICENSE)
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+COLIPRI is shared for research purposes only.
+It is **not meant to be used for clinical practice**.
+
+<!-- ### Downstream use -->
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+The encoders be plugged to other models, or used independently or jointly for many downstream tasks, such as:
+
+- Image classification with text prompts
+- Image clustering
+- Text clustering
+- Text-to-image retrieval
+- Image-to-image retrieval
+- Image-to-text retrieval
+- Text-to-text retrieval
+- Image classification with a classifier
+- Text classification with a classifier
+- Image segmentation with a decoder
+- Report generation with a language decoder
+
+Fine-tuning COLIPRI is typically not necessary to obtain good performance in downstream tasks.
+
+<!-- ### Out-of-scope use -->
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+## Biases, risks, and limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+COLIPRI was trained with data from Turkey and the USA only, therefore it might be biased towards population in the training data.
+Underlying biases of the training datasets may not be well characterized.
+
+## Installation
+
+```shell
+pip install git+https://huggingface.co/microsoft/colipri.git
+```
+
+## Usage examples
+
+First, let's get a 3D chest CT we can use for demonstration.
+The plotted slices intersect a lung nodule near the heart.
+
+```python
+>>> from colipri import load_sample_ct
+>>> image = load_sample_ct()
+>>> image
+ScalarImage(shape: (1, 512, 512, 139); spacing: (0.76, 0.76, 2.50); orientation: LPS+; dtype: torch.IntTensor; memory: 139.0 MiB)
+>>> nodule_indices_ct = 180, 341, 76
+>>> image.plot(indices=nodule_indices_ct)
+```
+
+![Input CT](assets/input.png)
+
+Now, let's instantiate the model and processor.
+
+```python
+>>> from colipri import get_model
+>>> from colipri import get_processor
+>>> model = get_model().cuda()
+>>> processor = get_processor()
+```
+
+### Feature extraction
+
+```python
+>>> import torch
+>>> preprocessed_images = processor.process_images(image)
+>>> preprocessed_images[0]
+ScalarImage(shape: (1, 192, 192, 192); spacing: (2.00, 2.00, 2.00); orientation: SAR+; dtype: torch.FloatTensor; memory: 27.0 MiB)
+>>> images_batch = processor.to_images_batch(preprocessed_images)
+images_batch.shape
+torch.Size([1, 1, 192, 192, 192])
+>>> with torch.no_grad():
+...     patch_embeddings = model.encode_image(images_batch)
+>>> patch_embeddings.shape
+torch.Size([1, 768, 24, 24, 24])
+>>> with torch.no_grad():
+...     pooled_embeddings = model.encode_image(images_batch, pool=True)
+>>> pooled_embeddings.shape
+torch.Size([1, 768])
+```
+
+### Zero-shot classification
+
+```python
+>>> from colipri import ZeroShotImageClassificationPipeline
+>>> pipeline = ZeroShotImageClassificationPipeline(model, processor)
+>>> pipeline(image, ["No lung nodules", "Lung nodules"])
+[
+    {'score': 0.005, 'label': 'No lung nodules'},
+    {'score': 0.995, 'label': 'Lung nodules'}
+]
+```
+
+## Environmental impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+<!-- Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700). -->
+
+- **Hardware type:** NVIDIA A100 GPUs
+- **Hours used:** 72 hours × 4 GPUs = 288 GPU-hours
+- **Cloud provider:** Azure
+- **Compute region:** West US 2
+- **Carbon emitted:** 21.6 kg CO₂ eq.
+
+### Compute infrastructure
+
+COLIPRI was trained on [Azure Machine Learning](https://azure.microsoft.com/en-us/products/machine-learning).
+
+#### Hardware
+
+| Stage | Node type | Num. nodes | GPU type | GPUs per node |
+| --- | --- | --- | --- | --- |
+| Pre-training | [`Standard_NC96ads_A100_v4`](https://learn.microsoft.com/en-us/azure/virtual-machines/sizes/gpu-accelerated/nca100v4-series?tabs=sizeaccelerators) | 1 | NVIDIA A100 (80 GB) | 4 |
+| Evaluation   | [`Standard_NC24ads_A100_v4`](https://learn.microsoft.com/en-us/azure/virtual-machines/sizes/gpu-accelerated/nca100v4-series?tabs=sizeaccelerators) | 1 | NVIDIA A100 (80 GB) | 1 |
+
+#### Software
+
+The main software libraries used in this work were [nnSSL](https://github.com/MIC-DKFZ/nnssl) for training, [TorchIO](https://torchio.org/) for preprocessing and augmentation, [`nifti-zarr-py`](https://github.com/neuroscales/nifti-zarr-py) for data loading, and [nnU-Net](https://github.com/MIC-DKFZ/nnUNet) for segmentation evaluation.
+
+## Citation
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+### BibTeX
+
+```bibtex
+@misc{
+    wald2026_colipri,
+    title={Comprehensive language-image pre-training for 3D medical image understanding},
+    author={Tassilo Wald and Ibrahim Ethem Hamamci and Yuan Gao and Sam Bond-Taylor and Harshita Sharma and Maximilian Ilse and Cynthia Lo and Olesya Melnichenko and Anton Schwaighofer and Noel C. F. Codella and Maria Teodora Wetscherek and Klaus H. Maier-Hein and Panagiotis Korfiatis and Valentina Salvatelli and Javier Alvarez-Valle and P{\'e}rez-Garc{\'i}a},
+    year={2026},
+    eprint={2510.15042},
+    archivePrefix={arXiv},
+    primaryClass={cs.CV},
+    url={https://arxiv.org/abs/2510.15042},
+}
+```
+
+### APA
+
+> Wald, T., Hamamci, I. E., Gao, Y., Bond-Taylor, S., Sharma, H., Ilse, M., Lo, C., Melnichenko, O., Schwaighofer, A., Codella, N. C. F., Wetscherek, M. T., Maier-Hein, K. H., Korfiatis, P., Salvatelli, V., Alvarez-Valle, J., & Pérez-García, F. (2026). Comprehensive language-image pre-training for 3D medical image understanding. arXiv. <https://doi.org/10.48550/ARXIV.2510.15042>
+
+## Model card contact
+
+Fernando Pérez-García ([`fperezgarcia@microsoft.com`](mailto:fperezgarcia@microsoft.com)).

justfile

@@ -0,0 +1,2 @@
+requirements:
+	uv export --quiet --no-hashes --no-emit-project --output-file requirements_demo.txt

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:97d66586507d1f48577b5149777d02e12d25787435a2a295cac5563df9291c14
+size 1033541380

model_card.md

@@ -0,0 +1,115 @@
+
+# Model card
+
+Date effective: September 2025
+Owner: ORA, Internal Policy & Practice
+
+## Model summary
+
+| | |
+| --- | --- |
+| Developer | Microsoft Research |
+| Description | COLIPRI is composed of a vision encoder for 3D CT scans and a text encoder for clinical reports. The encoders may be leveraged for downstream tasks such as classification or segmentation. The models were trained using CT-RATE and NLST. COLIPRI is not a generative model. |
+| Model architecture | The vision encoder is Primus-M, a 3D vision transformer. The text encoder architecture is BERT. |
+| Parameters | 1-500M |
+| Inputs | 3D chest CT scans and/or radiology reports |
+| Context length | N/A (not an LLM) |
+| Outputs | Image and/or text embeddings, not directly human-interpretable. |
+| GPUs | 4 x A100 GPUs |
+| Training time | 1.75 Days |
+| Public data summary (or summaries) | Training Data: National Lung Screening Trial (NLST) dataset: Chest CT images from approximately 26k patients CT-RATE dataset: Approximately 25.7k CT acquisitions with associate reports |
+| Training dates | Dates of training: August 2024. Intended model weight release date: Dec 5, 2025 |
+| Status | Static |
+| Release date | Intended release date: Dec 5, 2025 |
+| Release date in the EU (if different) | Intended release date: Dec 5, 2025 |
+| License | Free/open source |
+| Model dependencies | COLIPRI’s text decoder is fine-tuned from microsoft/BiomedVLP-CXR-BERT-specialized · Hugging Face. |
+| List and link to any additional related assets | N/A |
+| Acceptable use policy | N/A |
+
+## Model overview
+
+COLIPRI is an encoder designed for 3D CT scans. Its architecture combines a 3D vision encoder (PRIMUS-M transformer) with a biomedical text encoder (BiomedVLP-CXR-BERT) and enhances existing models by adding additional training objectives and optimizations. These include the inclusion of new training objectives and optimizations to the CLIP paradigm to account for limited data availability. This multi-objective training approach optimizes both global semantic alignment and dense feature learning. After pre-training, the model is evaluated comprehensively on downstream medical imaging tasks such as classification, retrieval, and semantic segmentation.
+COLIPRI is not a generative model. It cannot generate any images or text.
+
+### Alignment approach
+
+N/A
+
+## Usage
+
+### Primary use cases
+
+COLIPRI is primarily intended for 3D medical imaging tasks such as segmentation and classification. It is intended for research only.
+
+### Out-of-scope use cases
+
+The COLIPRI model family was trained on 3D CT volumes and is therefore best suited for inputs of this type. It is intended for research only. The model is for research use only; clinical or deployed use cases are out of scope.
+
+### Distribution channels
+
+The model will be released with open access via Hugging Face.
+
+### Input formats
+
+The vision encoder takes 3D chest CT scans, and the text encoder takes radiology reports.
+
+### Technical requirements and integration guidance
+
+The model requires a Python-based environment with PyTorch. The only hardware requirement is one GPU.
+
+### Responsible AI considerations
+
+The model may underperform or produce unreliable outputs for non-English languages, out-of-distribution or adversarial inputs, or clinical scenarios not well represented in the training data. The model should be used in an AI-assisted research setup with human oversight, rather than as a sole diagnostic tool.
+
+## Data overview
+
+### Training, testing, and validation datasets
+
+The COLIPRI model family was pre-trained on the CT-RATE dataset, consisting of approximately 25.7k CT acquisitions with associated reports, and the National Lung Screening Trial (NLST) dataset, consisting of chest CT images from approximately 26k patients.
+Classification performance was evaluated on a withheld test set of CT-RATE and the publicly available subset of RAD-ChestCT, which comprises 3.6k chest CT volumes with 16 multi-abnormality labels. Semantic segmentation was evaluated by training a five-fold cross validation on four datasets:
+
+- LiTS: Task 3 of the Medical Segmentation Decathlon (MSD), containing segmentations for liver and liver tumors
+- Lung: Task 6 of MSD, containing segmentations of primary lung cancers
+- HVS: Task 8 of the MSD, containing segmentations of hepatic vessels and tumors next to such vessels
+- KiTS23: containing segmentations of tumors, cysts, and the kidney
+
+## Quality and performance evaluation
+
+### Results summary
+
+#### Classification Probes
+
+The COLIPRI family exceeded all third-party models on both the CT-RATE and RAD-ChestCT datasets across all metrics.
+
+#### Zero-Shot Classification
+
+COLIPRI encoders exceed the state of the art using short prompts on the CT-RATE and RAD-ChestCT datasets.
+
+#### Report-to-Image Retrieval
+
+COLIPRI encoders are substantially better in retrieving the associated image given a report, evaluated on the CT-RATE test set.
+
+#### Segmentation
+
+Across four datasets (LiTs, Lung, HVS, KiTS23), COLIPRI-CM and COLIPRI-CRM perform on-par or better than the state of the art.
+
+#### Qualitative Analysis
+
+COLIPRI embeddings generates sharper and more coherent features than third-party models.
+
+### Long context
+
+N/A – not an LLM.
+
+### Safety evaluation and red-teaming
+
+N/A - model is limited to the medical imaging domain. Model outputs were compared to ground-truth images. The standard quality metrics were computed.
+
+## Tracked capability evaluations
+
+N/A This model is not a frontier model.
+
+## Additional information
+
+Requests for additional information may be directed to [MSFTAIActRequest@microsoft.com](mailto:MSFTAIActRequest@microsoft.com).

pyproject.toml

@@ -0,0 +1,70 @@
+[project]
+name = "colipri"
+version = "0.1.0"
+description = "Vision-language encoder for chest CT scans and reports."
+readme = "README.md"
+authors = [
+    { name = "Microsoft Health Futures", email = "innereyedev@microsoft.com" },
+]
+requires-python = ">=3.10"
+dependencies = [
+    "accelerate",
+    "dynamic-network-architectures@git+https://github.com/fepegar/dynamic-network-architectures.git@88ad3c9",
+    "huggingface-hub",
+    "hydra-core",
+    "jaxtyping",
+    "safetensors",
+    "timm<1.0.23",  # RuntimeError related to position embeddings in 1.0.23
+    "torchio>=0.21.1",
+    "tqdm",
+    "transformers",
+    "typer",
+]
+
+[project.urls]
+Homepage = "https://huggingface.co/microsoft/colipri"
+Source = "https://huggingface.co/microsoft/colipri"
+"Issue tracker" = "https://huggingface.co/microsoft/colipri/discussions/new"
+Documentation = "https://huggingface.co/microsoft/colipri/blob/main/README.md"
+
+[project.optional-dependencies]
+demo = [
+    "lovely-tensors",
+    "matplotlib",
+    "scikit-learn",
+    "torchinfo",
+]
+
+[dependency-groups]
+dev = [
+    "ipykernel",
+    "ipywidgets",
+]
+types = ["ty"]
+
+[build-system]
+requires = ["uv_build"]
+build-backend = "uv_build"
+
+[tool.ruff.lint]
+# Defaults from https://docs.astral.sh/ruff/linter/#rule-selection
+select = [
+    # pycodestyle
+    "E",
+    # Pyflakes
+    "F",
+    # pyupgrade
+    "UP",
+    # flake8-bugbear
+    "B",
+    # flake8-simplify
+    "SIM",
+    # isort
+    "I",
+]
+ignore = [
+    "F722", # https://docs.kidger.site/jaxtyping/faq/#flake8-or-ruff-are-throwing-an-error
+]
+
+[tool.ruff.lint.isort]
+force-single-line = true

src/colipri/__init__.py

@@ -0,0 +1,11 @@
+from .model.multimodal import get_model
+from .pipelines import ZeroShotImageClassificationPipeline
+from .processor import get_processor
+from .sample_data import load_sample_ct
+
+__all__ = [
+    "get_model",
+    "get_processor",
+    "load_sample_ct",
+    "ZeroShotImageClassificationPipeline",
+]

src/colipri/checkpoint.py

@@ -0,0 +1,54 @@
+from pathlib import Path
+
+from huggingface_hub import hf_hub_download
+from hydra import compose
+from hydra import initialize_config_dir
+from omegaconf import DictConfig
+from torch import Tensor
+
+from .defaults import REPO_NAME
+from .defaults import REPO_USER
+from .defaults import REVISION
+from .defaults import ROOT_CONFIG_FILENAME
+from .defaults import WEIGHTS_FILENAME
+from .defaults import get_configs_dir
+
+TypeStateDict = dict[str, Tensor]
+
+
+def download_weights(**kwargs) -> Path:
+    weights_path = _download_from_hugging_face(**kwargs)
+    return weights_path
+
+
+def _download_from_hugging_face(
+    repo_user: str = REPO_USER,
+    repo_name: str = REPO_NAME,
+    revision: str | None = REVISION,
+    filename: str = WEIGHTS_FILENAME,
+) -> Path:
+    repo_id = f"{repo_user}/{repo_name}"
+    try:
+        weights_path = hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            revision=revision,
+        )
+    except Exception as e:
+        msg = f'Failed to download "{filename}" from Hugging Face Hub repo "{repo_id}".'
+        raise RuntimeError(msg) from e
+    return Path(weights_path)
+
+
+def _load_config(**kwargs) -> DictConfig:
+    with initialize_config_dir(str(get_configs_dir()), version_base=None):
+        config = compose(ROOT_CONFIG_FILENAME, **kwargs)
+    return config
+
+
+def load_model_config(**kwargs) -> DictConfig:
+    return _load_config(**kwargs)["model"]
+
+
+def load_processor_config(**kwargs) -> DictConfig:
+    return _load_config(**kwargs)["processor"]

src/colipri/configs/config.yaml

@@ -0,0 +1,10 @@
+defaults:
+  - model: default
+  - processor: default
+  - _self_
+
+input_size: 192
+spacing: 2 # mm
+padding_mode: minimum
+image_embed_dim: 864
+projection_dim: 768

src/colipri/configs/model/default.yaml

@@ -0,0 +1,5 @@
+defaults:
+  - image_encoder: default
+  - text_encoder: default
+
+_target_: colipri.model.multimodal.Model

src/colipri/configs/model/image_encoder/backbone/default.yaml

@@ -0,0 +1,20 @@
+_target_: dynamic_network_architectures.architectures.primus.Primus
+input_channels: 1
+num_classes: 1
+eva_depth: 16
+eva_numheads: 12
+embed_dim: ${image_embed_dim}
+patch_embed_size:
+  - 8
+  - 8
+  - 8
+input_shape:
+  - ${input_size}
+  - ${input_size}
+  - ${input_size}
+use_rot_pos_emb: True
+use_abs_pos_embed: False
+drop_path_rate: 0.2
+init_values: 0.1
+scale_attn_inner: True
+num_register_tokens: 0

src/colipri/configs/model/image_encoder/default.yaml

@@ -0,0 +1,6 @@
+defaults:
+  - backbone: default
+  - pooler: default
+  - projector: default
+
+_target_: colipri.model.image.ImageEncoder

src/colipri/configs/model/image_encoder/pooler/default.yaml

@@ -0,0 +1,3 @@
+_target_: colipri.pooling.AttentionPool1D
+embed_dim: ${projection_dim}
+num_heads: 12

src/colipri/configs/model/image_encoder/projector/default.yaml

@@ -0,0 +1,4 @@
+_target_: torch.nn.Conv3d
+in_channels: ${image_embed_dim}
+out_channels: ${projection_dim}
+kernel_size: 1

src/colipri/configs/model/text_encoder/backbone/default.yaml

@@ -0,0 +1,4 @@
+_target_: transformers.AutoModel.from_pretrained
+pretrained_model_name_or_path: microsoft/BiomedVLP-CXR-BERT-specialized
+trust_remote_code: yes
+revision: 5157bdba1437a3aed316dacb1a5b68edf96b9902

src/colipri/configs/model/text_encoder/default.yaml

@@ -0,0 +1,5 @@
+defaults:
+  - backbone: default
+  - pooler: default
+
+_target_: colipri.model.text.TextEncoder

src/colipri/configs/model/text_encoder/pooler/default.yaml

@@ -0,0 +1,3 @@
+_target_: colipri.pooling.AttentionPool1D
+embed_dim: ${projection_dim}
+num_heads: 12

src/colipri/configs/processor/default.yaml

@@ -0,0 +1,5 @@
+defaults:
+  - image_transform: default
+  - tokenizer: default
+
+_target_: colipri.processor.Processor

src/colipri/configs/processor/image_transform/default.yaml

@@ -0,0 +1,49 @@
+# Exported automatically from nnssl: https://github.com/microsoft/nnssl/pull/41, then
+# adapted for Colipri needs.
+_target_: torchio.transforms.augmentation.composition.Compose
+transforms:
+  - _target_: torchio.transforms.preprocessing.spatial.to_orientation.ToOrientation
+    copy: true
+    exclude: null
+    include: null
+    orientation: SAR
+  - _target_: torchio.transforms.preprocessing.spatial.resample.Resample
+    antialias: false
+    copy: true
+    exclude: null
+    image_interpolation: linear
+    include: null
+    label_interpolation: nearest
+    pre_affine_name: null
+    scalars_only: false
+    target: ${spacing}
+  - _target_: torchio.transforms.preprocessing.intensity.clamp.Clamp
+    out_min: -1000
+    out_max: 1000
+  - _target_: torchio.transforms.preprocessing.intensity.rescale.RescaleIntensity
+    copy: true
+    exclude: null
+    in_min_max:
+      - -1000
+      - 1000
+    include: null
+    masking_method: null
+    out_min_max:
+      - -1
+      - 1
+    percentiles:
+      - 0
+      - 100
+  - _target_: torchio.transforms.preprocessing.spatial.crop_or_pad.CropOrPad
+    copy: true
+    exclude: null
+    include: null
+    labels: null
+    mask_name: null
+    only_crop: false
+    only_pad: false
+    padding_mode: ${padding_mode}
+    target_shape:
+      - ${input_size}
+      - ${input_size}
+      - ${input_size}

src/colipri/configs/processor/tokenizer/default.yaml

@@ -0,0 +1,5 @@
+_target_: transformers.AutoTokenizer.from_pretrained
+pretrained_model_name_or_path: microsoft/BiomedVLP-CXR-BERT-specialized
+trust_remote_code: yes
+revision: 5157bdba1437a3aed316dacb1a5b68edf96b9902
+model_max_length: 512

src/colipri/defaults.py

@@ -0,0 +1,21 @@
+import importlib.resources
+from pathlib import Path
+
+package_name = __package__
+assert isinstance(package_name, str)
+
+REPO_USER = "microsoft"
+REPO_NAME = package_name
+REVISION = None
+WEIGHTS_FILENAME = "model.safetensors"
+ROOT_CONFIG_FILENAME = "config.yaml"
+
+
+def get_configs_dir(pkg: str | None = None) -> Path:
+    package_dir = importlib.resources.files(package_name)
+    return Path(package_dir / "configs")
+
+
+__all__ = [
+    "get_configs_dir",
+]

src/colipri/model/base.py

@@ -0,0 +1,8 @@
+import torch
+
+
+class ModelMixin:
+    @property
+    def device(self) -> torch.device:
+        one_parameter = next(self.parameters())
+        return one_parameter.device

src/colipri/model/image.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+import torchio as tio
+from dynamic_network_architectures.architectures.primus import Primus
+from einops import rearrange
+from torch import nn
+from tqdm.auto import tqdm
+
+from ..pooling import AttentionPool1D
+from ..types import TypeImageEmbeddings
+from ..types import TypeImagesTensor
+from ..types import TypeIntOrTripletInt
+from .base import ModelMixin
+
+
+class ImageEncoder(nn.Module, ModelMixin):
+    def __init__(
+        self,
+        backbone: Primus,
+        projector: nn.Conv3d,
+        pooler: AttentionPool1D,
+    ):
+        super().__init__()
+        self.backbone = self._remove_decoder(backbone)
+        self.projector = projector
+        self.pooler = pooler
+
+    def _remove_decoder(self, backbone: Primus) -> Primus:
+        if hasattr(backbone, "up_projection"):
+            backbone.up_projection = nn.Identity()  # type: ignore
+        return backbone
+
+    @property
+    def patch_size(self) -> tuple[int, int, int]:
+        patch_embedder: nn.Conv3d = self.backbone.down_projection.proj  # type: ignore[reportAssignmentType]
+        patch_size: tuple[int, int, int] = patch_embedder.stride  # type: ignore[reportAssignmentType]
+        return patch_size
+
+    def encode(
+        self,
+        images: TypeImagesTensor,
+    ) -> TypeImageEmbeddings:
+        images = images.to(self.device)
+        embeddings: TypeImageEmbeddings = self.backbone(images)
+        return embeddings
+
+    def encode_sliding_window(
+        self,
+        images: TypeImagesTensor,
+        window_size: TypeIntOrTripletInt,
+        overlap: int = 0,
+    ) -> TypeImageEmbeddings:
+        if len(set(self.patch_size)) > 1:
+            msg = (
+                "Sliding window encoding is only supported for models with cubic"
+                " patch sizes for now."
+            )
+            raise NotImplementedError(msg)
+        else:
+            patch_size = self.patch_size[0]
+        image_key = "image"  # could be anything
+        embeddings = []
+        for image in images:
+            grid_sampler = tio.inference.GridSampler(
+                tio.Subject(**{image_key: tio.ScalarImage(tensor=image)}),
+                window_size,
+                overlap,
+            )
+            patch_loader = tio.SubjectsLoader(grid_sampler)  # type: ignore[reportArgumentType]
+            aggregator = tio.data.GridAggregator(
+                grid_sampler,
+                downsampling_factor=patch_size,
+            )
+            for patches_batch in tqdm(patch_loader):
+                input_tensor = patches_batch[image_key][tio.DATA].to(self.device)
+                locations = patches_batch[tio.LOCATION]
+                outputs = self.backbone(input_tensor)
+                aggregator.add_batch(outputs, locations)
+            embeddings.append(aggregator.get_output_tensor())
+        return torch.stack(embeddings).to(self.device)
+
+    def forward(
+        self,
+        images: TypeImagesTensor,
+        *,
+        project: bool,
+        pool: bool,
+        normalize: bool,
+        window_size: TypeIntOrTripletInt | None = None,
+    ) -> TypeImageEmbeddings:
+        if pool and not project:
+            msg = "Pooling requires projection to be enabled. Set project=True."
+            raise NotImplementedError(msg)
+        if window_size is None:
+            embeddings = self.encode(images)
+        else:
+            embeddings = self.encode_sliding_window(images, window_size)
+        if project:
+            embeddings = self.projector(embeddings)
+            if pool:
+                sequence = rearrange(embeddings, "b c x y z -> b (x y z) c")
+                embeddings = self.pooler(sequence)
+            else:
+                embeddings = self.pooler.to_dense()(embeddings)
+        if normalize:
+            embeddings = F.normalize(embeddings, dim=1)
+        return embeddings

src/colipri/model/multimodal.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import torch
+from accelerate import init_empty_weights
+from accelerate import load_checkpoint_and_dispatch
+from einops import rearrange
+from hydra.utils import instantiate
+from safetensors.torch import load_model
+from safetensors.torch import save_model
+from torch import nn
+from transformers.utils.logging import get_verbosity
+from transformers.utils.logging import set_verbosity
+from transformers.utils.logging import set_verbosity_error
+
+from ..checkpoint import download_weights
+from ..checkpoint import load_model_config
+from ..types import TypeImageEmbeddings
+from ..types import TypeImagesTensor
+from ..types import TypeIntOrTripletInt
+from ..types import TypePatchEmbeddings
+from ..types import TypePatchLogits
+from ..types import TypePatchProbabilities
+from ..types import TypePath
+from ..types import TypePooledEmbeddings
+from ..types import TypePooledLogits
+from ..types import TypePooledProbabilities
+from ..types import TypeTextAttentionMask
+from ..types import TypeTextEmbeddings
+from ..types import TypeTextTokenIds
+from .base import ModelMixin
+from .image import ImageEncoder
+from .text import TextEncoder
+
+
+def get_model(
+    checkpoint_path: TypePath | None = None,
+    *,
+    pretrained: bool = True,
+    image_only: bool = False,
+    **kwargs,
+) -> Model:
+    if pretrained and checkpoint_path is None:
+        checkpoint_path = download_weights()
+
+    overrides = []
+    for key, value in kwargs.items():
+        overrides.append(f"{key}={value}")
+    config = load_model_config(overrides=overrides)
+
+    if image_only:
+        config.text_encoder = None
+
+    transformers_verbosity = get_verbosity()
+    set_verbosity_error()
+    if checkpoint_path is None:
+        model = instantiate(config)
+    else:
+        with init_empty_weights():
+            model = instantiate(config)
+        accelerate_logger = logging.getLogger("accelerate.utils.modeling")
+        old_level = accelerate_logger.getEffectiveLevel()
+        accelerate_logger.setLevel(logging.ERROR)
+        model = load_checkpoint_and_dispatch(model, checkpoint_path)
+        accelerate_logger.setLevel(old_level)
+    set_verbosity(transformers_verbosity)
+
+    assert isinstance(model, Model)
+    return model.eval()
+
+
+class Model(nn.Module, ModelMixin):
+    def __init__(
+        self,
+        image_encoder: ImageEncoder,
+        text_encoder: TextEncoder,
+        temperature: float = 1,
+    ):
+        super().__init__()
+        self.image_encoder = image_encoder
+        self.text_encoder = text_encoder
+        self.register_buffer("softmax_temperature", torch.tensor(temperature))
+
+    @property
+    def patch_size(self) -> tuple[int, int, int]:
+        return self.image_encoder.patch_size
+
+    def encode_image(
+        self,
+        images: TypeImagesTensor,
+        *,
+        pool: bool = False,
+        project: bool = False,
+        normalize: bool = False,
+        window_size: TypeIntOrTripletInt | None = None,
+    ) -> TypeImageEmbeddings:
+        return self.image_encoder(
+            images,
+            project=project,
+            pool=pool,
+            normalize=normalize,
+            window_size=window_size,
+        )
+
+    def encode_text(
+        self,
+        token_ids: TypeTextTokenIds,
+        attention_mask: TypeTextAttentionMask | None = None,
+        *,
+        pool: bool = True,
+        normalize: bool = True,
+    ) -> TypeTextEmbeddings:
+        return self.text_encoder(
+            token_ids,
+            attention_mask,
+            pool=pool,
+            normalize=normalize,
+        )
+
+    def compute_similarities(
+        self,
+        image_embeddings: TypePooledEmbeddings | TypePatchEmbeddings,
+        text_embeddings: TypePooledEmbeddings,
+    ) -> TypePatchLogits | TypePooledLogits:
+        text_embeddings = rearrange(text_embeddings, "num_prompts c -> c num_prompts")
+        is_grid = image_embeddings.ndim == 5
+        if is_grid:
+            num_images, _, x, y, z = image_embeddings.shape
+            image_embeddings = rearrange(
+                image_embeddings,
+                "num_images c x y z -> (num_images x y z) c",
+            )
+            similarities_flat = image_embeddings @ text_embeddings
+            similarities = rearrange(
+                similarities_flat,
+                "(num_images x y z) num_prompts -> num_images num_prompts x y z",
+                num_images=num_images,
+                x=x,
+                y=y,
+                z=z,
+            )
+        else:
+            similarities = image_embeddings @ text_embeddings
+        return similarities
+
+    def classify(
+        self,
+        images: TypePooledEmbeddings | TypePatchEmbeddings,
+        text: TypePooledEmbeddings,
+    ) -> TypePooledProbabilities | TypePatchProbabilities:
+        logits = self.compute_similarities(images, text)
+        assert isinstance(self.softmax_temperature, torch.Tensor)
+        probabilities = (logits / self.softmax_temperature).softmax(dim=-1)
+        return probabilities
+
+    def save_weights(self, path: TypePath) -> None:
+        path = Path(path)
+        if path.suffix == ".safetensors":
+            save_model(self, str(path))
+        else:
+            weights = self.state_dict()
+            torch.save(weights, path)
+
+    def load_weights(
+        self,
+        path: TypePath,
+        device: torch.device | str | None = None,
+    ) -> None:
+        path = Path(path)
+        if not path.exists():
+            raise FileNotFoundError(f"Checkpoint file {path} does not exist.")
+        if device is None:
+            device = torch.device("cpu")
+        if path.suffix == ".safetensors":
+            if device is not None:
+                device = str(device)
+            missing, unexpected = load_model(self, path, device=device)
+            if missing or unexpected:
+                raise RuntimeError("TODO")
+        else:
+            weights = torch.load(path, map_location=device)
+            self.load_state_dict(weights, strict=False)

src/colipri/model/text.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import torch.nn.functional as F
+from torch import nn
+from transformers import BertModel
+
+from ..types import TypeTextAttentionMask
+from ..types import TypeTextEmbeddings
+from ..types import TypeTextTokenIds
+from .base import ModelMixin
+
+
+class TextEncoder(nn.Module, ModelMixin):
+    def __init__(self, backbone: nn.Module, pooler: nn.Module):
+        super().__init__()
+        if hasattr(backbone, "bert"):
+            bert = backbone.bert
+            assert isinstance(bert, BertModel)
+            bert.pooler = None  # we use our own pooler
+            backbone = bert
+        self.backbone = backbone
+        self.pooler = pooler
+
+    def forward(
+        self,
+        token_ids: TypeTextTokenIds,
+        attention_mask: TypeTextAttentionMask | None = None,
+        *,
+        pool: bool = True,
+        normalize: bool = True,
+    ) -> TypeTextEmbeddings:
+        token_ids = token_ids.to(self.device)
+        if attention_mask is not None:
+            attention_mask = attention_mask.to(self.device)
+        # We didn't use a projector during pre-training
+        output = self.backbone(token_ids, attention_mask=attention_mask)
+        embeddings = output["last_hidden_state"]
+        if pool:
+            embeddings = self.pooler(embeddings)
+        if normalize:
+            embeddings = F.normalize(embeddings, dim=1)
+        return embeddings

src/colipri/pipelines.py

@@ -0,0 +1,65 @@
+import torch
+
+from .model.multimodal import Model
+from .processor import Processor
+from .types import TypeImage
+from .types import TypePath
+from .types import TypeScores
+
+
+class ZeroShotImageClassificationPipeline:
+    def __init__(
+        self,
+        model: Model,
+        processor: Processor,
+        batch_size: int = 1,  # TODO
+        num_workers: int = 0,  # TODO
+    ):
+        self._model = model.eval()
+        self._processor = processor
+
+    @property
+    def device(self) -> torch.device:
+        return self._model.device
+
+    # TODO: add support for multiple prompts per class
+    @torch.no_grad()
+    def __call__(
+        self,
+        images: TypePath | TypeImage | list[TypePath | TypeImage],
+        prompts: str | list[str],
+    ) -> TypeScores | list[TypeScores]:
+        preprocessed_images = self._processor.process_images(images)
+        images_batch = self._processor.to_images_batch(preprocessed_images)
+
+        text_token_ids, text_attention_mask = self._processor.process_text(prompts)
+
+        image_embeddings_batch = self._model.encode_image(
+            images_batch,
+            project=True,
+            pool=True,
+            normalize=True,
+        )
+        text_embeddings_batch = self._model.encode_text(
+            text_token_ids,
+            text_attention_mask,
+        )
+
+        probabilities = self._model.classify(
+            image_embeddings_batch,
+            text_embeddings_batch,
+        )
+
+        if not isinstance(prompts, list):
+            prompts = [prompts]
+
+        all_results = []
+        for image_probabilities in probabilities:
+            image_results = []
+            for prompt, score in zip(prompts, image_probabilities, strict=True):
+                image_results.append({"score": score.item(), "label": prompt})
+            all_results.append(image_results)
+
+        if len(images_batch) == 1 and not isinstance(images, list):
+            all_results = all_results[0]
+        return all_results

src/colipri/pooling.py

@@ -0,0 +1,74 @@
+import torch
+from einops import rearrange
+from torch import nn
+
+from .types import TypePooledEmbeddings
+from .types import TypeSequenceEmbeddings
+
+
+class AttentionPool1D(nn.Module):
+    def __init__(self, embed_dim: int, num_heads: int):
+        super().__init__()
+        self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
+
+    def forward(self, x: TypeSequenceEmbeddings) -> TypePooledEmbeddings:
+        query = x.mean(dim=1, keepdim=True)
+        key = value = x
+        pooled, _ = self.attn(query, key, value)
+        return rearrange(pooled, "batch 1 embed_dim -> batch embed_dim")
+
+    def to_dense(self):
+        v_proj_in_weight_qkv = self.get_parameter("attn.in_proj_weight")
+        v_proj_in_bias_qkv = self.get_parameter("attn.in_proj_bias")
+        v_proj_out_weight = self.get_parameter("attn.out_proj.weight")
+        v_proj_out_bias = self.get_parameter("attn.out_proj.bias")
+        dim = v_proj_in_weight_qkv.shape[0] // 3
+        v_proj_in_weight_v = v_proj_in_weight_qkv[2 * dim :]
+        v_proj_in_bias_v = v_proj_in_bias_qkv[2 * dim :]
+
+        value_projection = nn.Conv3d(
+            in_channels=dim,
+            out_channels=dim,
+            kernel_size=1,
+        )
+        value_projection.weight.data = rearrange(
+            v_proj_in_weight_v,
+            "c_out c_in -> c_out c_in 1 1 1",
+        )
+        assert value_projection.bias is not None
+        value_projection.bias.data = v_proj_in_bias_v
+
+        out_projection = nn.Conv3d(
+            in_channels=dim,
+            out_channels=dim,
+            kernel_size=1,
+        )
+        out_projection.weight.data = rearrange(
+            v_proj_out_weight,
+            "c_out c_in -> c_out c_in 1 1 1",
+        )
+        assert out_projection.bias is not None
+        out_projection.bias.data = v_proj_out_bias
+
+        return nn.Sequential(
+            value_projection,
+            out_projection,
+        )
+
+
+class MultiLearnedQueryAttentionPool1D(AttentionPool1D):
+    def __init__(self, embed_dim: int, num_heads: int):
+        super().__init__(embed_dim, num_heads)
+        # 4 Queries instead of 1
+        self.query = nn.Parameter(torch.randn(1, 4, embed_dim) / embed_dim**0.5)
+
+    def forward(self, x: TypeSequenceEmbeddings) -> TypePooledEmbeddings:
+        """
+        x: [B, T, D] — sequence of token embeddings
+        returns: [B, D] — pooled representation
+        """
+        B, T, D = x.shape
+        query = self.query.expand(B, -1, -1)  # [B, 4, D]
+        pooled, _ = self.attn(query, x, x)  # [B, 4, D]
+        # pooled: [4, B, D_out], want [B, D_out] by pooling over queries (mean)
+        return pooled.mean(dim=1)  # [B, D]

src/colipri/processor.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import torch
+import torchio as tio
+from hydra.utils import instantiate
+from transformers import BertTokenizer
+from transformers.tokenization_utils_base import BatchEncoding
+
+from .checkpoint import load_processor_config
+from .types import TypeImage
+from .types import TypeImages
+from .types import TypeImagesTensor
+from .types import TypePath
+from .types import TypeStringOrStrings
+from .types import TypeTextAttentionMask
+from .types import TypeTextTokenIds
+
+
+def get_processor(*, image_only: bool = False, **kwargs) -> Processor:
+    overrides = []
+    for key, value in kwargs.items():
+        overrides.append(f"{key}={value}")
+    config = load_processor_config(overrides=overrides)
+    if image_only:
+        config.tokenizer = None
+    return instantiate(config)
+
+
+class Processor:
+    def __init__(
+        self,
+        image_transform: tio.Transform,
+        tokenizer: BertTokenizer,
+    ):
+        self._image_transform = image_transform
+        self._text_tokenizer = tokenizer
+
+    def __repr__(self) -> str:
+        lines = [
+            f"{self.__class__.__name__}(",
+            f"  image_transform={self._image_transform},",
+            f"  text_tokenizer={self._text_tokenizer},",
+            ")",
+        ]
+        return "\n".join(lines)
+
+    def process_images(
+        self,
+        inputs: TypePath | TypeImage | list[TypePath | TypeImage],
+    ) -> TypeImages:
+        if not isinstance(inputs, list):
+            inputs = [inputs]
+        images = []
+        for image_or_path in inputs:
+            is_image = isinstance(image_or_path, tio.ScalarImage)
+            if is_image:
+                image = image_or_path
+            else:
+                path = image_or_path
+                image = tio.ScalarImage(path)
+            images.append(image)
+        images = [self._image_transform(image) for image in images]
+        return images
+
+    def to_images_batch(
+        self,
+        images: TypeImages,
+    ) -> TypeImagesTensor:
+        if not isinstance(images, list):
+            msg = f"Expected images to be a list, got {type(images)}"
+            raise TypeError(msg)
+        images_tensor = torch.stack([image.data for image in images])
+        return images_tensor
+
+    def process_text(
+        self,
+        text: TypeStringOrStrings,
+        **kwargs,
+    ) -> tuple[TypeTextTokenIds, TypeTextAttentionMask]:
+        encodings: BatchEncoding = self._text_tokenizer(
+            text,
+            max_length=512,  # TODO
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt",
+            **kwargs,
+        )
+        token_ids: TypeTextTokenIds = encodings["input_ids"]  # type: ignore[reportAssignmentType]
+        attention_mask: TypeTextAttentionMask = encodings["attention_mask"]  # type: ignore[reportAssignmentType]
+        return token_ids, attention_mask

src/colipri/sample_data.py

@@ -0,0 +1,24 @@
+from pathlib import Path
+
+import torchio as tio
+
+
+def load_sample_ct(clip: bool = True) -> tio.ScalarImage:
+    """Get a sample CT image."""
+    ct = get_sample_ct()
+    if clip:
+        clamp = tio.Clamp(-1000, 1000)
+        ct = clamp(ct)
+    else:
+        ct.load()
+    return ct
+
+
+def get_sample_ct() -> tio.ScalarImage:
+    """Download a sample CT image if not already present."""
+    return tio.datasets.Slicer("CTChest").CT_chest  # type: ignore[reportAttributeAccessIssue]
+
+
+def get_sample_ct_path() -> Path:
+    """Get the path to a sample CT image."""
+    return get_sample_ct().path  # type: ignore[reportReturnType]

src/colipri/types.py

@@ -0,0 +1,42 @@
+import os
+from typing import TypeAlias
+
+import torchio as tio
+from jaxtyping import Float
+from jaxtyping import Int64
+from torch import Tensor
+
+TypePath: TypeAlias = str | os.PathLike[str]
+TypePaths: TypeAlias = list[TypePath]
+
+# Raw inputs
+TypeImage: TypeAlias = tio.ScalarImage
+TypeImages: TypeAlias = list[tio.ScalarImage]
+TypeImageOrImages: TypeAlias = TypeImage | TypeImages
+
+TypeString: TypeAlias = str
+TypeStrings: TypeAlias = list[str]
+TypeStringOrStrings: TypeAlias = TypeString | TypeStrings
+
+TypeIntOrTripletInt: TypeAlias = int | tuple[int, int, int]
+
+# Processed inputs
+TypeImagesTensor: TypeAlias = Float[Tensor, "batch 1 x_in y_in z_in"]
+
+TypeTextTokenIds: TypeAlias = Int64[Tensor, "batch text_length"]
+TypeTextAttentionMask: TypeAlias = Int64[Tensor, "batch text_length"]
+
+# Outputs
+TypeSequenceEmbeddings: TypeAlias = Float[Tensor, "batch sequence_length embed_dim"]
+TypePooledEmbeddings: TypeAlias = Float[Tensor, "batch embed_dim"]
+
+TypePatchEmbeddings: TypeAlias = Float[Tensor, "batch embed_dim x_out y_out z_out"]
+TypeImageEmbeddings: TypeAlias = TypePatchEmbeddings | TypePooledEmbeddings
+
+TypeTextEmbeddings: TypeAlias = TypeSequenceEmbeddings | TypePooledEmbeddings
+
+TypePatchLogits: TypeAlias = Float[Tensor, "num_images num_prompts x_out y_out z_out"]
+TypePooledLogits: TypeAlias = Float[Tensor, "num_images num_prompts"]
+TypePatchProbabilities: TypeAlias = TypePatchLogits
+TypePooledProbabilities: TypeAlias = TypePooledLogits
+TypeScores: TypeAlias = list[dict[str, float | str]]