kaiko-eva 0.1.1__tar.gz → 0.1.5__tar.gz

{kaiko_eva-0.1.1 → kaiko_eva-0.1.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: kaiko-eva
-Version: 0.1.1
+Version: 0.1.5
 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>
@@ -216,14 +216,14 @@ Project-URL: Homepage, https://kaiko-ai.github.io/eva/dev/
 Project-URL: Repository, https://github.com/kaiko-ai/eva
 Project-URL: Documentation, https://kaiko-ai.github.io/eva/dev/
 Requires-Python: >=3.10
-Requires-Dist: torch==2.3.0
-Requires-Dist: lightning>=2.2.2
-Requires-Dist: jsonargparse[omegaconf]==4.31.0
+Requires-Dist: torch>=2.3.0
+Requires-Dist: lightning>=2.2.0
+Requires-Dist: jsonargparse[omegaconf]>=4.30.0
 Requires-Dist: tensorboard>=2.16.2
 Requires-Dist: loguru>=0.7.2
-Requires-Dist: pandas>=2.2.0
+Requires-Dist: pandas>=2.0.0
 Requires-Dist: transformers>=4.38.2
-Requires-Dist: onnxruntime>=1.17.1
+Requires-Dist: onnxruntime>=1.15.1
 Requires-Dist: onnx>=1.16.0
 Requires-Dist: toolz>=0.12.1
 Requires-Dist: rich>=13.7.1
@@ -468,41 +468,10 @@ and [tutorials](https://kaiko-ai.github.io/eva/dev/user-guide/advanced/replicate
 
 ## Leaderboards
 
-In this section you will find model benchmarks which were generated with _`eva`_.
+The following table shows the FMs we have evaluated with _`eva`_. For more detailed information about the evaluation process, please refer to our [documentation](https://kaiko-ai.github.io/eva/main/leaderboards/).
 
-### Table I: WSI and microscopy image tasks
+![Pathology Leaderboard](./docs/images/leaderboard.svg)
 
-<br />
-
-<div align="center">
-
-| Model   | BACH  | CRC   | MHIST | PCam   | Camelyon16 | PANDA | CoNSeP | MoNuSAC |
-|---------|-------|-------|-------|--------|------------|-------|------------|-------|
-| ViT-S/16 _(random)_	<sup>[1]</sup> | 0.411|0.613|0.5|0.752|0.551|0.347|0.489|0.394|
-| ViT-S/16 _(ImageNet)_ <sup>[1]</sup> | 0.675|0.936|0.827|0.861|0.751|0.676|0.54|0.512|
-| DINO<sub>(p=16)</sub> <sup>[2]</sup> | 0.77|0.936|0.751|0.905|0.869|0.737|0.625|0.549|
-| Phikon <sup>[3]</sup>                | 0.715|0.942|0.766|0.925|0.879|0.784|0.68|0.554|
-| UNI <sup>[4]</sup>                   | 0.797|0.95|0.835|0.939|0.933|0.774|0.67|0.575|
-| ViT-S/16 _(kaiko.ai)_ <sup>[5]</sup> | 0.8|0.949|0.831|0.902|0.897|0.77|0.622|0.573|
-| ViT-S/8 _(kaiko.ai)_ <sup>[5]</sup>  | 0.825|0.948|0.826|0.887|0.879|0.741|0.677|0.617|
-| ViT-B/16 _(kaiko.ai)_	<sup>[5]</sup> | 0.846|0.959|0.839|0.906|0.891|0.753|0.647|0.572|
-| ViT-B/8 _(kaiko.ai)_ <sup>[5]</sup>  | 0.867|0.952|0.814|0.921|0.939|0.761|0.706|0.661|
-| ViT-L/14 _(kaiko.ai)_ <sup>[5]</sup> | 0.862|0.935|0.822|0.907|0.941|0.769|0.686|0.599|
-
-_Table I: Linear probing evaluation of FMs on patch-level downstream datasets.<br> We report balanced accuracy
- for classification tasks and generalized Dice score for semgetnation tasks, averaged over 5 runs. Results are 
- reported on the "test" split if available and otherwise on the "validation" split._
-
-</div>
-
-<br />
-
-_References_:
-1. _"Emerging properties in self-supervised vision transformers”_, [arXiv](https://arxiv.org/abs/2104.14294)
-2. _"Benchmarking self-supervised learning on diverse pathology datasets”_, [arXiv](https://arxiv.org/abs/2212.04690)
-3. _"Scaling self-supervised learning for histopathology with masked image modeling”_, [medRxiv](https://www.medrxiv.org/content/10.1101/2023.07.21.23292757v1)
-4. _"A General-Purpose Self-Supervised Model for Computational Pathology”_, [arXiv](https://arxiv.org/abs/2308.15474)
-5. _"Towards Training Large-Scale Pathology Foundation Models: from TCGA to Hospital Scale”_, [arXiv](https://arxiv.org/pdf/2404.15217)
 
 ## Contributing
 

{kaiko_eva-0.1.1 → kaiko_eva-0.1.5}/README.md

@@ -212,41 +212,10 @@ and [tutorials](https://kaiko-ai.github.io/eva/dev/user-guide/advanced/replicate
 
 ## Leaderboards
 
-In this section you will find model benchmarks which were generated with _`eva`_.
+The following table shows the FMs we have evaluated with _`eva`_. For more detailed information about the evaluation process, please refer to our [documentation](https://kaiko-ai.github.io/eva/main/leaderboards/).
 
-### Table I: WSI and microscopy image tasks
+![Pathology Leaderboard](./docs/images/leaderboard.svg)
 
-<br />
-
-<div align="center">
-
-| Model   | BACH  | CRC   | MHIST | PCam   | Camelyon16 | PANDA | CoNSeP | MoNuSAC |
-|---------|-------|-------|-------|--------|------------|-------|------------|-------|
-| ViT-S/16 _(random)_	<sup>[1]</sup> | 0.411|0.613|0.5|0.752|0.551|0.347|0.489|0.394|
-| ViT-S/16 _(ImageNet)_ <sup>[1]</sup> | 0.675|0.936|0.827|0.861|0.751|0.676|0.54|0.512|
-| DINO<sub>(p=16)</sub> <sup>[2]</sup> | 0.77|0.936|0.751|0.905|0.869|0.737|0.625|0.549|
-| Phikon <sup>[3]</sup>                | 0.715|0.942|0.766|0.925|0.879|0.784|0.68|0.554|
-| UNI <sup>[4]</sup>                   | 0.797|0.95|0.835|0.939|0.933|0.774|0.67|0.575|
-| ViT-S/16 _(kaiko.ai)_ <sup>[5]</sup> | 0.8|0.949|0.831|0.902|0.897|0.77|0.622|0.573|
-| ViT-S/8 _(kaiko.ai)_ <sup>[5]</sup>  | 0.825|0.948|0.826|0.887|0.879|0.741|0.677|0.617|
-| ViT-B/16 _(kaiko.ai)_	<sup>[5]</sup> | 0.846|0.959|0.839|0.906|0.891|0.753|0.647|0.572|
-| ViT-B/8 _(kaiko.ai)_ <sup>[5]</sup>  | 0.867|0.952|0.814|0.921|0.939|0.761|0.706|0.661|
-| ViT-L/14 _(kaiko.ai)_ <sup>[5]</sup> | 0.862|0.935|0.822|0.907|0.941|0.769|0.686|0.599|
-
-_Table I: Linear probing evaluation of FMs on patch-level downstream datasets.<br> We report balanced accuracy
- for classification tasks and generalized Dice score for semgetnation tasks, averaged over 5 runs. Results are 
- reported on the "test" split if available and otherwise on the "validation" split._
-
-</div>
-
-<br />
-
-_References_:
-1. _"Emerging properties in self-supervised vision transformers”_, [arXiv](https://arxiv.org/abs/2104.14294)
-2. _"Benchmarking self-supervised learning on diverse pathology datasets”_, [arXiv](https://arxiv.org/abs/2212.04690)
-3. _"Scaling self-supervised learning for histopathology with masked image modeling”_, [medRxiv](https://www.medrxiv.org/content/10.1101/2023.07.21.23292757v1)
-4. _"A General-Purpose Self-Supervised Model for Computational Pathology”_, [arXiv](https://arxiv.org/abs/2308.15474)
-5. _"Towards Training Large-Scale Pathology Foundation Models: from TCGA to Hospital Scale”_, [arXiv](https://arxiv.org/pdf/2404.15217)
 
 ## Contributing
 

{kaiko_eva-0.1.1 → kaiko_eva-0.1.5}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "pdm.backend"
 
 [project]
 name = "kaiko-eva"
-version = "0.1.1"
+version = "0.1.5"
 description = "Evaluation Framework for oncology foundation models."
 keywords = [
     "machine-learning",
@@ -34,14 +34,14 @@ maintainers = [
 ]
 requires-python = ">=3.10"
 dependencies = [
-    "torch==2.3.0",
-    "lightning>=2.2.2",
-    "jsonargparse[omegaconf]==4.31.0",
+    "torch>=2.3.0",
+    "lightning>=2.2.0",
+    "jsonargparse[omegaconf]>=4.30.0",
     "tensorboard>=2.16.2",
     "loguru>=0.7.2",
-    "pandas>=2.2.0",
+    "pandas>=2.0.0",
     "transformers>=4.38.2",
-    "onnxruntime>=1.17.1",
+    "onnxruntime>=1.15.1",
     "onnx>=1.16.0",
     "toolz>=0.12.1",
     "rich>=13.7.1",

{kaiko_eva-0.1.1 → kaiko_eva-0.1.5}/src/eva/core/callbacks/writers/embeddings/base.py

@@ -172,15 +172,14 @@ class EmbeddingsWriter(callbacks.BasePredictionWriter, abc.ABC):
 
     def _check_if_exists(self) -> None:
         """Checks if the output directory already exists and if it should be overwritten."""
-        try:
-            os.makedirs(self._output_dir, exist_ok=self._overwrite)
-        except FileExistsError as e:
+        os.makedirs(self._output_dir, exist_ok=True)
+        if os.path.exists(os.path.join(self._output_dir, "manifest.csv")) and not self._overwrite:
             raise FileExistsError(
                 f"The embeddings output directory already exists: {self._output_dir}. This "
                 "either means that they have been computed before or that a wrong output "
                 "directory is being used. Consider using `eva fit` instead, selecting a "
                 "different output directory or setting overwrite=True."
-            ) from e
+            )
         os.makedirs(self._output_dir, exist_ok=True)
 
 

{kaiko_eva-0.1.1 → kaiko_eva-0.1.5}/src/eva/core/data/dataloaders/dataloader.py

@@ -38,7 +38,7 @@ class DataLoader:
     Mutually exclusive with `batch_size`, `shuffle`, `sampler` and `drop_last`.
     """
 
-    num_workers: int = multiprocessing.cpu_count()
+    num_workers: int | None = None
     """How many workers to use for loading the data.
 
     By default, it will use the number of CPUs available.
@@ -71,7 +71,7 @@ class DataLoader:
             shuffle=self.shuffle,
             sampler=self.sampler,
             batch_sampler=self.batch_sampler,
-            num_workers=self.num_workers,
+            num_workers=self.num_workers or multiprocessing.cpu_count(),
             collate_fn=self.collate_fn,
             pin_memory=self.pin_memory,
             drop_last=self.drop_last,

{kaiko_eva-0.1.1 → kaiko_eva-0.1.5}/src/eva/core/data/splitting/random.py

@@ -24,12 +24,13 @@ def random_split(
     Returns:
         The indices of the train, validation, and test sets as lists.
     """
-    if train_ratio + val_ratio + (test_ratio or 0) != 1:
-        raise ValueError("The sum of the ratios must be equal to 1.")
+    total_ratio = train_ratio + val_ratio + test_ratio
+    if total_ratio > 1.0:
+        raise ValueError("The sum of the ratios must be lower or equal to 1.")
 
-    np.random.seed(seed)
-    n_samples = len(samples)
-    indices = np.random.permutation(n_samples)
+    random_generator = np.random.default_rng(seed)
+    n_samples = int(total_ratio * len(samples))
+    indices = random_generator.permutation(len(samples))[:n_samples]
 
     n_train = int(np.floor(train_ratio * n_samples))
     n_val = n_samples - n_train if test_ratio == 0.0 else int(np.floor(val_ratio * n_samples)) or 1

{kaiko_eva-0.1.1 → kaiko_eva-0.1.5}/src/eva/core/data/splitting/stratified.py

@@ -28,10 +28,11 @@ def stratified_split(
     """
     if len(samples) != len(targets):
         raise ValueError("The number of samples and targets must be equal.")
-    if train_ratio + val_ratio + (test_ratio or 0) != 1:
-        raise ValueError("The sum of the ratios must be equal to 1.")
+    if train_ratio + val_ratio + (test_ratio or 0) > 1.0:
+        raise ValueError("The sum of the ratios must be lower or equal to 1.")
 
-    np.random.seed(seed)
+    use_all_samples = train_ratio + val_ratio + test_ratio == 1
+    random_generator = np.random.default_rng(seed)
     unique_classes, y_indices = np.unique(targets, return_inverse=True)
     n_classes = unique_classes.shape[0]
 
@@ -39,18 +40,23 @@ def stratified_split(
 
     for c in range(n_classes):
         class_indices = np.where(y_indices == c)[0]
-        np.random.shuffle(class_indices)
+        random_generator.shuffle(class_indices)
 
         n_train = int(np.floor(train_ratio * len(class_indices))) or 1
         n_val = (
             len(class_indices) - n_train
-            if test_ratio == 0.0
+            if test_ratio == 0.0 and use_all_samples
             else int(np.floor(val_ratio * len(class_indices))) or 1
         )
 
         train_indices.extend(class_indices[:n_train])
         val_indices.extend(class_indices[n_train : n_train + n_val])
         if test_ratio > 0.0:
-            test_indices.extend(class_indices[n_train + n_val :])
+            n_test = (
+                len(class_indices) - n_train - n_val
+                if use_all_samples
+                else int(np.floor(test_ratio * len(class_indices))) or 1
+            )
+            test_indices.extend(class_indices[n_train + n_val : n_train + n_val + n_test])
 
     return train_indices, val_indices, test_indices or None

kaiko_eva-0.1.5/src/eva/core/losses/__init__.py

@@ -0,0 +1,5 @@
+"""Loss functions API."""
+
+from eva.core.losses.cross_entropy import CrossEntropyLoss
+
+__all__ = ["CrossEntropyLoss"]

kaiko_eva-0.1.5/src/eva/core/losses/cross_entropy.py

@@ -0,0 +1,27 @@
+"""Cross-entropy based loss function."""
+
+from typing import Sequence
+
+import torch
+from torch import nn
+
+
+class CrossEntropyLoss(nn.CrossEntropyLoss):
+    """A wrapper around torch.nn.CrossEntropyLoss that accepts weights in list format.
+
+    Needed for .yaml file loading & class instantiation with jsonarparse.
+    """
+
+    def __init__(
+        self, *args, weight: Sequence[float] | torch.Tensor | None = None, **kwargs
+    ) -> None:
+        """Initialize the loss function.
+
+        Args:
+            args: Positional arguments from the base class.
+            weight: A list of weights to assign to each class.
+            kwargs: Key-word arguments from the base class.
+        """
+        if weight is not None and not isinstance(weight, torch.Tensor):
+            weight = torch.tensor(weight)
+        super().__init__(*args, **kwargs, weight=weight)

{kaiko_eva-0.1.1 → kaiko_eva-0.1.5}/src/eva/core/metrics/__init__.py

@@ -3,8 +3,6 @@
 from eva.core.metrics.average_loss import AverageLoss
 from eva.core.metrics.binary_balanced_accuracy import BinaryBalancedAccuracy
 from eva.core.metrics.defaults import BinaryClassificationMetrics, MulticlassClassificationMetrics
-from eva.core.metrics.generalized_dice import GeneralizedDiceScore
-from eva.core.metrics.mean_iou import MeanIoU
 from eva.core.metrics.structs import Metric, MetricCollection, MetricModule, MetricsSchema
 
 __all__ = [
@@ -12,8 +10,6 @@ __all__ = [
     "BinaryBalancedAccuracy",
     "BinaryClassificationMetrics",
     "MulticlassClassificationMetrics",
-    "GeneralizedDiceScore",
-    "MeanIoU",
     "Metric",
     "MetricCollection",
     "MetricModule",

{kaiko_eva-0.1.1 → kaiko_eva-0.1.5}/src/eva/core/metrics/defaults/__init__.py

@@ -4,10 +4,8 @@ from eva.core.metrics.defaults.classification import (
     BinaryClassificationMetrics,
     MulticlassClassificationMetrics,
 )
-from eva.core.metrics.defaults.segmentation import MulticlassSegmentationMetrics
 
 __all__ = [
     "MulticlassClassificationMetrics",
     "BinaryClassificationMetrics",
-    "MulticlassSegmentationMetrics",
 ]

{kaiko_eva-0.1.1 → kaiko_eva-0.1.5}/src/eva/core/models/modules/module.py

@@ -1,10 +1,10 @@
 """Base model module."""
 
+import os
 from typing import Any, Mapping
 
 import lightning.pytorch as pl
 import torch
-from lightning.pytorch.strategies.single_device import SingleDeviceStrategy
 from lightning.pytorch.utilities import memory
 from lightning.pytorch.utilities.types import STEP_OUTPUT
 from typing_extensions import override
@@ -49,14 +49,14 @@ class ModelModule(pl.LightningModule):
 
     @property
     def metrics_device(self) -> torch.device:
-        """Returns the device by which the metrics should be calculated.
-
-        We allocate the metrics to CPU when operating on single device, as
-        it is much faster, but to GPU when employing multiple ones, as DDP
-        strategy requires the metrics to be allocated to the module's GPU.
-        """
-        move_to_cpu = isinstance(self.trainer.strategy, SingleDeviceStrategy)
-        return torch.device("cpu") if move_to_cpu else self.device
+        """Returns the device by which the metrics should be calculated."""
+        device = os.getenv("METRICS_DEVICE", None)
+        if device is not None:
+            return torch.device(device)
+        elif self.device.type == "mps":
+            # mps seems to have compatibility issues with segmentation metrics
+            return torch.device("cpu")
+        return self.device
 
     @override
     def on_fit_start(self) -> None:

kaiko_eva-0.1.5/src/eva/core/models/transforms/extract_cls_features.py

@@ -0,0 +1,41 @@
+"""Transforms for extracting the CLS output from a model output."""
+
+import torch
+from transformers import modeling_outputs
+
+
+class ExtractCLSFeatures:
+    """Extracts the CLS token from a ViT model output."""
+
+    def __init__(
+        self, cls_index: int = 0, num_register_tokens: int = 0, include_patch_tokens: bool = False
+    ) -> None:
+        """Initializes the transformation.
+
+        Args:
+            cls_index: The index of the CLS token in the output tensor.
+            num_register_tokens: The number of register tokens in the model output.
+            include_patch_tokens: Whether to concat the mean aggregated patch tokens with
+                the cls token.
+        """
+        self._cls_index = cls_index
+        self._num_register_tokens = num_register_tokens
+        self._include_patch_tokens = include_patch_tokens
+
+    def __call__(
+        self, tensor: torch.Tensor | modeling_outputs.BaseModelOutputWithPooling
+    ) -> torch.Tensor:
+        """Call method for the transformation.
+
+        Args:
+            tensor: The tensor representing the model output.
+        """
+        if isinstance(tensor, modeling_outputs.BaseModelOutputWithPooling):
+            tensor = tensor.last_hidden_state
+
+        cls_token = tensor[:, self._cls_index, :]
+        if self._include_patch_tokens:
+            patch_tokens = tensor[:, 1 + self._num_register_tokens :, :]
+            return torch.cat([cls_token, patch_tokens.mean(1)], dim=-1)
+
+        return cls_token

kaiko_eva-0.1.5/src/eva/core/models/transforms/extract_patch_features.py

@@ -0,0 +1,59 @@
+"""Transforms for extracting the patch features from a model output."""
+
+import math
+from typing import List
+
+import torch
+from transformers import modeling_outputs
+
+
+class ExtractPatchFeatures:
+    """Extracts the patch features from a ViT model output."""
+
+    def __init__(
+        self,
+        has_cls_token: bool = True,
+        num_register_tokens: int = 0,
+        ignore_remaining_dims: bool = False,
+    ) -> None:
+        """Initializes the transformation.
+
+        Args:
+            has_cls_token: If set to `True`, the model output is expected to have
+                a classification token.
+            num_register_tokens: The number of register tokens in the model output.
+            ignore_remaining_dims: If set to `True`, ignore the remaining dimensions
+                of the patch grid if it is not a square number.
+        """
+        self._has_cls_token = has_cls_token
+        self._num_register_tokens = num_register_tokens
+        self._ignore_remaining_dims = ignore_remaining_dims
+
+    def __call__(
+        self, tensor: torch.Tensor | modeling_outputs.BaseModelOutputWithPooling
+    ) -> List[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_skip = int(self._has_cls_token) + self._num_register_tokens
+        if isinstance(tensor, modeling_outputs.BaseModelOutputWithPooling):
+            features = tensor.last_hidden_state[:, num_skip:, :].permute(0, 2, 1)
+        else:
+            features = tensor[:, num_skip:, :].permute(0, 2, 1)
+
+        batch_size, hidden_size, patch_grid = features.shape
+        height = width = int(math.sqrt(patch_grid))
+        if height * width != patch_grid:
+            if self._ignore_remaining_dims:
+                features = features[:, :, -height * width :]
+            else:
+                raise ValueError(f"Patch grid size must be a square number {patch_grid}.")
+        patch_embeddings = features.view(batch_size, hidden_size, height, width)
+
+        return [patch_embeddings]

kaiko_eva-0.1.5/src/eva/core/utils/io/__init__.py

@@ -0,0 +1,6 @@
+"""Core I/O utilities."""
+
+from eva.core.utils.io.dataframe import read_dataframe
+from eva.core.utils.io.gz import gunzip_file
+
+__all__ = ["read_dataframe", "gunzip_file"]

kaiko_eva-0.1.5/src/eva/core/utils/io/gz.py

@@ -0,0 +1,28 @@
+"""Utils for .gz files."""
+
+import gzip
+import os
+
+
+def gunzip_file(path: str, unpack_dir: str | None = None, keep: bool = True) -> str:
+    """Unpacks a .gz file to the provided directory.
+
+    Args:
+        path: Path to the .gz file to extract.
+        unpack_dir: Directory to extract the file to. If `None`, it will use the
+            same directory as the compressed file.
+        keep: Whether to keep the compressed .gz file.
+
+    Returns:
+        The path to the extracted file.
+    """
+    unpack_dir = unpack_dir or os.path.dirname(path)
+    os.makedirs(unpack_dir, exist_ok=True)
+    save_path = os.path.join(unpack_dir, os.path.basename(path).replace(".gz", ""))
+    if not os.path.isfile(save_path):
+        with gzip.open(path, "rb") as f_in:
+            with open(save_path, "wb") as f_out:
+                f_out.write(f_in.read())
+        if not keep:
+            os.remove(path)
+    return save_path

kaiko_eva-0.1.5/src/eva/core/utils/multiprocessing.py

@@ -0,0 +1,89 @@
+"""Multiprocessing utilities."""
+
+import multiprocessing
+import sys
+import traceback
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, TypeVar
+
+from eva.core.utils.progress_bar import tqdm
+
+
+class Process(multiprocessing.Process):
+    """Multiprocessing wrapper with logic to propagate exceptions to the parent process.
+
+    Source: https://stackoverflow.com/a/33599967/4992248
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        """Initialize the process."""
+        multiprocessing.Process.__init__(self, *args, **kwargs)
+
+        self._parent_conn, self._child_conn = multiprocessing.Pipe()
+        self._exception = None
+
+    def run(self) -> None:
+        """Run the process."""
+        try:
+            multiprocessing.Process.run(self)
+            self._child_conn.send(None)
+        except Exception as e:
+            tb = traceback.format_exc()
+            self._child_conn.send((e, tb))
+
+    @property
+    def exception(self):
+        """Property that contains exception information from the process."""
+        if self._parent_conn.poll():
+            self._exception = self._parent_conn.recv()
+        return self._exception
+
+    def check_exceptions(self) -> None:
+        """Check for exception propagate it to the parent process."""
+        if not self.is_alive():
+            if self.exception:
+                error, traceback = self.exception
+                sys.stderr.write(traceback + "\n")
+                raise error
+
+
+R = TypeVar("R")
+
+
+def run_with_threads(
+    func: Callable[..., R],
+    items: Iterable[Tuple[Any, ...]],
+    kwargs: Dict[str, Any] | None = None,
+    num_workers: int = 8,
+    progress_desc: Optional[str] = None,
+    show_progress: bool = True,
+    return_results: bool = True,
+) -> List[R] | None:
+    """Process items with multiple threads using ThreadPoolExecutor.
+
+    Args:
+        func: Function to execute for each item
+        items: Iterable of items to process. Each item should be a tuple of
+            arguments to pass to func.
+        kwargs: Additional keyword arguments to pass to func.
+        num_workers: Number of worker threads
+        progress_desc: Description for progress bar
+        show_progress: Whether to show progress bar
+        return_results: Whether to return the results. If False, the function
+            will return None.
+
+    Returns:
+        List of results if return_results is True, otherwise None
+    """
+    results: List[Any] = []
+
+    with ThreadPoolExecutor(max_workers=num_workers) as executor:
+        futures = [executor.submit(func, *args, **(kwargs or {})) for args in items]
+        pbar = tqdm(total=len(futures), desc=progress_desc, disable=not show_progress, leave=False)
+        for future in as_completed(futures):
+            if return_results:
+                results.append(future.result())
+            pbar.update(1)
+        pbar.close()
+
+    return results if return_results else None

kaiko_eva-0.1.5/src/eva/core/utils/progress_bar.py

@@ -0,0 +1,15 @@
+"""Progress bar utility functions."""
+
+import os
+
+from tqdm import tqdm as _tqdm
+
+
+def tqdm(*args, **kwargs) -> _tqdm:
+    """Wrapper function for `tqdm.tqdm`."""
+    refresh_rate = os.environ.get("TQDM_REFRESH_RATE")
+    refresh_rate = int(refresh_rate) if refresh_rate is not None else None
+    disable = bool(int(os.environ.get("TQDM_DISABLE", 0))) or (refresh_rate == 0)
+    kwargs.setdefault("disable", disable)
+    kwargs.setdefault("miniters", refresh_rate)
+    return _tqdm(*args, **kwargs)

{kaiko_eva-0.1.1 → kaiko_eva-0.1.5}/src/eva/vision/callbacks/loggers/batch/segmentation.py

@@ -128,7 +128,7 @@ def _draw_semantic_mask(tensor: torch.Tensor) -> torch.Tensor:
     integer values which represent the pixel class id.
 
     Args:
-        tensor: An image tensor of range [0., 1.].
+        tensor: An image tensor of range [0., N_CLASSES].
 
     Returns:
         The image as a tensor of range [0., 255.].
@@ -136,9 +136,11 @@ def _draw_semantic_mask(tensor: torch.Tensor) -> torch.Tensor:
     tensor = torch.squeeze(tensor)
     height, width = tensor.shape[-2], tensor.shape[-1]
     red, green, blue = torch.zeros((3, height, width), dtype=torch.uint8)
-    for class_id, color in colormap.COLORMAP.items():
+    class_ids = torch.unique(tensor)
+    colors = colormap.get_colors(max(class_ids))
+    for class_id in class_ids:
         indices = tensor == class_id
-        red[indices], green[indices], blue[indices] = color
+        red[indices], green[indices], blue[indices] = colors[int(class_id)]
     return torch.stack([red, green, blue])
 
 
@@ -157,8 +159,9 @@ def _overlay_mask(image: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
         from the predefined colormap.
     """
     binary_masks = functional.one_hot(mask).permute(2, 0, 1).to(dtype=torch.bool)
+    colors = colormap.get_colors(binary_masks.shape[0] + 1)
     return torchvision.utils.draw_segmentation_masks(
-        image, binary_masks[1:], alpha=0.65, colors=colormap.COLORS[1:]  # type: ignore
+        image, binary_masks[1:], alpha=0.65, colors=colors[1:]  # type: ignore
     )
 
 

{kaiko_eva-0.1.1 → kaiko_eva-0.1.5}/src/eva/vision/data/datasets/__init__.py

@@ -6,6 +6,7 @@ from eva.vision.data.datasets.classification import (
     MHIST,
     PANDA,
     Camelyon16,
+    PANDASmall,
     PatchCamelyon,
     WsiClassificationDataset,
 )
@@ -15,6 +16,7 @@ from eva.vision.data.datasets.segmentation import (
     EmbeddingsSegmentationDataset,
     ImageSegmentation,
     LiTS,
+    LiTSBalanced,
     MoNuSAC,
     TotalSegmentator2D,
 )
@@ -27,6 +29,7 @@ __all__ = [
     "CRC",
     "MHIST",
     "PANDA",
+    "PANDASmall",
     "Camelyon16",
     "PatchCamelyon",
     "WsiClassificationDataset",
@@ -34,6 +37,7 @@ __all__ = [
     "EmbeddingsSegmentationDataset",
     "ImageSegmentation",
     "LiTS",
+    "LiTSBalanced",
     "MoNuSAC",
     "TotalSegmentator2D",
     "VisionDataset",