Duplicate from microsoft/rad-dino

Co-authored-by: Fernando Pérez-García <fepegar@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl 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
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+training_images.csv filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,17 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+
+.vscode/
+
+*.ipynb
+uv.lock
+*.txt
+.python-version

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

@@ -0,0 +1,219 @@
+---
+license_name: mit
+pipeline_tag: image-feature-extraction
+---
+
+# Model card for RAD-DINO
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+RAD-DINO is a vision transformer model trained to encode chest X-rays using the self-supervised learning method [DINOv2](https://openreview.net/forum?id=a68SUt6zFt).
+
+## Model description
+
+<!-- Provide a longer summary of what this model is. -->
+
+RAD-DINO is described in detail in [Exploring Scalable Medical Image Encoders Beyond Text Supervision (F. Pérez-García, H. Sharma, S. Bond-Taylor, et al., 2025)](https://www.nature.com/articles/s42256-024-00965-w).
+
+- **Developed by:** Microsoft Health Futures
+- **Model type:** Vision transformer
+- **License:** [MIT](./LICENSE)
+- **Finetuned from model:** [`dinov2-base`](https://huggingface.co/facebook/dinov2-base)
+
+## 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. -->
+
+RAD-DINO 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 model is a vision backbone that can be plugged to other models for downstream tasks.
+Some potential uses are:
+
+- Image classification, with a classifier trained on top of the `CLS` token
+- Image segmentation, with a decoder trained using the patch tokens
+- Clustering, using the image embeddings directly
+- Image retrieval, using nearest neighbors of the CLS token
+- Report generation, with a language model to decode text
+
+Fine-tuning RAD-DINO 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. -->
+
+RAD-DINO was trained with data from three countries, 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 rad-dino
+```
+
+## Usage
+
+### Encode an image
+
+```python
+>>> from rad_dino import RadDino
+>>> from rad_dino.utils import download_sample_image
+>>> encoder = RadDino()
+>>> image = download_sample_image()
+>>> image
+<PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=2765x2505 at 0x7CCD5C014050>
+>>> cls_embeddings, patch_embeddings = encoder.extract_features(image)
+>>> cls_embeddings.shape, patch_embeddings.shape
+(torch.Size([1, 768]), torch.Size([1, 768, 37, 37]))
+```
+
+### Weights for fine-tuning
+
+We have released a checkpoint compatible with [the original DINOv2 code](https://github.com/facebookresearch/dinov2) to help researchers fine-tune our model.
+
+We can use the hub model and load the RAD-DINO weights.
+Let's clone the DINOv2 repository so we can import the code for the head.
+
+```shell
+git clone https://github.com/facebookresearch/dinov2.git
+```
+
+```python
+>>> import torch
+>>> from rad_dino.utils import safetensors_to_state_dict
+>>> rad_dino_gh = torch.hub.load("./dinov2", "dinov2_vitb14")
+>>> backbone_state_dict = safetensors_to_state_dict("backbone_compatible.safetensors")
+>>> rad_dino_gh.load_state_dict(backbone_state_dict, strict=True)
+<All keys matched successfully>
+```
+
+The weights of the head are also released:
+
+```python
+>>> from dinov2.layers import DINOHead
+>>> rad_dino_head_gh = DINOHead(
+...    in_dim=768,
+...    out_dim=65536,
+...    hidden_dim=2048,
+...    bottleneck_dim=256,
+...    nlayers=3,
+... )
+>>> head_state_dict = safetensors_to_state_dict("dino_head.safetensors")
+>>> rad_dino_head_gh.load_state_dict(head_state_dict, strict=True)
+<All keys matched successfully>
+```
+
+### Configs and augmentation
+
+The configuration files [`ssl_default_config.yaml`](./ssl_default_config.yaml) and [`vitb14_cxr.yaml`](./vitb14_cxr.yaml), and the [`augmentations`](./augmentations.py) module are also available in the repository to help researchers reproduce the training procedure with our hyperparameters.
+
+## Training details
+
+### Training data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+We used images from five public, deidentified chest X-ray datasets to train this checkpoint of RAD-DINO.
+
+| Dataset   | Num. images |
+| --------- | ----------: |
+| [MIMIC-CXR](https://www.nature.com/articles/s41597-019-0322-0) | 368 960 |
+| [CheXpert](https://ojs.aaai.org/index.php/AAAI/article/view/3834) | 223 648 |
+| [NIH-CXR](https://openaccess.thecvf.com/content_cvpr_2017/html/Wang_ChestX-ray8_Hospital-Scale_Chest_CVPR_2017_paper.html) | 112 120 |
+| [PadChest](https://www.sciencedirect.com/science/article/abs/pii/S1361841520301614) | 136 787 |
+| [BRAX](https://www.nature.com/articles/s41597-022-01608-8) | 41 260 |
+| **TOTAL** | 882 775 |
+
+Images in the validation and test sets used to train [MAIRA](https://arxiv.org/abs/2311.13668) were excluded from the training set of RAD-DINO.
+The list of image files used for training is available at [`./training_images.csv`](./training_images.csv).
+
+Note this checkpoint is different from the one in the paper, where some private data was used (and fewer GPUs).
+The checkpoint shared here is trained for 35 000 iterations (the total number of iterations in the run was 100 000, but we selected this checkpoint using linear probing on the validation sets of the evaluation datasets described in the paper).
+We used 16 nodes with 4 A100 GPUs each, and a batch size of 40 images per GPU.
+
+### Training procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+We refer to the [manuscript](https://www.nature.com/articles/s42256-024-00965-w) for a detailed description of the training procedure.
+
+#### Preprocessing
+
+All DICOM files were resized using B-spline interpolation so that their shorter size was 518, min-max scaled to [0, 255], and stored as PNG files.
+
+#### Training hyperparameters
+
+- **Training regime:** fp16 using PyTorch-FSDP mixed-precision.
+
+<!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+Our evaluation is best described in the [manuscript](https://www.nature.com/articles/s42256-024-00965-w).
+
+## 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: A100 PCIe -->
+<!-- Hours: 1d 16h = 40h -->
+<!-- Cloud provider: Azure -->
+<!-- Region: Italy North -->
+
+- **Hardware type:** NVIDIA A100 GPUs
+- **Hours used:** 40 hours/GPU × 16 nodes × 4 GPUs/node = 2560 GPU-hours
+- **Cloud provider:** Azure
+- **Compute region:** West US 2
+- **Carbon emitted:** 222 kg CO₂ eq.
+
+### Compute infrastructure
+
+RAD-DINO was trained on [Azure Machine Learning](https://azure.microsoft.com/en-us/products/machine-learning).
+
+#### Hardware
+
+We used 16 `Standard_NC96ads_A100_v4` nodes with four NVIDIA A100 (80 GB) GPUs each.
+
+#### Software
+
+We leveraged the code in [DINOv2](https://openreview.net/forum?id=a68SUt6zFt) for training.
+We used [SimpleITK](https://simpleitk.org/) and [Pydicom](https://pydicom.github.io/) for processing of DICOM files.
+
+## 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
+@article{perez-garcia_exploring_2025,
+ title = {Exploring scalable medical image encoders beyond text supervision},
+ issn = {2522-5839},
+ url = {https://doi.org/10.1038/s42256-024-00965-w},
+ doi = {10.1038/s42256-024-00965-w},
+ journal = {Nature Machine Intelligence},
+ author = {P{\'e}rez-Garc{\'i}a, Fernando and Sharma, Harshita and Bond-Taylor, Sam and Bouzid, Kenza and Salvatelli, Valentina and Ilse, Maximilian and Bannur, Shruthi and Castro, Daniel C. and Schwaighofer, Anton and Lungren, Matthew P. and Wetscherek, Maria Teodora and Codella, Noel and Hyland, Stephanie L. and Alvarez-Valle, Javier and Oktay, Ozan},
+ month = jan,
+ year = {2025},
+}
+```
+
+**APA:**
+
+> Pérez-García, F., Sharma, H., Bond-Taylor, S., Bouzid, K., Salvatelli, V., Ilse, M., Bannur, S., Castro, D. C., Schwaighofer, A., Lungren, M. P., Wetscherek, M. T., Codella, N., Hyland, S. L., Alvarez-Valle, J., & Oktay, O. (2025). *Exploring scalable medical image encoders beyond text supervision*. In Nature Machine Intelligence. Springer Science and Business Media LLC. <https://doi.org/10.1038/s42256-024-00965-w>
+
+## Model card contact
+
+Fernando Pérez-García ([`fperezgarcia@microsoft.com`](mailto:fperezgarcia@microsoft.com)).

augmentations.py

@@ -0,0 +1,147 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# See LICENSE in the repo root for license information.
+#
+# Portions:
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+import logging
+
+from PIL import Image
+from torchvision import transforms
+
+from .transforms import (
+    GaussianBlur,
+    MaybeToTensor,
+    make_normalize_transform,
+)
+
+
+logger = logging.getLogger("dinov2")
+
+
+class DataAugmentationDINO(object):
+    def __init__(
+        self,
+        global_crops_scale,
+        local_crops_scale,
+        local_crops_number,
+        global_crops_size=224,
+        local_crops_size=96,
+    ):
+        self.global_crops_scale = global_crops_scale
+        self.local_crops_scale = local_crops_scale
+        self.local_crops_number = local_crops_number
+        self.global_crops_size = global_crops_size
+        self.local_crops_size = local_crops_size
+
+        logger.info("###################################")
+        logger.info("Using data augmentation parameters:")
+        logger.info(f"global_crops_scale: {global_crops_scale}")
+        logger.info(f"local_crops_scale: {local_crops_scale}")
+        logger.info(f"local_crops_number: {local_crops_number}")
+        logger.info(f"global_crops_size: {global_crops_size}")
+        logger.info(f"local_crops_size: {local_crops_size}")
+        logger.info("###################################")
+
+        # random resized crop and flip
+        self.geometric_augmentation_global = transforms.Compose(
+            [
+                transforms.RandomResizedCrop(
+                    global_crops_size, scale=global_crops_scale, interpolation=transforms.InterpolationMode.BICUBIC
+                ),
+                transforms.RandomHorizontalFlip(p=0.5),
+            ]
+        )
+
+        self.geometric_augmentation_local = transforms.Compose(
+            [
+                transforms.RandomResizedCrop(
+                    local_crops_size, scale=local_crops_scale, interpolation=transforms.InterpolationMode.BICUBIC
+                ),
+                transforms.RandomHorizontalFlip(p=0.5),
+            ]
+        )
+
+        # color distorsions / blurring
+        color_jittering = transforms.Compose(
+            [
+                transforms.RandomApply(
+                    [transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1)],
+                    p=0.8,
+                ),
+                transforms.RandomGrayscale(p=0.2),
+            ]
+        )
+
+        global_transfo1_extra = GaussianBlur(p=0.5)
+
+        global_transfo2_extra = transforms.Compose(
+            [
+                GaussianBlur(p=0.1),
+            ]
+        )
+
+        local_transfo_extra = GaussianBlur(p=0.5)
+
+        # normalization
+        self.normalize = transforms.Compose(
+            [
+                MaybeToTensor(),
+                make_normalize_transform(),
+            ]
+        )
+
+        self.global_transfo1 = transforms.Compose([color_jittering, global_transfo1_extra, self.normalize])
+        self.global_transfo2 = transforms.Compose([color_jittering, global_transfo2_extra, self.normalize])
+        self.local_transfo = transforms.Compose([color_jittering, local_transfo_extra, self.normalize])
+
+    def __call__(self, image):
+        output = {}
+
+        # global crops:
+        im1_base = self.geometric_augmentation_global(image)
+        global_crop_1 = self.global_transfo1(im1_base)
+
+        im2_base = self.geometric_augmentation_global(image)
+        global_crop_2 = self.global_transfo2(im2_base)
+
+        output["global_crops"] = [global_crop_1, global_crop_2]
+
+        # global crops for teacher:
+        output["global_crops_teacher"] = [global_crop_1, global_crop_2]
+
+        # local crops:
+        local_crops = [
+            self.local_transfo(self.geometric_augmentation_local(image)) for _ in range(self.local_crops_number)
+        ]
+        output["local_crops"] = local_crops
+        output["offsets"] = ()
+
+        return output
+
+
+def get_online_classification_augmentation_from_config(cfg) -> transforms.Compose:
+    augmentation_config = cfg.evaluation.online.augmentation
+    interpolation = getattr(Image.Resampling, augmentation_config.interpolation)
+    resize_size = crop_size = cfg.crops.global_crops_size
+    resize = transforms.Resize(resize_size, interpolation=interpolation)
+    crop = transforms.CenterCrop(crop_size)
+    affine = transforms.RandomAffine(
+        degrees=augmentation_config.degrees,
+        scale=augmentation_config.scale,
+        shear=augmentation_config.shear,
+        interpolation=interpolation,
+    )
+    transforms_list = [
+        resize,
+        crop,
+        affine,
+        MaybeToTensor(),
+        make_normalize_transform(),
+    ]
+    if augmentation_config.horizontal_flip:
+        transforms_list.append(transforms.RandomHorizontalFlip())
+    return transforms.Compose(transforms_list)

backbone_compatible.safetensors

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

config.json

@@ -0,0 +1,47 @@
+{
+  "apply_layernorm": true,
+  "architectures": [
+    "Dinov2Model"
+  ],
+  "attention_probs_dropout_prob": 0.0,
+  "drop_path_rate": 0.0,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.0,
+  "hidden_size": 768,
+  "image_size": 518,
+  "initializer_range": 0.02,
+  "layer_norm_eps": 1e-06,
+  "layerscale_value": 1.0,
+  "mlp_ratio": 4,
+  "model_type": "dinov2",
+  "num_attention_heads": 12,
+  "num_channels": 3,
+  "num_hidden_layers": 12,
+  "out_features": [
+    "stage12"
+  ],
+  "out_indices": [
+    12
+  ],
+  "patch_size": 14,
+  "qkv_bias": true,
+  "reshape_hidden_states": true,
+  "stage_names": [
+    "stem",
+    "stage1",
+    "stage2",
+    "stage3",
+    "stage4",
+    "stage5",
+    "stage6",
+    "stage7",
+    "stage8",
+    "stage9",
+    "stage10",
+    "stage11",
+    "stage12"
+  ],
+  "torch_dtype": "float32",
+  "transformers_version": "4.41.0",
+  "use_swiglu_ffn": false
+}

dino_head.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b3599663464ed1054f7777f547db02f518581acc5becdd3eddffc8c507f3778
+size 92554920

model.safetensors

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

preprocessor_config.json

@@ -0,0 +1,44 @@
+{
+  "_valid_processor_keys": [
+    "images",
+    "do_resize",
+    "size",
+    "resample",
+    "do_center_crop",
+    "crop_size",
+    "do_rescale",
+    "rescale_factor",
+    "do_normalize",
+    "image_mean",
+    "image_std",
+    "do_convert_rgb",
+    "return_tensors",
+    "data_format",
+    "input_data_format"
+  ],
+  "crop_size": {
+    "height": 518,
+    "width": 518
+  },
+  "do_center_crop": true,
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.5307,
+    0.5307,
+    0.5307
+  ],
+  "image_processor_type": "BitImageProcessor",
+  "image_std": [
+    0.2583,
+    0.2583,
+    0.2583
+  ],
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "shortest_edge": 518
+  }
+}

pyproject.toml

@@ -0,0 +1,34 @@
+[project]
+name = "rad-dino"
+version = "0.1.1"
+description = "Vision encoder for chest X-rays."
+readme = "README.md"
+authors = [
+    { name = "Microsoft Health Futures", email = "innereyedev@microsoft.com" },
+]
+requires-python = ">=3.10"
+dependencies = [
+    "einops",
+    "jaxtyping",
+    "pillow",
+    "requests",
+    "safetensors",
+    "transformers[torch]",
+    "typer",
+]
+
+[project.urls]
+Homepage = "https://huggingface.co/microsoft/rad-dino"
+Source = "https://huggingface.co/microsoft/rad-dino"
+"Issue tracker" = "https://huggingface.co/microsoft/rad-dino/discussions/new"
+Documentation = "https://huggingface.co/microsoft/rad-dino/blob/main/README.md"
+
+[build-system]
+requires = ["uv_build"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "ipykernel",
+    "ipywidgets",
+]

src/rad_dino/__init__.py

@@ -0,0 +1,74 @@
+import torch
+from einops import rearrange
+from jaxtyping import Float
+from PIL import Image
+from torch import Tensor
+from torch import nn
+from transformers import AutoImageProcessor
+from transformers import AutoModel
+from transformers.feature_extraction_utils import BatchFeature
+
+
+__version__ = "0.1.0"
+
+TypeClsToken = Float[Tensor, "batch_size embed_dim"]
+TypePatchTokensFlat = Float[Tensor, "batch_size (height width) embed_dim"]
+TypePatchTokens = Float[Tensor, "batch_size embed_dim height width"]
+TypeInputImages = Image.Image | list[Image.Image]
+
+
+class RadDino(nn.Module):
+    _REPO = "microsoft/rad-dino"
+
+    def __init__(self):
+        super().__init__()
+        self.model = AutoModel.from_pretrained(self._REPO).eval()
+        self.processor = AutoImageProcessor.from_pretrained(self._REPO, use_fast=False)
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.model.parameters()).device
+
+    def preprocess(self, image_or_images: TypeInputImages) -> BatchFeature:
+        return self.processor(image_or_images, return_tensors="pt")
+
+    def encode(self, inputs: BatchFeature) -> tuple[TypeClsToken, TypePatchTokensFlat]:
+        outputs = self.model(**inputs)
+        cls_token = outputs.last_hidden_state[:, 0]
+        patch_tokens = outputs.last_hidden_state[:, 1:]
+        return cls_token, patch_tokens
+
+    def reshape_patch_tokens(
+        self,
+        patch_tokens_flat: TypePatchTokensFlat,
+    ) -> TypePatchTokens:
+        input_size = self.processor.crop_size["height"]
+        patch_size = self.model.config.patch_size
+        embeddings_size = input_size // patch_size
+        patches_grid = rearrange(
+            patch_tokens_flat,
+            "batch (height width) embed_dim -> batch embed_dim height width",
+            height=embeddings_size,
+        )
+        return patches_grid
+
+    @torch.inference_mode()
+    def extract_features(
+        self,
+        image_or_images: TypeInputImages,
+    ) -> tuple[TypeClsToken, TypePatchTokens]:
+        inputs = self.preprocess(image_or_images).to(self.device)
+        cls_token, patch_tokens_flat = self.encode(inputs)
+        patch_tokens = self.reshape_patch_tokens(patch_tokens_flat)
+        return cls_token, patch_tokens
+
+    def extract_cls_token(self, image_or_images: TypeInputImages) -> TypeClsToken:
+        cls_token, _ = self.extract_features(image_or_images)
+        return cls_token
+
+    def extract_patch_tokens(self, image_or_images: TypeInputImages) -> TypePatchTokens:
+        _, patch_tokens = self.extract_features(image_or_images)
+        return patch_tokens
+
+    def forward(self, *args) -> tuple[TypeClsToken, TypePatchTokens]:
+        return self.extract_features(*args)

src/rad_dino/utils.py

@@ -0,0 +1,22 @@
+import requests
+import torch
+from PIL import Image
+from safetensors import safe_open
+
+
+def download_sample_image() -> Image.Image:
+    """Download chest X-ray with CC license."""
+    base_url = "https://upload.wikimedia.org/wikipedia/commons"
+    path = "2/20/Chest_X-ray_in_influenza_and_Haemophilus_influenzae.jpg"
+    image_url = f"{base_url}/{path}"
+    headers = {"User-Agent": "RAD-DINO"}
+    response = requests.get(image_url, headers=headers, stream=True)
+    return Image.open(response.raw)
+
+
+def safetensors_to_state_dict(checkpoint_path: str) -> dict[str, torch.Tensor]:
+    state_dict = {}
+    with safe_open(checkpoint_path, framework="pt") as ckpt_file:
+        for key in ckpt_file.keys():
+            state_dict[key] = ckpt_file.get_tensor(key)
+    return state_dict

ssl_default_config.yaml

@@ -0,0 +1,135 @@
+MODEL:
+  WEIGHTS: ''
+compute_precision:
+  grad_scaler: true
+  teacher:
+    backbone:
+      sharding_strategy: SHARD_GRAD_OP
+      mixed_precision:
+        param_dtype: fp16
+        reduce_dtype: fp16
+        buffer_dtype: fp32
+    dino_head:
+      sharding_strategy: SHARD_GRAD_OP
+      mixed_precision:
+        param_dtype: fp16
+        reduce_dtype: fp16
+        buffer_dtype: fp32
+    ibot_head:
+      sharding_strategy: SHARD_GRAD_OP
+      mixed_precision:
+        param_dtype: fp16
+        reduce_dtype: fp16
+        buffer_dtype: fp32
+  student:
+    backbone:
+      sharding_strategy: SHARD_GRAD_OP
+      mixed_precision:
+        param_dtype: fp16
+        reduce_dtype: fp16
+        buffer_dtype: fp32
+    dino_head:
+      sharding_strategy: SHARD_GRAD_OP
+      mixed_precision:
+        param_dtype: fp16
+        reduce_dtype: fp32
+        buffer_dtype: fp32
+    ibot_head:
+      sharding_strategy: SHARD_GRAD_OP
+      mixed_precision:
+        param_dtype: fp16
+        reduce_dtype: fp32
+        buffer_dtype: fp32
+dino:
+  loss_weight: 1.0
+  head_n_prototypes: 65536
+  head_bottleneck_dim: 256
+  head_nlayers: 3
+  head_hidden_dim: 2048
+  koleo_loss_weight: 0.1
+ibot:
+  loss_weight: 1.0
+  mask_sample_probability: 0.5
+  mask_ratio_min_max:
+  - 0.1
+  - 0.5
+  separate_head: false
+  head_n_prototypes: 65536
+  head_bottleneck_dim: 256
+  head_nlayers: 3
+  head_hidden_dim: 2048
+train:
+  batch_size_per_gpu: 64
+  dataset_path: ImageNet:split=TRAIN
+  output_dir: .
+  saveckp_every_n_epoch: 5
+  seed: 0
+  num_workers: 10
+  OFFICIAL_EPOCH_LENGTH: 0 # automatic rescaling based on the dataset len is applied if this is set to 0
+  cache_dataset: true
+  centering: "centering" # or "sinkhorn_knopp"
+student:
+  arch: vit_large
+  patch_size: 16
+  drop_block_rate: 0.0
+  drop_path_rate: 0.3
+  layerscale: 1.0e-05
+  drop_path_uniform: true
+  pretrained_weights: ''
+  ffn_layer: "mlp"
+  block_chunks: 0
+  qkv_bias: true
+  proj_bias: true
+  ffn_bias: true
+  num_register_tokens: 0
+  interpolate_antialias: false
+  interpolate_offset: 0.1
+  load_weights: true
+  checkpoints_dir: null
+teacher:
+  momentum_teacher: 0.992
+  final_momentum_teacher: 1
+  warmup_teacher_temp: 0.04
+  teacher_temp: 0.07
+  warmup_teacher_temp_epochs: 30
+optim:
+  epochs: 100
+  weight_decay: 0.04
+  weight_decay_end: 0.4
+  base_lr: 0.004  # learning rate for a batch size of 1024
+  lr: 0.  # will be set after applying scaling rule
+  warmup_epochs: 10
+  min_lr: 1.0e-06
+  clip_grad: 3.0
+  freeze_last_layer_epochs: 1
+  scaling_rule: sqrt_wrt_1024
+  patch_embed_lr_mult: 0.2
+  layerwise_decay: 0.9
+  adamw_beta1: 0.9
+  adamw_beta2: 0.999
+crops:
+  global_crops_scale:
+  - 0.32
+  - 1.0
+  local_crops_number: 8
+  local_crops_scale:
+  - 0.05
+  - 0.32
+  global_crops_size: 224
+  local_crops_size: 96
+evaluation:
+  eval_period_iterations: 12500
+  dataset_str: None
+  online:  # see dinov2.eval.linear_callback for documentation
+    learning_rate: 1e-6  # will be multiplied by batch size and number of devices
+    num_last_blocks: 1
+    add_avg_pool: true
+    num_update_epochs_per_eval: 3
+    augmentation:
+      degrees: 30
+      scale:
+      - 0.8
+      - 1.2
+      shear: 15
+      interpolation: BICUBIC
+      horizontal_flip: true

training_images.csv

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

vitb14_cxr.yaml

@@ -0,0 +1,31 @@
+# this corresponds to the CXR config
+train:
+  batch_size_per_gpu: 40  # For nodes with v100s (32 GB), use 20.
+  saveckp_every_n_epoch: 25
+student:
+  arch: vit_base
+  block_chunks: 4
+  patch_size: 14
+  drop_block_rate: 0.00
+  drop_path_rate: 0.30
+teacher:
+  warmup_teacher_temp_epochs: 50
+optim:
+  epochs: 100
+  warmup_epochs: 5
+  base_lr: 0.001
+evaluation:
+  eval_period_iterations: 300
+  tasks:  # from the metadata.csv file of the CANDID processed dataset
+  - pneumothorax
+crops:
+    global_crops_size: 518
+    local_crops_size: 196
+    global_crops_scale:
+    - 0.50
+    - 1.00
+    local_crops_number: 8
+    local_crops_scale:
+    - 0.20
+    - 0.50
+pretrained: true