added unit tests

.gitattributes

@@ -9,3 +9,4 @@ tests/eva/assets/**/*.npy filter=lfs diff=lfs merge=lfs -text
 tests/eva/assets/**/*.xml filter=lfs diff=lfs merge=lfs -text
 tests/eva/assets/**/*.mat filter=lfs diff=lfs merge=lfs -text
 tests/eva/assets/**/*.nii filter=lfs diff=lfs merge=lfs -text
+tests/eva/assets/**/*.nii.gz filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -179,6 +179,3 @@ cython_debug/
 
 # numpy data
 *.npy
-
-# NiFti data
-*.nii.gz

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

@@ -24,6 +24,11 @@
 class KiTS23(base.ImageSegmentation):
     """KiTS23 - The 2023 Kidney and Kidney Tumor Segmentation challenge.
 
+    To optimize data loading, the dataset is preprocessed by reorienting the images
+    from IPL to LAS and uncompressing them. The reorientation is necessary, because
+    loading slices from the first dimension is significantly slower than from the last,
+    due to data not being stored in a contiguous manner on disk accross all dimensions.
+
     Webpage: https://kits-challenge.org/kits23/
     """
 
@@ -215,6 +220,8 @@ def _download_dataset(self) -> None:
             _download_case_with_retry(case_id, image_path, segmentation_path)
 
     def _preprocess(self):
+        """Reorienting the images to LPS and uncompressing them."""
+
         def _reorient_and_save(path: Path) -> None:
             relative_path = str(path.relative_to(self._root))
             save_path = os.path.join(self._processed_root, relative_path.rstrip(".gz"))
@@ -227,7 +234,7 @@ def _reorient_and_save(path: Path) -> None:
         multiprocessing.run_with_threads(
             _reorient_and_save,
             [(path,) for path in compressed_paths],
-            num_workers=self._num_workers,
+            num_workers=1,
             progress_desc=">> Preprocessing dataset",
             return_results=False,
         )

src/eva/vision/utils/io/nifti.py

@@ -61,7 +61,7 @@ def reorient(
     """
     orig_ornt = nib.io_orientation(nii.affine)
     targ_ornt = orientations.axcodes2ornt(orientation)
-    if orig_ornt == targ_ornt:
+    if np.all(orig_ornt == targ_ornt):
         return nii
     transform = orientations.ornt_transform(orig_ornt, targ_ornt)
     reoriented_nii = nii.as_reoriented(transform)

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

@@ -0,0 +1,83 @@
+"""KiTS23 dataset tests."""
+
+import os
+import shutil
+from typing import Literal
+from unittest.mock import patch
+
+import pytest
+from torchvision import tv_tensors
+
+from eva.vision.data import datasets
+
+
+@pytest.mark.parametrize(
+    "split, expected_length",
+    [(None, 8)],
+)
+def test_length(kits23_dataset: datasets.KiTS23, expected_length: int) -> None:
+    """Tests the length of the dataset."""
+    assert len(kits23_dataset) == expected_length
+
+
+@pytest.mark.parametrize(
+    "split, index",
+    [
+        (None, 0),
+    ],
+)
+def test_sample(kits23_dataset: datasets.KiTS23, index: int) -> None:
+    """Tests the format of a dataset sample."""
+    # assert data sample is a tuple
+    sample = kits23_dataset[index]
+    assert isinstance(sample, tuple)
+    assert len(sample) == 3
+    # assert the format of the `image` and `mask`
+    image, mask, metadata = sample
+    assert isinstance(image, tv_tensors.Image)
+    assert image.shape == (1, 512, 512)
+    assert isinstance(mask, tv_tensors.Mask)
+    assert mask.shape == (512, 512)
+    assert isinstance(metadata, dict)
+    assert "slice_index" in metadata
+
+
+@pytest.mark.parametrize("split", [None])
+def test_processed_dir_exists(kits23_dataset: datasets.KiTS23) -> None:
+    """Tests the existence of the processed directory."""
+    assert os.path.isdir(kits23_dataset._processed_root)
+
+    for index in ["00036", "00240"]:
+        assert os.path.isfile(
+            os.path.join(kits23_dataset._processed_root, f"case_{index}/master_{index}.nii")
+        )
+        assert os.path.isfile(
+            os.path.join(kits23_dataset._processed_root, f"case_{index}/segmentation.nii")
+        )
+
+
+@pytest.fixture(scope="function")
+def kits23_dataset(split: Literal["train", "val", "test"] | None, assets_path: str):
+    """KiTS23 dataset fixture."""
+    dataset = datasets.KiTS23(
+        root=os.path.join(
+            assets_path,
+            "vision",
+            "datasets",
+            "kits23",
+        ),
+        split=split,
+    )
+    dataset.prepare_data()
+    dataset.configure()
+    yield dataset
+
+    if os.path.isdir(dataset._processed_root):
+        shutil.rmtree(dataset._processed_root)
+
+
+@pytest.fixture(autouse=True)
+def mock_indices():
+    """Mocks the download function to avoid downloading resources when running tests."""
+    with patch.object(datasets.KiTS23, "_get_split_indices", return_value=[36, 240]):
+        yield