Add option to WSI datasets to save patch coordinates

configs/vision/pathology/offline/classification/camelyon16.yaml

@@ -34,7 +34,7 @@ trainer:
             0: train
             1: val
             2: test
-          metadata_keys: ["wsi_id"]
+          metadata_keys: ["wsi_id", "x", "y", "width", "height", "level_idx"]
           backbone:
             class_path: eva.vision.models.ModelFromRegistry
             init_args:

configs/vision/pathology/offline/classification/camelyon16_small.yaml

@@ -34,7 +34,7 @@ trainer:
             0: train
             1: val
             2: test
-          metadata_keys: ["wsi_id"]
+          metadata_keys: ["wsi_id", "x", "y", "width", "height", "level_idx"]
           backbone:
             class_path: eva.vision.models.ModelFromRegistry
             init_args:

configs/vision/pathology/offline/classification/panda.yaml

@@ -33,7 +33,7 @@ trainer:
             0: train
             1: val
             2: test
-          metadata_keys: ["wsi_id"]
+          metadata_keys: ["wsi_id", "x", "y", "width", "height", "level_idx"]
           backbone:
             class_path: eva.vision.models.ModelFromRegistry
             init_args:

configs/vision/pathology/offline/classification/panda_small.yaml

@@ -33,7 +33,7 @@ trainer:
             0: train
             1: val
             2: test
-          metadata_keys: ["wsi_id"]
+          metadata_keys: ["wsi_id", "x", "y", "width", "height", "level_idx"]
           backbone:
             class_path: eva.vision.models.ModelFromRegistry
             init_args:

configs/vision/tests/offline/panda.yaml

@@ -14,7 +14,7 @@ trainer:
             0: train
             1: val
             2: test
-          metadata_keys: ["wsi_id"]
+          metadata_keys: ["wsi_id", "x", "y", "width", "height", "level_idx"]
           backbone:
             class_path: eva.models.ModelFromFunction
             init_args:

src/eva/vision/data/datasets/classification/camelyon16.py

@@ -207,7 +207,9 @@ def load_target(self, index: int) -> torch.Tensor:
 
     @override
     def load_metadata(self, index: int) -> Dict[str, Any]:
-        return {"wsi_id": self.filename(index).split(".")[0]}
+        dataset_index, sample_index = self._get_dataset_idx(index), self._get_sample_idx(index)
+        patch_metadata = self.datasets[dataset_index].load_metadata(sample_index)
+        return {"wsi_id": self.filename(index).split(".")[0]} | patch_metadata
 
     def _load_file_paths(self, split: Literal["train", "val", "test"] | None = None) -> List[str]:
         """Loads the file paths of the corresponding dataset split."""

src/eva/vision/data/datasets/classification/panda.py

@@ -132,7 +132,9 @@ def load_target(self, index: int) -> torch.Tensor:
 
     @override
     def load_metadata(self, index: int) -> Dict[str, Any]:
-        return {"wsi_id": self.filename(index).split(".")[0]}
+        dataset_index, sample_index = self._get_dataset_idx(index), self._get_sample_idx(index)
+        patch_metadata = self.datasets[dataset_index].load_metadata(sample_index)
+        return {"wsi_id": self.filename(index).split(".")[0]} | patch_metadata
 
     def _load_file_paths(self, split: Literal["train", "val", "test"] | None = None) -> List[str]:
         """Loads the file paths of the corresponding dataset split."""

src/eva/vision/data/datasets/classification/wsi.py

@@ -88,7 +88,9 @@ def load_target(self, index: int) -> np.ndarray:
 
     @override
     def load_metadata(self, index: int) -> Dict[str, Any]:
-        return {"wsi_id": self.filename(index).split(".")[0]}
+        dataset_index, sample_index = self._get_dataset_idx(index), self._get_sample_idx(index)
+        patch_metadata = self.datasets[dataset_index].load_metadata(sample_index)
+        return {"wsi_id": self.filename(index).split(".")[0]} | patch_metadata
 
     def _load_manifest(self, manifest_path: str) -> pd.DataFrame:
         df = pd.read_csv(manifest_path)

src/eva/vision/data/datasets/wsi.py

@@ -2,7 +2,7 @@
 
 import bisect
 import os
-from typing import Callable, List
+from typing import Any, Callable, Dict, List
 
 from loguru import logger
 from torch.utils.data import dataset as torch_datasets
@@ -85,6 +85,17 @@ def __getitem__(self, index: int) -> tv_tensors.Image:
         patch = self._apply_transforms(patch)
         return patch
 
+    def load_metadata(self, index: int) -> Dict[str, Any]:
+        """Loads the metadata for the patch at the specified index."""
+        x, y = self._coords.x_y[index]
+        return {
+            "x": x,
+            "y": y,
+            "width": self._coords.width,
+            "height": self._coords.height,
+            "level_idx": self._coords.level_idx,
+        }
+
     def _apply_transforms(self, image: tv_tensors.Image) -> tv_tensors.Image:
         if self._image_transforms is not None:
             image = self._image_transforms(image)
@@ -185,3 +196,7 @@ def _load_datasets(self) -> list[WsiDataset]:
 
     def _get_dataset_idx(self, index: int) -> int:
         return bisect.bisect_right(self.cumulative_sizes, index)
+
+    def _get_sample_idx(self, index: int) -> int:
+        dataset_idx = self._get_dataset_idx(index)
+        return index if dataset_idx == 0 else index - self.cumulative_sizes[dataset_idx - 1]

tests/eva/core/callbacks/writers/embeddings/test_classification.py

@@ -29,7 +29,7 @@
         (5, 7, None, None),
         (5, 7, ["wsi_id"], None),
         (8, 16, None, None),
-        (8, 32, ["wsi_id"], ["slide_1", "slide_2"]),
+        (8, 32, ["wsi_id", "x", "y"], ["slide_1", "slide_2"]),
     ],
 )
 def test_embeddings_writer(datamodule: datamodules.DataModule, model: modules.HeadModule) -> None:

tests/eva/vision/data/datasets/classification/test_camelyon16.py

@@ -69,6 +69,11 @@ def _check_batch_shape(batch: Any):
     assert isinstance(target, torch.Tensor)
     assert isinstance(metadata, dict)
     assert "wsi_id" in metadata
+    assert "x" in metadata
+    assert "y" in metadata
+    assert "width" in metadata
+    assert "height" in metadata
+    assert "level_idx" in metadata
 
 
 @pytest.fixture

tests/eva/vision/data/datasets/classification/test_panda.py

@@ -102,6 +102,11 @@ def _check_batch_shape(batch: Any):
     assert isinstance(target, torch.Tensor)
     assert isinstance(metadata, dict)
     assert "wsi_id" in metadata
+    assert "x" in metadata
+    assert "y" in metadata
+    assert "width" in metadata
+    assert "height" in metadata
+    assert "level_idx" in metadata
 
 
 @pytest.fixture

tests/eva/vision/data/datasets/classification/test_wsi.py

@@ -79,6 +79,11 @@ def _check_batch_shape(batch: Any):
 
     assert isinstance(metadata, dict)
     assert "wsi_id" in metadata
+    assert "x" in metadata
+    assert "y" in metadata
+    assert "width" in metadata
+    assert "height" in metadata
+    assert "level_idx" in metadata
 
 
 @pytest.fixture