Split code in blocks

README.md

@@ -51,31 +51,36 @@ Fine-tuning RAD-DINO is typically not necessary to obtain good performance in do
 <!-- 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 characterised.
+Underlying biases of the training datasets may not be well characterized.
 
 ## Getting started
 
+Let us first write an auxiliary function to download a chest X-ray.
+
 ```python
->>> import torch
+>>> import requests
 >>> from PIL import Image
->>> from transformers import AutoModel
->>> from transformers import AutoImageProcessor
->>>
->>> # Define a small function to get a sample image
 >>> def download_sample_image() -> Image.Image:
 ...     """Download chest X-ray with CC license."""
-...     import requests
-...     from PIL import Image
 ...     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": "fperezgarcia@microsoft.com"}
+...     headers = {"User-Agent": "RAD-DINO"}
 ...     response = requests.get(image_url, headers=headers, stream=True)
 ...     return Image.open(response.raw)
 ...
+```
+
+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"
 >>> model = 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
@@ -95,8 +100,12 @@ Underlying biases of the training datasets may not be well characterised.
 >>> cls_embeddings = outputs.pooler_output
 >>> cls_embeddings.shape  # (batch_size, num_channels)
 torch.Size([1, 768])
->>>
->>> # Look at the patch embeddings (needs `pip install einops`)
+```
+
+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