keep axis same

kaiko-ai · Apr 22, 2024 · 5f8ba05 · 5f8ba05
1 parent c5ce41e
commit 5f8ba05
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/eva/vision/data/datasets/segmentation/total_segmentator.py b/src/eva/vision/data/datasets/segmentation/total_segmentator.py
@@ -136,7 +136,7 @@ def load_masks(self, index: int) -> tv_tensors.Mask:
         slice_index = self._get_sample_slice_index(index)
         mask_paths = (os.path.join(masks_dir, label + ".nii.gz") for label in self.classes)
         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)
+        masks = np.concatenate(list_of_mask_arrays, axis=2)
         return tv_tensors.Mask(masks.transpose(2, 0, 1))
 
     def _get_masks_dir(self, index: int) -> str: