kaiko-eva 0.1.6__py3-none-any.whl → 0.1.8__py3-none-any.whl

eva/core/models/wrappers/from_torchhub.py

@@ -0,0 +1,93 @@
+"""Model wrapper for torch.hub models."""
+
+from typing import Any, Callable, Dict, List, Tuple
+
+import torch
+import torch.nn as nn
+from typing_extensions import override
+
+from eva.core.models import wrappers
+from eva.core.models.wrappers import _utils
+
+
+class TorchHubModel(wrappers.BaseModel):
+    """Model wrapper for `torch.hub` models."""
+
+    def __init__(
+        self,
+        model_name: str,
+        repo_or_dir: str,
+        pretrained: bool = True,
+        checkpoint_path: str = "",
+        out_indices: int | Tuple[int, ...] | None = None,
+        norm: bool = False,
+        trust_repo: bool = True,
+        model_kwargs: Dict[str, Any] | None = None,
+        tensor_transforms: Callable | None = None,
+    ) -> None:
+        """Initializes the encoder.
+
+        Args:
+            model_name: Name of model to instantiate.
+            repo_or_dir: The torch.hub repository or local directory to load the model from.
+            pretrained: If set to `True`, load pretrained ImageNet-1k weights.
+            checkpoint_path: Path of checkpoint to load.
+            out_indices: Returns last n blocks if `int`, all if `None`, select
+                matching indices if sequence.
+            norm: Wether to apply norm layer to all intermediate features. Only
+                used when `out_indices` is not `None`.
+            trust_repo: If set to `False`, a prompt will ask the user whether the
+                repo should be trusted.
+            model_kwargs: Extra model arguments.
+            tensor_transforms: The transforms to apply to the output tensor
+                produced by the model.
+        """
+        super().__init__(tensor_transforms=tensor_transforms)
+
+        self._model_name = model_name
+        self._repo_or_dir = repo_or_dir
+        self._pretrained = pretrained
+        self._checkpoint_path = checkpoint_path
+        self._out_indices = out_indices
+        self._norm = norm
+        self._trust_repo = trust_repo
+        self._model_kwargs = model_kwargs or {}
+
+        self.load_model()
+
+    @override
+    def load_model(self) -> None:
+        """Builds and loads the torch.hub model."""
+        self._model: nn.Module = torch.hub.load(
+            repo_or_dir=self._repo_or_dir,
+            model=self._model_name,
+            trust_repo=self._trust_repo,
+            pretrained=self._pretrained,
+            **self._model_kwargs,
+        )  # type: ignore
+
+        if self._checkpoint_path:
+            _utils.load_model_weights(self._model, self._checkpoint_path)
+
+        TorchHubModel.__name__ = self._model_name
+
+    @override
+    def model_forward(self, tensor: torch.Tensor) -> torch.Tensor | List[torch.Tensor]:
+        if self._out_indices is not None:
+            if not hasattr(self._model, "get_intermediate_layers"):
+                raise ValueError(
+                    "Only models with `get_intermediate_layers` are supported "
+                    "when using `out_indices`."
+                )
+
+            return list(
+                self._model.get_intermediate_layers(
+                    tensor,
+                    self._out_indices,
+                    reshape=True,
+                    return_class_token=False,
+                    norm=self._norm,
+                )
+            )
+
+        return self._model(tensor)

eva/core/trainers/functional.py

@@ -96,11 +96,13 @@ def fit_and_validate(
         A tuple of with the validation and the test metrics (if exists).
     """
     trainer.fit(model, datamodule=datamodule)
-    validation_scores = trainer.validate(datamodule=datamodule, verbose=verbose)
+    validation_scores = trainer.validate(
+        datamodule=datamodule, verbose=verbose, ckpt_path=trainer.checkpoint_type
+    )
     test_scores = (
         None
         if datamodule.datasets.test is None
-        else trainer.test(datamodule=datamodule, verbose=verbose)
+        else trainer.test(datamodule=datamodule, verbose=verbose, ckpt_path=trainer.checkpoint_type)
     )
     return validation_scores, test_scores
 

eva/core/trainers/trainer.py

@@ -1,7 +1,7 @@
 """Core trainer module."""
 
 import os
-from typing import Any
+from typing import Any, Literal
 
 import loguru
 from lightning.pytorch import loggers as pl_loggers
@@ -28,6 +28,7 @@ class Trainer(pl_trainer.Trainer):
         *args: Any,
         default_root_dir: str = "logs",
         n_runs: int = 1,
+        checkpoint_type: Literal["best", "last"] = "best",
         **kwargs: Any,
     ) -> None:
         """Initializes the trainer.
@@ -40,11 +41,14 @@ class Trainer(pl_trainer.Trainer):
                 Unlike in ::class::`lightning.pytorch.Trainer`, this path would be the
                 prioritized destination point.
             n_runs: The amount of runs (fit and evaluate) to perform in an evaluation session.
+            checkpoint_type: Wether to load the "best" or "last" checkpoint saved by the checkpoint
+                callback for evaluations on validation & test sets.
             kwargs: Kew-word arguments of ::class::`lightning.pytorch.Trainer`.
         """
         super().__init__(*args, default_root_dir=default_root_dir, **kwargs)
 
-        self._n_runs = n_runs
+        self.checkpoint_type = checkpoint_type
+        self.n_runs = n_runs
 
         self._session_id: str = _logging.generate_session_id()
         self._log_dir: str = self.default_log_dir
@@ -106,6 +110,6 @@ class Trainer(pl_trainer.Trainer):
             base_trainer=self,
             base_model=model,
             datamodule=datamodule,
-            n_runs=self._n_runs,
-            verbose=self._n_runs > 1,
+            n_runs=self.n_runs,
+            verbose=self.n_runs > 1,
         )

eva/vision/data/datasets/segmentation/_total_segmentator.py

@@ -0,0 +1,91 @@
+"""Utils for TotalSegmentator dataset classes."""
+
+from typing import Dict
+
+reduced_class_mappings: Dict[str, str] = {
+    # Abdominal Organs
+    "spleen": "spleen",
+    "kidney_right": "kidney",
+    "kidney_left": "kidney",
+    "gallbladder": "gallbladder",
+    "liver": "liver",
+    "stomach": "stomach",
+    "pancreas": "pancreas",
+    "small_bowel": "small_bowel",
+    "duodenum": "duodenum",
+    "colon": "colon",
+    # Endocrine System
+    "adrenal_gland_right": "adrenal_gland",
+    "adrenal_gland_left": "adrenal_gland",
+    "thyroid_gland": "thyroid_gland",
+    # Respiratory System
+    "lung_upper_lobe_left": "lungs",
+    "lung_lower_lobe_left": "lungs",
+    "lung_upper_lobe_right": "lungs",
+    "lung_middle_lobe_right": "lungs",
+    "lung_lower_lobe_right": "lungs",
+    "trachea": "trachea",
+    "esophagus": "esophagus",
+    # Urogenital System
+    "urinary_bladder": "urogenital_system",
+    "prostate": "urogenital_system",
+    "kidney_cyst_left": "kidney_cyst",
+    "kidney_cyst_right": "kidney_cyst",
+    # Vertebral Column
+    **{f"vertebrae_{v}": "vertebrae" for v in ["C1", "C2", "C3", "C4", "C5", "C6", "C7"]},
+    **{f"vertebrae_{v}": "vertebrae" for v in [f"T{i}" for i in range(1, 13)]},
+    **{f"vertebrae_{v}": "vertebrae" for v in [f"L{i}" for i in range(1, 6)]},
+    "vertebrae_S1": "vertebrae",
+    "sacrum": "sacral_spine",
+    # Cardiovascular System
+    "heart": "heart",
+    "aorta": "aorta",
+    "pulmonary_vein": "veins",
+    "brachiocephalic_trunk": "arteries",
+    "subclavian_artery_right": "arteries",
+    "subclavian_artery_left": "arteries",
+    "common_carotid_artery_right": "arteries",
+    "common_carotid_artery_left": "arteries",
+    "brachiocephalic_vein_left": "veins",
+    "brachiocephalic_vein_right": "veins",
+    "atrial_appendage_left": "atrial_appendage",
+    "superior_vena_cava": "veins",
+    "inferior_vena_cava": "veins",
+    "portal_vein_and_splenic_vein": "veins",
+    "iliac_artery_left": "arteries",
+    "iliac_artery_right": "arteries",
+    "iliac_vena_left": "veins",
+    "iliac_vena_right": "veins",
+    # Upper Extremity Bones
+    "humerus_left": "humerus",
+    "humerus_right": "humerus",
+    "scapula_left": "scapula",
+    "scapula_right": "scapula",
+    "clavicula_left": "clavicula",
+    "clavicula_right": "clavicula",
+    # Lower Extremity Bones
+    "femur_left": "femur",
+    "femur_right": "femur",
+    "hip_left": "hip",
+    "hip_right": "hip",
+    # Muscles
+    "gluteus_maximus_left": "gluteus",
+    "gluteus_maximus_right": "gluteus",
+    "gluteus_medius_left": "gluteus",
+    "gluteus_medius_right": "gluteus",
+    "gluteus_minimus_left": "gluteus",
+    "gluteus_minimus_right": "gluteus",
+    "autochthon_left": "autochthon",
+    "autochthon_right": "autochthon",
+    "iliopsoas_left": "iliopsoas",
+    "iliopsoas_right": "iliopsoas",
+    # Central Nervous System
+    "brain": "brain",
+    "spinal_cord": "spinal_cord",
+    # Skull and Thoracic Cage
+    "skull": "skull",
+    **{f"rib_left_{i}": "ribs" for i in range(1, 13)},
+    **{f"rib_right_{i}": "ribs" for i in range(1, 13)},
+    "costal_cartilages": "ribs",
+    "sternum": "sternum",
+}

eva/vision/data/datasets/segmentation/consep.py

@@ -20,9 +20,12 @@ from eva.vision.utils import io
 class CoNSeP(wsi.MultiWsiDataset, base.ImageSegmentation):
     """Dataset class for CoNSeP semantic segmentation task.
 
-    We combine classes 3 (healthy epithelial) & 4 (dysplastic/malignant epithelial)
+    As in [1], we combine classes 3 (healthy epithelial) & 4 (dysplastic/malignant epithelial)
     into the epithelial class and 5 (fibroblast), 6 (muscle) & 7 (endothelial) into
     the spindle-shaped class.
+
+    [1] Graham, Simon, et al. "Hover-net: Simultaneous segmentation and classification of
+        nuclei in multi-tissue histology images." https://arxiv.org/abs/1802.04712
     """
 
     _expected_dataset_lengths: Dict[str | None, int] = {

eva/vision/data/datasets/segmentation/lits.py

@@ -76,7 +76,7 @@ class LiTS(base.ImageSegmentation):
     @property
     @override
     def classes(self) -> List[str]:
-        return ["liver", "tumor"]
+        return ["background", "liver", "tumor"]
 
     @functools.cached_property
     @override
@@ -105,8 +105,8 @@ class LiTS(base.ImageSegmentation):
         _validators.check_dataset_integrity(
             self,
             length=self._expected_dataset_lengths.get(self._split, 0),
-            n_classes=2,
-            first_and_last_labels=("liver", "tumor"),
+            n_classes=3,
+            first_and_last_labels=("background", "tumor"),
         )
 
     @override

eva/vision/data/datasets/segmentation/total_segmentator_2d.py

@@ -1,7 +1,9 @@
 """TotalSegmentator 2D segmentation dataset class."""
 
 import functools
+import hashlib
 import os
+import re
 from glob import glob
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Literal, Tuple
@@ -16,7 +18,7 @@ from typing_extensions import override
 from eva.core.utils import io as core_io
 from eva.core.utils import multiprocessing
 from eva.vision.data.datasets import _validators, structs
-from eva.vision.data.datasets.segmentation import base
+from eva.vision.data.datasets.segmentation import _total_segmentator, base
 from eva.vision.utils import io
 
 
@@ -66,6 +68,7 @@ class TotalSegmentator2D(base.ImageSegmentation):
         version: Literal["small", "full"] | None = "full",
         download: bool = False,
         classes: List[str] | None = None,
+        class_mappings: Dict[str, str] | None = _total_segmentator.reduced_class_mappings,
         optimize_mask_loading: bool = True,
         decompress: bool = True,
         num_workers: int = 10,
@@ -85,6 +88,8 @@ class TotalSegmentator2D(base.ImageSegmentation):
                 exist yet on disk.
             classes: Whether to configure the dataset with a subset of classes.
                 If `None`, it will use all of them.
+            class_mappings: A dictionary that maps the original class names to a
+                reduced set of classes. If `None`, it will use the original classes.
             optimize_mask_loading: Whether to pre-process the segmentation masks
                 in order to optimize the loading time. In the `setup` method, it
                 will reformat the binary one-hot masks to a semantic mask and store
@@ -109,11 +114,10 @@ class TotalSegmentator2D(base.ImageSegmentation):
         self._optimize_mask_loading = optimize_mask_loading
         self._decompress = decompress
         self._num_workers = num_workers
+        self._class_mappings = class_mappings
 
-        if self._optimize_mask_loading and self._classes is not None:
-            raise ValueError(
-                "To use customize classes please set the optimize_mask_loading to `False`."
-            )
+        if self._classes and self._class_mappings:
+            raise ValueError("Both 'classes' and 'class_mappings' cannot be set at the same time.")
 
         self._samples_dirs: List[str] = []
         self._indices: List[Tuple[int, int]] = []
@@ -125,16 +129,21 @@ class TotalSegmentator2D(base.ImageSegmentation):
             """Returns the filename from the full path."""
             return os.path.basename(path).split(".")[0]
 
-        first_sample_labels = os.path.join(
-            self._root, self._samples_dirs[0], "segmentations", "*.nii.gz"
-        )
+        first_sample_labels = os.path.join(self._root, "s0011", "segmentations", "*.nii.gz")
         all_classes = sorted(map(get_filename, glob(first_sample_labels)))
         if self._classes:
             is_subset = all(name in all_classes for name in self._classes)
             if not is_subset:
-                raise ValueError("Provided class names are not subset of the dataset onces.")
-
-        return all_classes if self._classes is None else self._classes
+                raise ValueError("Provided class names are not subset of the original ones.")
+            classes = sorted(self._classes)
+        elif self._class_mappings:
+            is_subset = all(name in all_classes for name in self._class_mappings.keys())
+            if not is_subset:
+                raise ValueError("Provided class names are not subset of the original ones.")
+            classes = sorted(set(self._class_mappings.values()))
+        else:
+            classes = all_classes
+        return ["background"] + classes
 
     @property
     @override
@@ -145,6 +154,10 @@ class TotalSegmentator2D(base.ImageSegmentation):
     def _file_suffix(self) -> str:
         return "nii" if self._decompress else "nii.gz"
 
+    @functools.cached_property
+    def _classes_hash(self) -> str:
+        return hashlib.md5(str(self.classes).encode(), usedforsecurity=False).hexdigest()
+
     @override
     def filename(self, index: int) -> str:
         sample_idx, _ = self._indices[index]
@@ -170,15 +183,22 @@ class TotalSegmentator2D(base.ImageSegmentation):
         if self._version is None or self._sample_every_n_slices is not None:
             return
 
+        if self._classes:
+            last_label = self._classes[-1]
+            n_classes = len(self._classes)
+        elif self._class_mappings:
+            classes = sorted(set(self._class_mappings.values()))
+            last_label = classes[-1]
+            n_classes = len(classes)
+        else:
+            last_label = "vertebrae_T9"
+            n_classes = 117
+
         _validators.check_dataset_integrity(
             self,
             length=self._expected_dataset_lengths.get(f"{self._split}_{self._version}", 0),
-            n_classes=len(self._classes) if self._classes else 117,
-            first_and_last_labels=(
-                (self._classes[0], self._classes[-1])
-                if self._classes
-                else ("adrenal_gland_left", "vertebrae_T9")
-            ),
+            n_classes=n_classes + 1,
+            first_and_last_labels=("background", last_label),
         )
 
     @override
@@ -190,32 +210,31 @@ class TotalSegmentator2D(base.ImageSegmentation):
         sample_index, slice_index = self._indices[index]
         image_path = self._get_image_path(sample_index)
         image_array = io.read_nifti(image_path, slice_index)
-        image_rgb_array = image_array.repeat(3, axis=2)
-        return tv_tensors.Image(image_rgb_array.transpose(2, 0, 1))
+        image_array = self._fix_orientation(image_array)
+        return tv_tensors.Image(image_array.copy().transpose(2, 0, 1))
 
     @override
     def load_mask(self, index: int) -> tv_tensors.Mask:
         if self._optimize_mask_loading:
-            return self._load_semantic_label_mask(index)
-        return self._load_mask(index)
+            mask = self._load_semantic_label_mask(index)
+        else:
+            mask = self._load_mask(index)
+        mask = self._fix_orientation(mask)
+        return tv_tensors.Mask(mask.copy().squeeze(), dtype=torch.int64)  # type: ignore
 
     @override
     def load_metadata(self, index: int) -> Dict[str, Any]:
         _, slice_index = self._indices[index]
         return {"slice_index": slice_index}
 
-    def _load_mask(self, index: int) -> tv_tensors.Mask:
+    def _load_mask(self, index: int) -> npt.NDArray[Any]:
         sample_index, slice_index = self._indices[index]
-        semantic_labels = self._load_masks_as_semantic_label(sample_index, slice_index)
-        return tv_tensors.Mask(semantic_labels.squeeze(), dtype=torch.int64)  # type: ignore[reportCallIssue]
+        return self._load_masks_as_semantic_label(sample_index, slice_index)
 
-    def _load_semantic_label_mask(self, index: int) -> tv_tensors.Mask:
+    def _load_semantic_label_mask(self, index: int) -> npt.NDArray[Any]:
         """Loads the segmentation mask from a semantic label NifTi file."""
         sample_index, slice_index = self._indices[index]
-        masks_dir = self._get_masks_dir(sample_index)
-        filename = os.path.join(masks_dir, "semantic_labels", "masks.nii")
-        semantic_labels = io.read_nifti(filename, slice_index)
-        return tv_tensors.Mask(semantic_labels.squeeze(), dtype=torch.int64)  # type: ignore[reportCallIssue]
+        return io.read_nifti(self._get_optimized_masks_file(sample_index), slice_index)
 
     def _load_masks_as_semantic_label(
         self, sample_index: int, slice_index: int | None = None
@@ -227,18 +246,39 @@ class TotalSegmentator2D(base.ImageSegmentation):
             slice_index: Whether to return only a specific slice.
         """
         masks_dir = self._get_masks_dir(sample_index)
-        mask_paths = [os.path.join(masks_dir, f"{label}.nii.gz") for label in self.classes]
+        classes = self._class_mappings.keys() if self._class_mappings else self.classes[1:]
+        mask_paths = [os.path.join(masks_dir, f"{label}.nii.gz") for label in classes]
         binary_masks = [io.read_nifti(path, slice_index) for path in mask_paths]
+
+        if self._class_mappings:
+            mapped_binary_masks = [np.zeros_like(binary_masks[0], dtype=np.bool_)] * len(
+                self.classes[1:]
+            )
+            for original_class, mapped_class in self._class_mappings.items():
+                mapped_index = self.class_to_idx[mapped_class] - 1
+                original_index = list(self._class_mappings.keys()).index(original_class)
+                mapped_binary_masks[mapped_index] = np.logical_or(
+                    mapped_binary_masks[mapped_index], binary_masks[original_index]
+                )
+            binary_masks = mapped_binary_masks
+
         background_mask = np.zeros_like(binary_masks[0])
         return np.argmax([background_mask] + binary_masks, axis=0)
 
     def _export_semantic_label_masks(self) -> None:
         """Exports the segmentation binary masks (one-hot) to semantic labels."""
+        mask_classes_file = os.path.join(f"{self._get_optimized_masks_root()}/classes.txt")
+        if os.path.isfile(mask_classes_file):
+            with open(mask_classes_file, "r") as file:
+                if file.read() != str(self.classes):
+                    raise ValueError(
+                        "Optimized masks hash doesn't match the current classes or mappings."
+                    )
+            return
+
         total_samples = len(self._samples_dirs)
-        masks_dirs = map(self._get_masks_dir, range(total_samples))
         semantic_labels = [
-            (index, os.path.join(directory, "semantic_labels", "masks.nii"))
-            for index, directory in enumerate(masks_dirs)
+            (index, self._get_optimized_masks_file(index)) for index in range(total_samples)
         ]
         to_export = filter(lambda x: not os.path.isfile(x[1]), semantic_labels)
 
@@ -255,6 +295,16 @@ class TotalSegmentator2D(base.ImageSegmentation):
             return_results=False,
         )
 
+        os.makedirs(os.path.dirname(mask_classes_file), exist_ok=True)
+        with open(mask_classes_file, "w") as file:
+            file.write(str(self.classes))
+
+    def _fix_orientation(self, array: npt.NDArray):
+        """Fixes orientation such that table is at the bottom & liver on the left."""
+        array = np.rot90(array)
+        array = np.flip(array, axis=1)
+        return array
+
     def _get_image_path(self, sample_index: int) -> str:
         """Returns the corresponding image path."""
         sample_dir = self._samples_dirs[sample_index]
@@ -265,10 +315,15 @@ class TotalSegmentator2D(base.ImageSegmentation):
         sample_dir = self._samples_dirs[sample_index]
         return os.path.join(self._root, sample_dir, "segmentations")
 
-    def _get_semantic_labels_filename(self, sample_index: int) -> str:
+    def _get_optimized_masks_root(self) -> str:
+        """Returns the directory of the optimized masks."""
+        return os.path.join(self._root, f"processed/masks/{self._classes_hash}")
+
+    def _get_optimized_masks_file(self, sample_index: int) -> str:
         """Returns the semantic label filename."""
-        masks_dir = self._get_masks_dir(sample_index)
-        return os.path.join(masks_dir, "semantic_labels", "masks.nii")
+        return os.path.join(
+            f"{self._get_optimized_masks_root()}/{self._samples_dirs[sample_index]}/masks.nii"
+        )
 
     def _get_number_of_slices_per_sample(self, sample_index: int) -> int:
         """Returns the total amount of slices of a sample."""
@@ -281,7 +336,7 @@ class TotalSegmentator2D(base.ImageSegmentation):
         sample_filenames = [
             filename
             for filename in os.listdir(self._root)
-            if os.path.isdir(os.path.join(self._root, filename))
+            if os.path.isdir(os.path.join(self._root, filename)) and re.match(r"^s\d{4}$", filename)
         ]
         return sorted(sample_filenames)
 

eva/vision/data/datasets/vision.py

@@ -9,7 +9,7 @@ DataSample = TypeVar("DataSample")
 """The data sample type."""
 
 
-class VisionDataset(base.Dataset, abc.ABC, Generic[DataSample]):
+class VisionDataset(base.MapDataset, abc.ABC, Generic[DataSample]):
     """Base dataset class for vision tasks."""
 
     @abc.abstractmethod
@@ -24,20 +24,3 @@ class VisionDataset(base.Dataset, abc.ABC, Generic[DataSample]):
         Returns:
             The filename of the `index`'th data sample.
         """
-
-    @abc.abstractmethod
-    def __getitem__(self, index: int) -> DataSample:
-        """Returns the `index`'th data sample.
-
-        Args:
-            index: The index of the data-sample to select.
-
-        Returns:
-            A data sample and its target.
-        """
-        raise NotImplementedError
-
-    @abc.abstractmethod
-    def __len__(self) -> int:
-        """Returns the total length of the data."""
-        raise NotImplementedError

eva/vision/losses/dice.py

@@ -45,9 +45,6 @@ class DiceLoss(losses.DiceLoss):  # type: ignore
         inputs, targets = _apply_ignore_index(inputs, targets, self.ignore_index)
         targets = _to_one_hot(targets, num_classes=inputs.shape[1])
 
-        if targets.ndim == 3:
-            targets = one_hot(targets[:, None, ...], num_classes=inputs.shape[1])
-
         return super().forward(inputs, targets)
 
 

eva/vision/metrics/__init__.py

@@ -1,11 +1,15 @@
 """Default metric collections API."""
 
 from eva.vision.metrics.defaults.segmentation import MulticlassSegmentationMetrics
+from eva.vision.metrics.segmentation.dice import DiceScore
 from eva.vision.metrics.segmentation.generalized_dice import GeneralizedDiceScore
 from eva.vision.metrics.segmentation.mean_iou import MeanIoU
+from eva.vision.metrics.segmentation.monai_dice import MonaiDiceScore
 
 __all__ = [
-    "MulticlassSegmentationMetrics",
+    "DiceScore",
     "GeneralizedDiceScore",
     "MeanIoU",
+    "MonaiDiceScore",
+    "MulticlassSegmentationMetrics",
 ]

eva/vision/metrics/defaults/segmentation/multiclass.py

@@ -1,7 +1,7 @@
 """Default metric collection for multiclass semantic segmentation tasks."""
 
 from eva.core.metrics import structs
-from eva.vision.metrics.segmentation import generalized_dice, mean_iou
+from eva.vision.metrics import segmentation
 
 
 class MulticlassSegmentationMetrics(structs.MetricCollection):
@@ -26,19 +26,43 @@ class MulticlassSegmentationMetrics(structs.MetricCollection):
             postfix: A string to add after the keys in the output dictionary.
         """
         super().__init__(
-            metrics=[
-                generalized_dice.GeneralizedDiceScore(
+            metrics={
+                "MonaiDiceScore": segmentation.MonaiDiceScore(
                     num_classes=num_classes,
                     include_background=include_background,
-                    weight_type="linear",
                     ignore_index=ignore_index,
+                    ignore_empty=True,
                 ),
-                mean_iou.MeanIoU(
+                "MonaiDiceScore (ignore_empty=False)": segmentation.MonaiDiceScore(
                     num_classes=num_classes,
                     include_background=include_background,
                     ignore_index=ignore_index,
+                    ignore_empty=False,
                 ),
-            ],
+                "DiceScore (micro)": segmentation.DiceScore(
+                    num_classes=num_classes,
+                    include_background=include_background,
+                    average="micro",
+                    ignore_index=ignore_index,
+                ),
+                "DiceScore (macro)": segmentation.DiceScore(
+                    num_classes=num_classes,
+                    include_background=include_background,
+                    average="macro",
+                    ignore_index=ignore_index,
+                ),
+                "DiceScore (weighted)": segmentation.DiceScore(
+                    num_classes=num_classes,
+                    include_background=include_background,
+                    average="weighted",
+                    ignore_index=ignore_index,
+                ),
+                "MeanIoU": segmentation.MeanIoU(
+                    num_classes=num_classes,
+                    include_background=include_background,
+                    ignore_index=ignore_index,
+                ),
+            },
             prefix=prefix,
             postfix=postfix,
         )

eva/vision/metrics/segmentation/__init__.py

@@ -1,9 +1,13 @@
 """Segmentation metrics API."""
 
+from eva.vision.metrics.segmentation.dice import DiceScore
 from eva.vision.metrics.segmentation.generalized_dice import GeneralizedDiceScore
 from eva.vision.metrics.segmentation.mean_iou import MeanIoU
+from eva.vision.metrics.segmentation.monai_dice import MonaiDiceScore
 
 __all__ = [
+    "DiceScore",
+    "MonaiDiceScore",
     "GeneralizedDiceScore",
     "MeanIoU",
 ]