kaiko-eva 0.1.1__py3-none-any.whl → 0.1.5__py3-none-any.whl

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

@@ -5,12 +5,14 @@ import glob
 import os
 from typing import Any, Callable, Dict, List, Literal, Tuple
 
+import numpy as np
+import numpy.typing as npt
 import torch
 from torchvision import tv_tensors
 from typing_extensions import override
 
 from eva.core import utils
-from eva.vision.data.datasets import _utils as data_utils
+from eva.core.data import splitting
 from eva.vision.data.datasets import _validators
 from eva.vision.data.datasets.segmentation import base
 from eva.vision.utils import io
@@ -20,22 +22,23 @@ class LiTS(base.ImageSegmentation):
     """LiTS - Liver Tumor Segmentation Challenge.
 
     Webpage: https://competitions.codalab.org/competitions/17094
-
-    For the splits we follow: https://arxiv.org/pdf/2010.01663v2
     """
 
-    _train_index_ranges: List[Tuple[int, int]] = [(0, 102)]
-    _val_index_ranges: List[Tuple[int, int]] = [(102, 117)]
-    _test_index_ranges: List[Tuple[int, int]] = [(117, 131)]
+    _train_ratio: float = 0.7
+    _val_ratio: float = 0.15
+    _test_ratio: float = 0.15
     """Index ranges per split."""
 
+    _fix_orientation: bool = True
+    """Whether to fix the orientation of the images to match the default for radiologists."""
+
     _sample_every_n_slices: int | None = None
     """The amount of slices to sub-sample per 3D CT scan image."""
 
     _expected_dataset_lengths: Dict[str | None, int] = {
-        "train": 39307,
-        "val": 12045,
-        "test": 7286,
+        "train": 38686,
+        "val": 11192,
+        "test": 8760,
         None: 58638,
     }
     """Dataset version and split to the expected size."""
@@ -51,6 +54,7 @@ class LiTS(base.ImageSegmentation):
         root: str,
         split: Literal["train", "val", "test"] | None = None,
         transforms: Callable | None = None,
+        seed: int = 8,
     ) -> None:
         """Initialize dataset.
 
@@ -60,12 +64,13 @@ class LiTS(base.ImageSegmentation):
             split: Dataset split to use.
             transforms: A function/transforms that takes in an image and a target
                 mask and returns the transformed versions of both.
+            seed: Seed used for generating the dataset splits.
         """
         super().__init__(transforms=transforms)
 
         self._root = root
         self._split = split
-
+        self._seed = seed
         self._indices: List[Tuple[int, int]] = []
 
     @property
@@ -90,10 +95,12 @@ class LiTS(base.ImageSegmentation):
 
     @override
     def validate(self) -> None:
-        if len(self._volume_files) != len(self._segmentation_files):
-            raise ValueError(
-                "The number of volume files does not match the number of the segmentation ones."
-            )
+        for i in range(len(self._volume_files)):
+            seg_path = self._segmentation_file(i)
+            if not os.path.exists(seg_path):
+                raise FileNotFoundError(
+                    f"Segmentation file {seg_path} not found for volume {self._volume_files[i]}."
+                )
 
         _validators.check_dataset_integrity(
             self,
@@ -107,15 +114,27 @@ class LiTS(base.ImageSegmentation):
         sample_index, slice_index = self._indices[index]
         volume_path = self._volume_files[sample_index]
         image_array = io.read_nifti(volume_path, slice_index)
+        if self._fix_orientation:
+            image_array = self._orientation(image_array, sample_index)
         return tv_tensors.Image(image_array.transpose(2, 0, 1))
 
     @override
     def load_mask(self, index: int) -> tv_tensors.Mask:
         sample_index, slice_index = self._indices[index]
-        segmentation_path = self._segmentation_files[sample_index]
+        segmentation_path = self._segmentation_file(sample_index)
         semantic_labels = io.read_nifti(segmentation_path, slice_index)
+        if self._fix_orientation:
+            semantic_labels = self._orientation(semantic_labels, sample_index)
         return tv_tensors.Mask(semantic_labels.squeeze(), dtype=torch.int64)  # type: ignore[reportCallIssue]
 
+    def _orientation(self, array: npt.NDArray, sample_index: int) -> npt.NDArray:
+        volume_path = self._volume_files[sample_index]
+        orientation = io.fetch_nifti_axis_direction_code(volume_path)
+        array = np.rot90(array, axes=(0, 1))
+        if orientation == "LPS":
+            array = np.flip(array, axis=0)
+        return array.copy()
+
     @override
     def load_metadata(self, index: int) -> Dict[str, Any]:
         _, slice_index = self._indices[index]
@@ -137,11 +156,10 @@ class LiTS(base.ImageSegmentation):
         files = glob.glob(files_pattern, recursive=True)
         return utils.numeric_sort(files)
 
-    @functools.cached_property
-    def _segmentation_files(self) -> List[str]:
-        files_pattern = os.path.join(self._root, "**", "segmentation-*.nii")
-        files = glob.glob(files_pattern, recursive=True)
-        return utils.numeric_sort(files)
+    def _segmentation_file(self, index: int) -> str:
+        volume_file_path = self._volume_files[index]
+        segmentation_file = os.path.basename(volume_file_path).replace("volume", "segmentation")
+        return os.path.join(os.path.dirname(volume_file_path), segmentation_file)
 
     def _create_indices(self) -> List[Tuple[int, int]]:
         """Builds the dataset indices for the specified split.
@@ -161,17 +179,19 @@ class LiTS(base.ImageSegmentation):
 
     def _get_split_indices(self) -> List[int]:
         """Returns the sample indices for the specified dataset split."""
-        split_index_ranges = {
-            "train": self._train_index_ranges,
-            "val": self._val_index_ranges,
-            "test": self._test_index_ranges,
-            None: [(0, len(self._volume_files))],
+        indices = list(range(len(self._volume_files)))
+        train_indices, val_indices, test_indices = splitting.random_split(
+            indices, self._train_ratio, self._val_ratio, self._test_ratio, seed=self._seed
+        )
+        split_indices_dict = {
+            "train": train_indices,
+            "val": val_indices,
+            "test": test_indices,
+            None: indices,
         }
-        index_ranges = split_index_ranges.get(self._split)
-        if index_ranges is None:
+        if self._split not in split_indices_dict:
             raise ValueError("Invalid data split. Use 'train', 'val', 'test' or `None`.")
-
-        return data_utils.ranges_to_indices(index_ranges)
+        return list(split_indices_dict[self._split])
 
     def _print_license(self) -> None:
         """Prints the dataset license."""

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

@@ -0,0 +1,93 @@
+"""Balanced LiTS dataset."""
+
+from typing import Callable, Dict, List, Literal, Tuple
+
+import numpy as np
+from typing_extensions import override
+
+from eva.vision.data.datasets.segmentation import lits
+from eva.vision.utils import io
+
+
+class LiTSBalanced(lits.LiTS):
+    """Balanced version of the LiTS - Liver Tumor Segmentation Challenge dataset.
+
+    For each volume in the dataset, we sample the same number of slices where
+    only the liver and where both liver and tumor are present.
+
+    Webpage: https://competitions.codalab.org/competitions/17094
+
+    For the splits we follow: https://arxiv.org/pdf/2010.01663v2
+    """
+
+    _expected_dataset_lengths: Dict[str | None, int] = {
+        "train": 5514,
+        "val": 1332,
+        "test": 1530,
+        None: 8376,
+    }
+    """Dataset version and split to the expected size."""
+
+    def __init__(
+        self,
+        root: str,
+        split: Literal["train", "val", "test"] | None = None,
+        transforms: Callable | None = None,
+        seed: int = 8,
+    ) -> None:
+        """Initialize dataset.
+
+        Args:
+            root: Path to the root directory of the dataset. The dataset will
+                be downloaded and extracted here, if it does not already exist.
+            split: Dataset split to use.
+            transforms: A function/transforms that takes in an image and a target
+                mask and returns the transformed versions of both.
+            seed: Seed used for generating the dataset splits and sampling of the slices.
+        """
+        super().__init__(root=root, split=split, transforms=transforms, seed=seed)
+
+    @override
+    def _create_indices(self) -> List[Tuple[int, int]]:
+        """Builds the dataset indices for the specified split.
+
+        Returns:
+            A list of tuples, where the first value indicates the
+            sample index which the second its corresponding slice
+            index.
+        """
+        split_indices = set(self._get_split_indices())
+        indices: List[Tuple[int, int]] = []
+        random_generator = np.random.default_rng(seed=self._seed)
+
+        for sample_idx in range(len(self._volume_files)):
+            if sample_idx not in split_indices:
+                continue
+
+            segmentation = io.read_nifti(self._segmentation_file(sample_idx))
+            tumor_filter = segmentation == 2
+            tumor_slice_filter = tumor_filter.sum(axis=(0, 1)) > 0
+
+            if tumor_filter.sum() == 0:
+                continue
+
+            liver_filter = segmentation == 1
+            liver_slice_filter = liver_filter.sum(axis=(0, 1)) > 0
+
+            liver_and_tumor_filter = liver_slice_filter & tumor_slice_filter
+            liver_only_filter = liver_slice_filter & ~tumor_slice_filter
+
+            n_slice_samples = min(liver_and_tumor_filter.sum(), liver_only_filter.sum())
+            tumor_indices = list(np.where(liver_and_tumor_filter)[0])
+            tumor_indices = list(
+                random_generator.choice(tumor_indices, size=n_slice_samples, replace=False)
+            )
+
+            liver_indices = list(np.where(liver_only_filter)[0])
+            liver_indices = list(
+                random_generator.choice(liver_indices, size=n_slice_samples, replace=False)
+            )
+
+            indices.extend([(sample_idx, slice_idx) for slice_idx in tumor_indices + liver_indices])
+
+        return list(indices)

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

@@ -10,12 +10,12 @@ import imagesize
 import numpy as np
 import numpy.typing as npt
 import torch
-import tqdm
 from skimage import draw
 from torchvision import tv_tensors
 from torchvision.datasets import utils
 from typing_extensions import override
 
+from eva.core.utils.progress_bar import tqdm
 from eva.vision.data.datasets import _validators, structs
 from eva.vision.data.datasets.segmentation import base
 from eva.vision.utils import io
@@ -84,7 +84,7 @@ class MoNuSAC(base.ImageSegmentation):
     @property
     @override
     def classes(self) -> List[str]:
-        return ["Epithelial", "Lymphocyte", "Neutrophil", "Macrophage"]
+        return ["Background", "Epithelial", "Lymphocyte", "Neutrophil", "Macrophage", "Ambiguous"]
 
     @functools.cached_property
     @override
@@ -107,8 +107,8 @@ class MoNuSAC(base.ImageSegmentation):
         _validators.check_dataset_integrity(
             self,
             length=self._expected_dataset_lengths.get(self._split, 0),
-            n_classes=4,
-            first_and_last_labels=("Epithelial", "Macrophage"),
+            n_classes=6,
+            first_and_last_labels=("Background", "Ambiguous"),
         )
 
     @override
@@ -161,7 +161,7 @@ class MoNuSAC(base.ImageSegmentation):
             for index, filename in enumerate(self._image_files)
         ]
         to_export = filter(lambda x: not os.path.isfile(x[1]), mask_files)
-        for sample_index, filename in tqdm.tqdm(
+        for sample_index, filename in tqdm(
             list(to_export),
             desc=">> Exporting semantic masks",
             leave=False,
@@ -199,9 +199,9 @@ class MoNuSAC(base.ImageSegmentation):
         semantic_labels = np.zeros((height, width), "uint8")  # type: ignore[reportCallIssue]
         for level in range(len(root)):
             label = [item.attrib["Name"] for item in root[level][0]][0]
-            class_id = self.class_to_idx.get(label, 254) + 1
+            class_id = self.class_to_idx.get(label, self.class_to_idx["Ambiguous"])
             # for the test dataset an additional class 'Ambiguous' was added for
-            # difficult regions with fuzzy boundaries - we return it as 255
+            # difficult regions with fuzzy boundaries
             regions = [item for child in root[level] for item in child if item.tag == "Region"]
             for region in regions:
                 vertices = np.array(

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

@@ -3,16 +3,18 @@
 import functools
 import os
 from glob import glob
+from pathlib import Path
 from typing import Any, Callable, Dict, List, Literal, Tuple
 
 import numpy as np
 import numpy.typing as npt
 import torch
-import tqdm
 from torchvision import tv_tensors
 from torchvision.datasets import utils
 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.utils import io
@@ -65,6 +67,8 @@ class TotalSegmentator2D(base.ImageSegmentation):
         download: bool = False,
         classes: List[str] | None = None,
         optimize_mask_loading: bool = True,
+        decompress: bool = True,
+        num_workers: int = 10,
         transforms: Callable | None = None,
     ) -> None:
         """Initialize dataset.
@@ -85,8 +89,15 @@ class TotalSegmentator2D(base.ImageSegmentation):
                 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
                 it on disk.
+            decompress: Whether to decompress the ct.nii.gz files when preparing the data.
+                The label masks won't be decompressed, but when enabling optimize_mask_loading
+                it will export the semantic label masks to a single file in uncompressed .nii
+                format.
+            num_workers: The number of workers to use for optimizing the masks &
+                decompressing the .gz files.
             transforms: A function/transforms that takes in an image and a target
                 mask and returns the transformed versions of both.
+
         """
         super().__init__(transforms=transforms)
 
@@ -96,6 +107,8 @@ class TotalSegmentator2D(base.ImageSegmentation):
         self._download = download
         self._classes = classes
         self._optimize_mask_loading = optimize_mask_loading
+        self._decompress = decompress
+        self._num_workers = num_workers
 
         if self._optimize_mask_loading and self._classes is not None:
             raise ValueError(
@@ -128,23 +141,29 @@ class TotalSegmentator2D(base.ImageSegmentation):
     def class_to_idx(self) -> Dict[str, int]:
         return {label: index for index, label in enumerate(self.classes)}
 
+    @property
+    def _file_suffix(self) -> str:
+        return "nii" if self._decompress else "nii.gz"
+
     @override
-    def filename(self, index: int, segmented: bool = True) -> str:
+    def filename(self, index: int) -> str:
         sample_idx, _ = self._indices[index]
         sample_dir = self._samples_dirs[sample_idx]
-        return os.path.join(sample_dir, "ct.nii.gz")
+        return os.path.join(sample_dir, f"ct.{self._file_suffix}")
 
     @override
     def prepare_data(self) -> None:
         if self._download:
             self._download_dataset()
+        if self._decompress:
+            self._decompress_files()
+        self._samples_dirs = self._fetch_samples_dirs()
+        if self._optimize_mask_loading:
+            self._export_semantic_label_masks()
 
     @override
     def configure(self) -> None:
-        self._samples_dirs = self._fetch_samples_dirs()
         self._indices = self._create_indices()
-        if self._optimize_mask_loading:
-            self._export_semantic_label_masks()
 
     @override
     def validate(self) -> None:
@@ -186,16 +205,15 @@ class TotalSegmentator2D(base.ImageSegmentation):
         return {"slice_index": slice_index}
 
     def _load_mask(self, index: int) -> tv_tensors.Mask:
-        """Loads and builds the segmentation mask from NifTi files."""
         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, dtype=torch.int64)  # type: ignore[reportCallIssue]
+        return tv_tensors.Mask(semantic_labels.squeeze(), dtype=torch.int64)  # type: ignore[reportCallIssue]
 
     def _load_semantic_label_mask(self, index: int) -> tv_tensors.Mask:
         """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.gz")
+        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]
 
@@ -209,7 +227,7 @@ 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, label + ".nii.gz") for label in self.classes]
+        mask_paths = [os.path.join(masks_dir, f"{label}.nii.gz") for label in self.classes]
         binary_masks = [io.read_nifti(path, slice_index) for path in mask_paths]
         background_mask = np.zeros_like(binary_masks[0])
         return np.argmax([background_mask] + binary_masks, axis=0)
@@ -219,24 +237,28 @@ class TotalSegmentator2D(base.ImageSegmentation):
         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.gz"))
+            (index, os.path.join(directory, "semantic_labels", "masks.nii"))
             for index, directory in enumerate(masks_dirs)
         ]
         to_export = filter(lambda x: not os.path.isfile(x[1]), semantic_labels)
 
-        for sample_index, filename in tqdm.tqdm(
-            list(to_export),
-            desc=">> Exporting optimized semantic masks",
-            leave=False,
-        ):
+        def _process_mask(sample_index: Any, filename: str) -> None:
             semantic_labels = self._load_masks_as_semantic_label(sample_index)
             os.makedirs(os.path.dirname(filename), exist_ok=True)
             io.save_array_as_nifti(semantic_labels, filename)
 
+        multiprocessing.run_with_threads(
+            _process_mask,
+            list(to_export),
+            num_workers=self._num_workers,
+            progress_desc=">> Exporting optimized semantic mask",
+            return_results=False,
+        )
+
     def _get_image_path(self, sample_index: int) -> str:
         """Returns the corresponding image path."""
         sample_dir = self._samples_dirs[sample_index]
-        return os.path.join(self._root, sample_dir, "ct.nii.gz")
+        return os.path.join(self._root, sample_dir, f"ct.{self._file_suffix}")
 
     def _get_masks_dir(self, sample_index: int) -> str:
         """Returns the directory of the corresponding masks."""
@@ -246,7 +268,7 @@ class TotalSegmentator2D(base.ImageSegmentation):
     def _get_semantic_labels_filename(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.gz")
+        return os.path.join(masks_dir, "semantic_labels", "masks.nii")
 
     def _get_number_of_slices_per_sample(self, sample_index: int) -> int:
         """Returns the total amount of slices of a sample."""
@@ -320,6 +342,16 @@ class TotalSegmentator2D(base.ImageSegmentation):
                 remove_finished=True,
             )
 
+    def _decompress_files(self) -> None:
+        compressed_paths = Path(self._root).rglob("*/ct.nii.gz")
+        multiprocessing.run_with_threads(
+            core_io.gunzip_file,
+            [(str(path),) for path in compressed_paths],
+            num_workers=self._num_workers,
+            progress_desc=">> Decompressing .gz files",
+            return_results=False,
+        )
+
     def _print_license(self) -> None:
         """Prints the dataset license."""
         print(f"Dataset license: {self._license}")

eva/vision/data/datasets/wsi.py

@@ -2,8 +2,9 @@
 
 import bisect
 import os
-from typing import Callable, List
+from typing import Any, Callable, Dict, List
 
+import pandas as pd
 from loguru import logger
 from torch.utils.data import dataset as torch_datasets
 from torchvision import tv_tensors
@@ -85,6 +86,17 @@ class WsiDataset(vision.VisionDataset):
         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)
@@ -105,6 +117,7 @@ class MultiWsiDataset(vision.VisionDataset):
         overwrite_mpp: float | None = None,
         backend: str = "openslide",
         image_transforms: Callable | None = None,
+        coords_path: str | None = None,
     ):
         """Initializes a new dataset instance.
 
@@ -118,6 +131,7 @@ class MultiWsiDataset(vision.VisionDataset):
             sampler: The sampler to use for sampling patch coordinates.
             backend: The backend to use for reading the whole-slide images.
             image_transforms: Transforms to apply to the extracted image patches.
+            coords_path: File path to save the patch coordinates as .csv.
         """
         super().__init__()
 
@@ -130,6 +144,7 @@ class MultiWsiDataset(vision.VisionDataset):
         self._sampler = sampler
         self._backend = backend
         self._image_transforms = image_transforms
+        self._coords_path = coords_path
 
         self._concat_dataset: torch_datasets.ConcatDataset
 
@@ -146,6 +161,7 @@ class MultiWsiDataset(vision.VisionDataset):
     @override
     def configure(self) -> None:
         self._concat_dataset = torch_datasets.ConcatDataset(datasets=self._load_datasets())
+        self._save_coords_to_file()
 
     @override
     def __len__(self) -> int:
@@ -159,6 +175,12 @@ class MultiWsiDataset(vision.VisionDataset):
     def filename(self, index: int) -> str:
         return os.path.basename(self._file_paths[self._get_dataset_idx(index)])
 
+    def load_metadata(self, index: int) -> Dict[str, Any]:
+        """Loads the metadata for the patch at the specified index."""
+        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_datasets(self) -> list[WsiDataset]:
         logger.info(f"Initializing dataset with {len(self._file_paths)} WSIs ...")
         wsi_datasets = []
@@ -185,3 +207,17 @@ class MultiWsiDataset(vision.VisionDataset):
 
     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]
+
+    def _save_coords_to_file(self):
+        if self._coords_path is not None:
+            coords = [
+                {"file": self._file_paths[i]} | dataset._coords.to_dict()
+                for i, dataset in enumerate(self.datasets)
+            ]
+            os.makedirs(os.path.abspath(os.path.join(self._coords_path, os.pardir)), exist_ok=True)
+            pd.DataFrame(coords).to_csv(self._coords_path, index=False)
+            logger.info(f"Saved patch coordinates to: {self._coords_path}")

eva/vision/data/wsi/patching/coordinates.py

@@ -2,7 +2,7 @@
 
 import dataclasses
 import functools
-from typing import List, Tuple
+from typing import Any, Dict, List, Tuple
 
 from eva.vision.data.wsi import backends
 from eva.vision.data.wsi.patching import samplers
@@ -75,6 +75,14 @@ class PatchCoordinates:
 
         return cls(x_y, scaled_width, scaled_height, level_idx, sample_args.get("mask"))
 
+    def to_dict(self, include_keys: List[str] | None = None) -> Dict[str, Any]:
+        """Convert the coordinates to a dictionary."""
+        include_keys = include_keys or ["x_y", "width", "height", "level_idx"]
+        coord_dict = dataclasses.asdict(self)
+        if include_keys:
+            coord_dict = {key: coord_dict[key] for key in include_keys}
+        return coord_dict
+
 
 @functools.lru_cache(LRU_CACHE_SIZE)
 def get_cached_coords(

eva/vision/data/wsi/patching/samplers/_utils.py

@@ -1,14 +1,8 @@
-import random
 from typing import Tuple
 
 import numpy as np
 
 
-def set_seed(seed: int) -> None:
-    random.seed(seed)
-    np.random.seed(seed)
-
-
 def get_grid_coords_and_indices(
     layer_shape: Tuple[int, int],
     width: int,
@@ -33,8 +27,8 @@ def get_grid_coords_and_indices(
 
     indices = list(range(len(x_y)))
     if shuffle:
-        set_seed(seed)
-        np.random.shuffle(indices)
+        random_generator = np.random.default_rng(seed)
+        random_generator.shuffle(indices)
     return x_y, indices
 
 

eva/vision/data/wsi/patching/samplers/random.py

@@ -18,6 +18,7 @@ class RandomSampler(base.Sampler):
         """Initializes the sampler."""
         self.seed = seed
         self.n_samples = n_samples
+        self.random_generator = random.Random(seed)  # nosec
 
     def sample(
         self,
@@ -33,9 +34,10 @@ class RandomSampler(base.Sampler):
             layer_shape: The shape of the layer.
         """
         _utils.validate_dimensions(width, height, layer_shape)
-        _utils.set_seed(self.seed)
 
         x_max, y_max = layer_shape[0], layer_shape[1]
         for _ in range(self.n_samples):
-            x, y = random.randint(0, x_max - width), random.randint(0, y_max - height)  # nosec
+            x, y = self.random_generator.randint(0, x_max - width), self.random_generator.randint(
+                0, y_max - height
+            )
             yield x, y

eva/vision/losses/__init__.py

@@ -1,5 +1,5 @@
 """Loss functions API."""
 
-from eva.vision.losses.dice import DiceLoss
+from eva.vision.losses.dice import DiceCELoss, DiceLoss
 
-__all__ = ["DiceLoss"]
+__all__ = ["DiceLoss", "DiceCELoss"]