Add package code update docs

.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

README.md

@@ -56,113 +56,43 @@ Fine-tuning RAD-DINO is typically not necessary to obtain good performance in do
 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.
 
-## Getting started
+## Installation
 
-### Get some data
-
-Let us first write an auxiliary function to download a chest X-ray.
-
-```python
->>> import requests
->>> from PIL import Image
->>> def download_sample_image() -> Image.Image:
-...     """Download chest X-ray with CC license."""
-...     base_url = "https://upload.wikimedia.org/wikipedia/commons"
-...     image_url = f"{base_url}/2/20/Chest_X-ray_in_influenza_and_Haemophilus_influenzae.jpg"
-...     headers = {"User-Agent": "RAD-DINO"}
-...     response = requests.get(image_url, headers=headers, stream=True)
-...     return Image.open(response.raw)
-...
+```shell
+pip install rad-dino
 ```
 
-### Load the model
-
-Now let us download the model and encode an image.
-
-```python
->>> import torch
->>> from transformers import AutoModel
->>> from transformers import AutoImageProcessor
->>>
->>> # Download the model
->>> repo = "microsoft/rad-dino"
->>> rad_dino = AutoModel.from_pretrained(repo)
->>>
->>> # The processor takes a PIL image, performs resizing, center-cropping, and
->>> # intensity normalization using stats from MIMIC-CXR, and returns a
->>> # dictionary with a PyTorch tensor ready for the encoder
->>> processor = AutoImageProcessor.from_pretrained(repo)
-```
+## Usage
 
 ### Encode an image
 
 ```python
->>> # Download and preprocess a chest X-ray
+>>> from rad_dino import RadDino
+>>> from rad_dino.utils import download_sample_image
+>>> encoder = RadDino()
 >>> image = download_sample_image()
->>> image.size  # (width, height)
-(2765, 2505)
->>> inputs = processor(images=image, return_tensors="pt")
->>>
->>> # Encode the image!
->>> with torch.inference_mode():
->>>     outputs = rad_dino(**inputs)
->>>
->>> # Look at the CLS embeddings
->>> cls_embeddings = outputs.pooler_output
->>> cls_embeddings.shape  # (batch_size, num_channels)
-torch.Size([1, 768])
-```
-
-If we are interested in the feature maps, we can reshape the patch embeddings into a grid.
-We will use [`einops`](https://einops.rocks/) (install with `pip install einops`) for this.
-
-```python
->>> def reshape_patch_embeddings(flat_tokens: torch.Tensor) -> torch.Tensor:
-...     """Reshape flat list of patch tokens into a nice grid."""
-...     from einops import rearrange
-...     image_size = processor.crop_size["height"]
-...     patch_size = model.config.patch_size
-...     embeddings_size = image_size // patch_size
-...     patches_grid = rearrange(flat_tokens, "b (h w) c -> b c h w", h=embeddings_size)
-...     return patches_grid
-...
->>> flat_patch_embeddings = outputs.last_hidden_state[:, 1:]  # first token is CLS
->>> reshaped_patch_embeddings = reshape_patch_embeddings(flat_patch_embeddings)
->>> reshaped_patch_embeddings.shape  # (batch_size, num_channels, height, width)
-torch.Size([1, 768, 37, 37])
+>>> image
+<PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=2765x2505 at 0x7CCD5C014050>
+>>> cls_token, patch_tokens = 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.
-
-First, let us write code to load a
-[`safetensors` checkpoint](https://huggingface.co/docs/safetensors).
-
-```python
->>> import safetensors
->>> 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
-...
-```
+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 now use the hub model and load the RAD-DINO weights.
+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
-cd dinov2
 ```
 
 ```python
 >>> import torch
->>> rad_dino_gh = torch.hub.load(".", "dinov2_vitb14")
+>>> 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>
@@ -272,20 +202,20 @@ We used [SimpleITK](https://simpleitk.org/) and [Pydicom](https://pydicom.github
 
 ```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},
+ 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
+> 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
 

pyproject.toml

@@ -0,0 +1,26 @@
+[project]
+name = "rad-dino"
+version = "0.1.0"
+description = "Vision encoder for chest X-rays."
+readme = "README.md"
+authors = [{ name = "Microsoft Health Futures" }]
+requires-python = ">=3.10"
+dependencies = [
+    "einops",
+    "jaxtyping",
+    "pillow",
+    "requests",
+    "safetensors",
+    "transformers[torch]",
+    "typer>=0.16.0",
+]
+
+[project.scripts]
+rad-dino = "rad_dino.__main__:main"
+
+[build-system]
+requires = ["uv_build>=0.8.3,<0.9.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = ["ipykernel", "ipywidgets"]

src/rad_dino/__init__.py

@@ -0,0 +1,70 @@
+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.image_processing_base 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)
+
+    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)
+        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/__main__.py

@@ -0,0 +1,6 @@
+def main():
+    print("Hello from rad-dino!")
+
+
+if __name__ == "__main__":
+    main()

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