update

README.md

@@ -8,14 +8,78 @@ tags:
 ### Usage
 
 ```python
-import timm
-from timm.data import resolve_data_config
-from timm.data.transforms_factory import create_transform
+from transformers import AutoImageProcessor, AutoModel
+from PIL import Image
+import requests
 from huggingface_hub import login
+from torchvision.transforms import v2
 
 login()  # login or use an access token
 
-model = timm.create_model("hf-hub:kaiko-ai/midnight", pretrained=True)
-transform = create_transform(**resolve_data_config(model.pretrained_cfg, model=model))
-model.eval()
+# FYI: here a natural image instead of a crop of a WSI for simplicity
+url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
+image = Image.open(requests.get(url, stream=True).raw)
+
+transform = v2.Compose(
+    [
+        v2.Resize(224),
+        v2.CenterCrop(224),
+        v2.ToTensor(),
+        v2.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
+    ]
+)
+model = AutoModel.from_pretrained('kaiko-ai/midnight')
+```
+
+
+### Extract embeddings for classification
+```python
+import torch
+from transformers import modeling_outputs
+from typing_extensions import override
+
+# for classification
+class ExtractConcatToken:
+    """Extracts the CLS with Mean Patch tokens from a model output."""
+    def __call__(self, tensor: torch.Tensor) -> torch.Tensor:
+        num_reg_tokens = 0
+        return torch.cat([tensor[:, 0, :], tensor[:, 1 + num_reg_tokens :, :].mean(1)], dim=-1)
+
+extract_embeddings = ExtractConcatToken()
+emb = extract_embeddings(model(transform(image)[None]).last_hidden_state)
+print(f"Embeddings shape: {emb.shape}")
+```
+
+
+### Extract embeddings for segmentation
+```python
+import math
+import torch
+from transformers import modeling_outputs
+from typing_extensions import override
+
+# for segmentation
+class ExtractPatchFeatures:
+    """Extracts the patch features from a model output."""
+    def __call__(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Call method for the transformation.
+
+        Args:
+            tensor: The raw embeddings of the model.
+
+        Returns:
+            A tensor (batch_size, hidden_size, n_patches_height, n_patches_width)
+            representing the model output.
+        """
+        num_reg_tokens = 0
+        num_skip = 1 + num_reg_tokens
+        features = tensor[:, num_skip:, :].permute(0, 2, 1)
+        batch_size, hidden_size, patch_grid = features.shape
+        height = width = int(math.sqrt(patch_grid))
+        assert height * width == patch_grid
+        return features.view(batch_size, hidden_size, height, width)
+
+extract_embeddings = ExtractPatchFeatures()
+emb = extract_embeddings(model(transform(image)[None]).last_hidden_state)
+print(f"Embeddings shape for segmentation: {emb[0].shape}")
 ```