kaiko-eva 0.1.7__tar.gz → 0.2.0__tar.gz

{kaiko_eva-0.1.7 → kaiko_eva-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: kaiko-eva
-Version: 0.1.7
+Version: 0.2.0
 Summary: Evaluation Framework for oncology foundation models.
 Keywords: machine-learning,evaluation-framework,oncology,foundation-models
 Author-Email: Ioannis Gatopoulos <ioannis@kaiko.ai>, =?utf-8?q?Nicolas_K=C3=A4nzig?= <nicolas@kaiko.ai>, Roman Moser <roman@kaiko.ai>
@@ -228,6 +228,7 @@ Requires-Dist: onnx>=1.16.0
 Requires-Dist: toolz>=0.12.1
 Requires-Dist: rich>=13.7.1
 Requires-Dist: torchmetrics>=1.6.0
+Requires-Dist: nibabel>=3.2.2
 Provides-Extra: vision
 Requires-Dist: h5py>=3.10.0; extra == "vision"
 Requires-Dist: nibabel>=4.0.1; extra == "vision"

{kaiko_eva-0.1.7 → kaiko_eva-0.2.0}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "pdm.backend"
 
 [project]
 name = "kaiko-eva"
-version = "0.1.7"
+version = "0.2.0"
 description = "Evaluation Framework for oncology foundation models."
 keywords = [
     "machine-learning",
@@ -46,6 +46,7 @@ dependencies = [
     "toolz>=0.12.1",
     "rich>=13.7.1",
     "torchmetrics>=1.6.0",
+    "nibabel>=3.2.2",
 ]
 
 [project.urls]

{kaiko_eva-0.1.7 → kaiko_eva-0.2.0}/src/eva/core/data/datasets/classification/embeddings.py

@@ -12,7 +12,7 @@ class EmbeddingsClassificationDataset(embeddings_base.EmbeddingsDataset[torch.Te
     """Embeddings dataset class for classification tasks."""
 
     @override
-    def _load_embeddings(self, index: int) -> torch.Tensor:
+    def load_embeddings(self, index: int) -> torch.Tensor:
         filename = self.filename(index)
         embeddings_path = os.path.join(self._root, filename)
         tensor = torch.load(embeddings_path, map_location="cpu")
@@ -25,7 +25,7 @@ class EmbeddingsClassificationDataset(embeddings_base.EmbeddingsDataset[torch.Te
         return tensor.squeeze(0)
 
     @override
-    def _load_target(self, index: int) -> torch.Tensor:
+    def load_target(self, index: int) -> torch.Tensor:
         target = self._data.at[index, self._column_mapping["target"]]
         return torch.tensor(target, dtype=torch.int64)
 

{kaiko_eva-0.1.7 → kaiko_eva-0.2.0}/src/eva/core/data/datasets/classification/multi_embeddings.py

@@ -66,7 +66,7 @@ class MultiEmbeddingsClassificationDataset(embeddings_base.EmbeddingsDataset[tor
         self._multi_ids = list(self._data[self._column_mapping["multi_id"]].unique())
 
     @override
-    def _load_embeddings(self, index: int) -> torch.Tensor:
+    def load_embeddings(self, index: int) -> torch.Tensor:
         """Loads and stacks all embedding corresponding to the `index`'th multi_id."""
         # Get all embeddings for the given index (multi_id)
         multi_id = self._multi_ids[index]
@@ -89,7 +89,7 @@ class MultiEmbeddingsClassificationDataset(embeddings_base.EmbeddingsDataset[tor
         return embeddings
 
     @override
-    def _load_target(self, index: int) -> np.ndarray:
+    def load_target(self, index: int) -> np.ndarray:
         """Returns the target corresponding to the `index`'th multi_id.
 
         This method assumes that all the embeddings corresponding to the same `multi_id`

{kaiko_eva-0.1.7 → kaiko_eva-0.2.0}/src/eva/core/data/datasets/embeddings.py

@@ -98,12 +98,12 @@ class EmbeddingsDataset(base.Dataset, Generic[TargetType]):
         Returns:
             A data sample and its target.
         """
-        embeddings = self._load_embeddings(index)
-        target = self._load_target(index)
+        embeddings = self.load_embeddings(index)
+        target = self.load_target(index)
         return self._apply_transforms(embeddings, target)
 
     @abc.abstractmethod
-    def _load_embeddings(self, index: int) -> torch.Tensor:
+    def load_embeddings(self, index: int) -> torch.Tensor:
         """Returns the `index`'th embedding sample.
 
         Args:
@@ -114,7 +114,7 @@ class EmbeddingsDataset(base.Dataset, Generic[TargetType]):
         """
 
     @abc.abstractmethod
-    def _load_target(self, index: int) -> TargetType:
+    def load_target(self, index: int) -> TargetType:
         """Returns the `index`'th target sample.
 
         Args:

{kaiko_eva-0.1.7 → kaiko_eva-0.2.0}/src/eva/core/data/samplers/classification/balanced.py

@@ -4,6 +4,7 @@ from collections import defaultdict
 from typing import Dict, Iterator, List
 
 import numpy as np
+from loguru import logger
 from typing_extensions import override
 
 from eva.core.data import datasets
@@ -33,6 +34,7 @@ class BalancedSampler(SamplerWithDataSource[int]):
         self._replacement = replacement
         self._class_indices: Dict[int, List[int]] = defaultdict(list)
         self._random_generator = np.random.default_rng(seed)
+        self._indices: List[int] = []
 
     def __len__(self) -> int:
         """Returns the total number of samples."""
@@ -44,18 +46,7 @@ class BalancedSampler(SamplerWithDataSource[int]):
         Returns:
             Iterator yielding dataset indices.
         """
-        indices = []
-
-        for class_idx in self._class_indices:
-            class_indices = self._class_indices[class_idx]
-            sampled_indices = self._random_generator.choice(
-                class_indices, size=self._num_samples, replace=self._replacement
-            ).tolist()
-            indices.extend(sampled_indices)
-
-        self._random_generator.shuffle(indices)
-
-        return iter(indices)
+        return iter(self._indices)
 
     @override
     def set_dataset(self, data_source: datasets.MapDataset):
@@ -72,13 +63,13 @@ class BalancedSampler(SamplerWithDataSource[int]):
         self._make_indices()
 
     def _make_indices(self):
-        """Builds indices for each class in the dataset."""
+        """Samples the indices for each class in the dataset."""
         self._class_indices.clear()
-
-        for idx in tqdm(
-            range(len(self.data_source)), desc="Fetching class indices for balanced sampler"
-        ):
-            _, target, _ = DataSample(*self.data_source[idx])
+        for idx in tqdm(range(len(self.data_source)), desc="Fetching class indices for sampler"):
+            if hasattr(self.data_source, "load_target"):
+                target = self.data_source.load_target(idx)  # type: ignore
+            else:
+                _, target, _ = DataSample(*self.data_source[idx])
             if target is None:
                 raise ValueError("The dataset must return non-empty targets.")
             if target.numel() != 1:
@@ -94,3 +85,13 @@ class BalancedSampler(SamplerWithDataSource[int]):
                         f"Class {class_idx} has only {len(indices)} samples, "
                         f"which is less than the required {self._num_samples} samples."
                     )
+
+        self._indices = []
+        for class_idx in self._class_indices:
+            class_indices = self._class_indices[class_idx]
+            sampled_indices = self._random_generator.choice(
+                class_indices, size=self._num_samples, replace=self._replacement
+            ).tolist()
+            self._indices.extend(sampled_indices)
+        self._random_generator.shuffle(self._indices)
+        logger.debug(f"Sampled indices: {self._indices}")

kaiko_eva-0.2.0/src/eva/core/loggers/utils/wandb.py

@@ -0,0 +1,33 @@
+# type: ignore
+"""Utility functions for logging with Weights & Biases."""
+
+from typing import Any, Dict
+
+from loguru import logger
+
+
+def rename_active_run(name: str) -> None:
+    """Renames the current run."""
+    import wandb
+
+    if wandb.run:
+        wandb.run.name = name
+        wandb.run.save()
+    else:
+        logger.warning("No active wandb run found that could be renamed.")
+
+
+def init_run(name: str, init_kwargs: Dict[str, Any]) -> None:
+    """Initializes a new run. If there is an active run, it will be renamed and reused."""
+    import wandb
+
+    init_kwargs["name"] = name
+    rename_active_run(name)
+    wandb.init(**init_kwargs)
+
+
+def finish_run() -> None:
+    """Finish the current run."""
+    import wandb
+
+    wandb.finish()

{kaiko_eva-0.1.7 → kaiko_eva-0.2.0}/src/eva/core/models/modules/head.py

@@ -1,4 +1,4 @@
-""""Neural Network Head Module."""
+"""Neural Network Head Module."""
 
 from typing import Any, Callable, Dict
 

{kaiko_eva-0.1.7 → kaiko_eva-0.2.0}/src/eva/core/trainers/functional.py

@@ -39,7 +39,7 @@ def run_evaluation_session(
             base_trainer,
             base_model,
             datamodule,
-            run_id=f"run_{run_index}",
+            run_id=run_index,
             verbose=not verbose,
         )
         recorder.update(validation_scores, test_scores)
@@ -51,7 +51,7 @@ def run_evaluation(
     base_model: modules.ModelModule,
     datamodule: datamodules.DataModule,
     *,
-    run_id: str | None = None,
+    run_id: int | None = None,
     verbose: bool = True,
 ) -> Tuple[_EVALUATE_OUTPUT, _EVALUATE_OUTPUT | None]:
     """Fits and evaluates a model out-of-place.
@@ -61,7 +61,6 @@ def run_evaluation(
         base_model: The model module to use but not modify.
         datamodule: The data module.
         run_id: The run id to be appended to the output log directory.
-            If `None`, it will use the log directory of the trainer as is.
         verbose: Whether to print the validation and test metrics
             in the end of the training.
 
@@ -70,8 +69,12 @@ def run_evaluation(
     """
     trainer, model = _utils.clone(base_trainer, base_model)
     model.configure_model()
-    trainer.setup_log_dirs(run_id or "")
-    return fit_and_validate(trainer, model, datamodule, verbose=verbose)
+
+    trainer.init_logger_run(run_id)
+    results = fit_and_validate(trainer, model, datamodule, verbose=verbose)
+    trainer.finish_logger_run(run_id)
+
+    return results
 
 
 def fit_and_validate(

{kaiko_eva-0.1.7 → kaiko_eva-0.2.0}/src/eva/core/trainers/trainer.py

@@ -12,6 +12,7 @@ from typing_extensions import override
 
 from eva.core import loggers as eva_loggers
 from eva.core.data import datamodules
+from eva.core.loggers.utils import wandb as wandb_utils
 from eva.core.models import modules
 from eva.core.trainers import _logging, functional
 
@@ -53,7 +54,7 @@ class Trainer(pl_trainer.Trainer):
         self._session_id: str = _logging.generate_session_id()
         self._log_dir: str = self.default_log_dir
 
-        self.setup_log_dirs()
+        self.init_logger_run(0)
 
     @property
     def default_log_dir(self) -> str:
@@ -65,31 +66,45 @@ class Trainer(pl_trainer.Trainer):
     def log_dir(self) -> str | None:
         return self.strategy.broadcast(self._log_dir)
 
-    def setup_log_dirs(self, subdirectory: str = "") -> None:
-        """Setups the logging directory of the trainer and experimental loggers in-place.
+    def init_logger_run(self, run_id: int | None) -> None:
+        """Setup the loggers & log directories when starting a new run.
 
         Args:
-            subdirectory: Whether to append a subdirectory to the output log.
+            run_id: The id of the current run.
         """
+        subdirectory = f"run_{run_id}" if run_id is not None else ""
         self._log_dir = os.path.join(self.default_root_dir, self._session_id, subdirectory)
 
         enabled_loggers = []
-        if isinstance(self.loggers, list) and len(self.loggers) > 0:
-            for logger in self.loggers:
-                if isinstance(logger, (pl_loggers.CSVLogger, pl_loggers.TensorBoardLogger)):
-                    if not cloud_io._is_local_file_protocol(self.default_root_dir):
-                        loguru.logger.warning(
-                            f"Skipped {type(logger).__name__} as remote storage is not supported."
-                        )
-                        continue
-                    else:
-                        logger._root_dir = self.default_root_dir
-                        logger._name = self._session_id
-                        logger._version = subdirectory
-                enabled_loggers.append(logger)
+        for logger in self.loggers or []:
+            if isinstance(logger, (pl_loggers.CSVLogger, pl_loggers.TensorBoardLogger)):
+                if not cloud_io._is_local_file_protocol(self.default_root_dir):
+                    loguru.logger.warning(
+                        f"Skipped {type(logger).__name__} as remote storage is not supported."
+                    )
+                    continue
+                else:
+                    logger._root_dir = self.default_root_dir
+                    logger._name = self._session_id
+                    logger._version = subdirectory
+            elif isinstance(logger, pl_loggers.WandbLogger):
+                task_name = self.default_root_dir.split("/")[-1]
+                run_name = os.getenv("WANDB_RUN_NAME", f"{task_name}_{self._session_id}")
+                wandb_utils.init_run(f"{run_name}_{run_id}", logger._wandb_init)
+            enabled_loggers.append(logger)
 
         self._loggers = enabled_loggers or [eva_loggers.DummyLogger(self._log_dir)]
 
+    def finish_logger_run(self, run_id: int | None) -> None:
+        """Finish the current run in the enabled loggers.
+
+        Args:
+            run_id: The id of the current run.
+        """
+        for logger in self.loggers or []:
+            if isinstance(logger, pl_loggers.WandbLogger):
+                wandb_utils.finish_run()
+
     def run_evaluation_session(
         self,
         model: modules.ModelModule,

{kaiko_eva-0.1.7 → kaiko_eva-0.2.0}/src/eva/vision/data/datasets/__init__.py

@@ -2,12 +2,16 @@
 
 from eva.vision.data.datasets.classification import (
     BACH,
+    BRACS,
     CRC,
     MHIST,
     PANDA,
+    BreaKHis,
     Camelyon16,
+    GleasonArvaniti,
     PANDASmall,
     PatchCamelyon,
+    UniToPatho,
     WsiClassificationDataset,
 )
 from eva.vision.data.datasets.segmentation import (
@@ -26,12 +30,16 @@ from eva.vision.data.datasets.wsi import MultiWsiDataset, WsiDataset
 __all__ = [
     "BACH",
     "BCSS",
+    "BreaKHis",
+    "BRACS",
     "CRC",
+    "GleasonArvaniti",
     "MHIST",
     "PANDA",
     "PANDASmall",
     "Camelyon16",
     "PatchCamelyon",
+    "UniToPatho",
     "WsiClassificationDataset",
     "CoNSeP",
     "EmbeddingsSegmentationDataset",

{kaiko_eva-0.1.7 → kaiko_eva-0.2.0}/src/eva/vision/data/datasets/classification/__init__.py

@@ -1,18 +1,27 @@
 """Image classification datasets API."""
 
 from eva.vision.data.datasets.classification.bach import BACH
+from eva.vision.data.datasets.classification.bracs import BRACS
+from eva.vision.data.datasets.classification.breakhis import BreaKHis
 from eva.vision.data.datasets.classification.camelyon16 import Camelyon16
 from eva.vision.data.datasets.classification.crc import CRC
+from eva.vision.data.datasets.classification.gleason_arvaniti import GleasonArvaniti
 from eva.vision.data.datasets.classification.mhist import MHIST
 from eva.vision.data.datasets.classification.panda import PANDA, PANDASmall
 from eva.vision.data.datasets.classification.patch_camelyon import PatchCamelyon
+from eva.vision.data.datasets.classification.unitopatho import UniToPatho
 from eva.vision.data.datasets.classification.wsi import WsiClassificationDataset
 
 __all__ = [
     "BACH",
+    "BreaKHis",
+    "BRACS",
+    "Camelyon16",
     "CRC",
+    "GleasonArvaniti",
     "MHIST",
     "PatchCamelyon",
+    "UniToPatho",
     "WsiClassificationDataset",
     "PANDA",
     "PANDASmall",

kaiko_eva-0.2.0/src/eva/vision/data/datasets/classification/bracs.py

@@ -0,0 +1,112 @@
+"""BRACS dataset class."""
+
+import os
+from typing import Callable, Dict, List, Literal, Tuple
+
+import torch
+from torchvision import tv_tensors
+from torchvision.datasets import folder
+from typing_extensions import override
+
+from eva.vision.data.datasets import _validators
+from eva.vision.data.datasets.classification import base
+from eva.vision.utils import io
+
+
+class BRACS(base.ImageClassification):
+    """Dataset class for BRACS images and corresponding targets."""
+
+    _expected_dataset_lengths: Dict[str, int] = {
+        "train": 3657,
+        "val": 312,
+        "test": 570,
+    }
+    """Expected dataset lengths for the splits and complete dataset."""
+
+    _license: str = "CC BY-NC 4.0 (https://creativecommons.org/licenses/by-nc/4.0/)"
+    """Dataset license."""
+
+    def __init__(
+        self,
+        root: str,
+        split: Literal["train", "val", "test"],
+        transforms: Callable | None = None,
+    ) -> None:
+        """Initializes the dataset.
+
+        Args:
+            root: Path to the root directory of the dataset.
+            split: Dataset split to use.
+            transforms: A function/transform which returns a transformed
+                version of the raw data samples.
+        """
+        super().__init__(transforms=transforms)
+
+        self._root = root
+        self._split = split
+
+        self._samples: List[Tuple[str, int]] = []
+
+    @property
+    @override
+    def classes(self) -> List[str]:
+        return ["0_N", "1_PB", "2_UDH", "3_FEA", "4_ADH", "5_DCIS", "6_IC"]
+
+    @property
+    @override
+    def class_to_idx(self) -> Dict[str, int]:
+        return {name: index for index, name in enumerate(self.classes)}
+
+    @override
+    def filename(self, index: int) -> str:
+        image_path, *_ = self._samples[index]
+        return os.path.relpath(image_path, self._dataset_path)
+
+    @override
+    def prepare_data(self) -> None:
+        _validators.check_dataset_exists(self._root, True)
+
+    @override
+    def configure(self) -> None:
+        self._samples = self._make_dataset()
+
+    @override
+    def validate(self) -> None:
+        _validators.check_dataset_integrity(
+            self,
+            length=self._expected_dataset_lengths[self._split],
+            n_classes=7,
+            first_and_last_labels=("0_N", "6_IC"),
+        )
+
+    @override
+    def load_image(self, index: int) -> tv_tensors.Image:
+        image_path, _ = self._samples[index]
+        return io.read_image_as_tensor(image_path)
+
+    @override
+    def load_target(self, index: int) -> torch.Tensor:
+        _, target = self._samples[index]
+        return torch.tensor(target, dtype=torch.long)
+
+    @override
+    def __len__(self) -> int:
+        return len(self._samples)
+
+    @property
+    def _dataset_path(self) -> str:
+        """Returns the full path of dataset directory."""
+        return os.path.join(self._root, "BRACS_RoI/latest_version")
+
+    def _make_dataset(self) -> List[Tuple[str, int]]:
+        """Builds the dataset for the specified split."""
+        dataset = folder.make_dataset(
+            directory=os.path.join(self._dataset_path, self._split),
+            class_to_idx=self.class_to_idx,
+            extensions=(".png"),
+        )
+        return dataset
+
+    def _print_license(self) -> None:
+        """Prints the dataset license."""
+        print(f"Dataset license: {self._license}")