updates

pyproject.toml

@@ -34,7 +34,7 @@ maintainers = [
 ]
 requires-python = ">=3.10"
 dependencies = [
-    "lightning>=2.2.1",
+    "lightning>=2.2.2",
     "jsonargparse[omegaconf]>=4.27.4",
     "tensorboard>=2.16.2",
     "loguru>=0.7.2",

src/eva/vision/data/datasets/segmentation/base.py

@@ -3,7 +3,6 @@
 import abc
 from typing import Any, Callable, Dict, List, Tuple
 
-import numpy as np
 from torchvision import tv_tensors
 from typing_extensions import override
 
@@ -54,7 +53,7 @@ def load_image(self, index: int) -> tv_tensors.Image:
             index: The index of the data sample to load.
 
         Returns:
-            An image torchvision tensor.
+            An image torchvision tensor (channels, height, width).
         """
 
     @abc.abstractmethod

src/eva/vision/data/datasets/segmentation/total_segmentator.py

@@ -138,7 +138,7 @@ def load_masks(self, index: int) -> tv_tensors.Mask:
         list_of_mask_arrays = [io.read_nifti_slice(path, slice_index) for path in mask_paths]
         masks = np.concatenate(list_of_mask_arrays, axis=-1)
         return tv_tensors.Mask(masks.transpose(2, 0, 1))
-    
+
     def _get_masks_dir(self, index: int) -> str:
         """Returns the directory of the corresponding masks."""
         sample_dir = self._get_sample_dir(index)

tests/eva/vision/data/datasets/segmentation/test_total_segmentator.py

@@ -3,8 +3,8 @@
 import os
 from typing import Literal
 
-import numpy as np
 import pytest
+from torchvision import tv_tensors
 
 from eva.vision.data import datasets
 
@@ -35,10 +35,10 @@ def test_sample(total_segmentator_dataset: datasets.TotalSegmentator2D, index: i
     assert len(sample) == 2
     # assert the format of the `image` and `mask`
     image, mask = sample
-    assert isinstance(image, np.ndarray)
-    assert image.shape == (16, 16, 3)
-    assert isinstance(mask, np.ndarray)
-    assert mask.shape == (16, 16, 3)
+    assert isinstance(image, tv_tensors.Image)
+    assert image.shape == (3, 16, 16)
+    assert isinstance(mask, tv_tensors.Mask)
+    assert mask.shape == (3, 16, 16)
 
 
 @pytest.fixture(scope="function")