Smaller Changes on Task Manager + Contextmanager for SITK IO

Author: Lars-Kraemer

.gitignore

@@ -5,6 +5,7 @@ stats/
 dataset_cfg/
 tests/temp/
 nodes.md
+playground/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

LICENSE

@@ -187,7 +187,7 @@
       same "printed page" as the copyright notice for easier
       identification within third-party archives.
 
-   Copyright {yyyy} {name of copyright owner}
+   Copyright [2025] Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany
 
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.

src/vidata/file_manager/file_manager.py

@@ -42,7 +42,7 @@ def __init__(
         recursive: bool = False,
         lazy_init: bool = False,
     ):
-        self.path = path
+        self.path = Path(path)
         self.file_type = file_type
         self.pattern = pattern
         self.include_names = include_names
@@ -177,7 +177,7 @@ def get_name(self, file: str | int, with_file_type=True) -> str:
         """Legacy alias for :meth:`name_from_path` (kept for backward compatibility)."""
         return self.name_from_path(file, with_file_type)
 
-    def name_from_path(self, file: str | int, include_ext: bool = True) -> str:
+    def name_from_path(self, file: Path | str | int, include_ext: bool = True) -> str:
         """
         Get the relative name of a file (e.g., 'subdir/sample.png').
 
@@ -194,14 +194,16 @@ def name_from_path(self, file: str | int, include_ext: bool = True) -> str:
             Relative file name.
         """
         if isinstance(file, int):
-            file = str(self.files[file])
-        name = (
-            str(Path(file).relative_to(self.path))
-            if not str(self.path).endswith(".json")
-            else str(file)
-        )
-        if not include_ext:
-            name = name.replace(self.file_type, "")
+            file = self.files[file]
+        if not isinstance(file, Path):
+            file = Path(file)
+
+        name_pl = file.relative_to(self.path) if self.path.suffix != ".json" else file
+        name = name_pl.as_posix()
+
+        if not include_ext and name.endswith(self.file_type):
+            name = name[: -len(self.file_type)]
+
         return name
 
     def path_from_name(self, name: str | Path, include_ext=True):
@@ -211,7 +213,7 @@ def path_from_name(self, name: str | Path, include_ext=True):
         rel = Path(name)
         if include_ext and rel.suffix != self.file_type:
             rel = rel.with_suffix(self.file_type)
-        if str(self.path).endswith(".json"):
+        if self.path.suffix == ".json":
             return rel
         else:
             return (self.path / rel).resolve()

src/vidata/io/sitk_io.py

@@ -1,3 +1,5 @@
+import locale
+from contextlib import contextmanager
 from pathlib import Path
 
 import numpy as np
@@ -7,6 +9,20 @@
 from vidata.utils.affine import build_affine
 
 
+@contextmanager
+def temporary_c_locale():
+    # Save current LC_NUMERIC
+    old_locale = locale.setlocale(locale.LC_NUMERIC, None)
+
+    try:
+        # Switch to safe C locale
+        locale.setlocale(locale.LC_NUMERIC, "C")
+        yield
+    finally:
+        # Restore original locale
+        locale.setlocale(locale.LC_NUMERIC, old_locale)
+
+
 @register_writer("image", ".nii.gz", ".nii", ".mha", ".nrrd", backend="sitk")
 @register_writer("mask", ".nii.gz", ".nii", ".mha", ".nrrd", backend="sitk")
 def save_sitk(data: np.ndarray, file: str | Path, metadata: dict | None = None) -> list[str]:
@@ -56,7 +72,9 @@ def load_sitk(file: str | Path) -> tuple[np.ndarray, dict]:
                 - "direction": orientation matrix (np.ndarray)
                 - "affine": computed affine matrix (np.ndarray)
     """
-    image = sitk.ReadImage(file)
+    with temporary_c_locale():
+        image = sitk.ReadImage(file)
+
     array = sitk.GetArrayFromImage(image)
     ndims = len(array.shape)
 

src/vidata/task_manager/multilabel_segmentation_manager.py

@@ -15,12 +15,20 @@ def empty(size: tuple[int, ...], num_classes: int) -> np.ndarray:
         return np.zeros((num_classes, *size), dtype=np.uint8)
 
     @staticmethod
-    def class_ids(data: np.ndarray) -> np.ndarray:
+    def class_ids(
+        data: np.ndarray, return_counts: bool = False
+    ) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
         """
         Return class indices that are present in the mask (i.e., where at least one pixel is non-zero).
         """
-        return np.flatnonzero(data.reshape(data.shape[0], -1).any(axis=1))
-        # return np.where(data.reshape(self.num_classes, -1).any(axis=1))[0]
+        class_ids = np.flatnonzero(data.reshape(data.shape[0], -1).any(axis=1))
+
+        if return_counts:
+            axes = tuple(range(1, data.ndim))
+            counts = np.count_nonzero(data, axis=axes)
+            return class_ids, counts[class_ids]
+
+        return class_ids
 
     @staticmethod
     def class_count(data: np.ndarray, class_id: int) -> int:
@@ -30,10 +38,14 @@ def class_count(data: np.ndarray, class_id: int) -> int:
         return int(np.sum(data[class_id]))
 
     @staticmethod
-    def class_location(data: np.ndarray, class_id: int) -> tuple[np.ndarray, ...]:
+    def class_location(
+        data: np.ndarray, class_id: int, return_mask: bool = False
+    ) -> tuple[np.ndarray, ...] | np.ndarray:
         """
         Return indices where the given class is active (non-zero).
         """
+        if return_mask:
+            return data[class_id]
         return np.where(data[class_id] > 0)
 
     @staticmethod
@@ -45,15 +57,3 @@ def spatial_dims(shape: np.ndarray) -> np.ndarray:
     def has_background():
         """if the task has a dedicated background class --> is class 0 the bg class?"""
         return False
-
-
-if __name__ == "__main__":
-    data = MultiLabelSegmentationManager.random((100, 100), 7)
-    data = MultiLabelSegmentationManager.empty((100, 100), 7)
-    data[1, 0, 5] = 1
-    data[0, 0, 0] = 1
-    data[5, 0, 0] = 1
-    print(np.unique(data))
-    print(data.shape)
-    print(MultiLabelSegmentationManager.class_ids(data))
-    print(MultiLabelSegmentationManager.class_location(data, 1))

src/vidata/task_manager/semantic_segmentation_manager.py

@@ -15,15 +15,23 @@ def empty(size: tuple[int, ...], num_classes: int) -> np.ndarray:
         return np.zeros(size, dtype=np.uint8)
 
     @staticmethod
-    def class_ids(data: np.ndarray) -> np.ndarray:
+    def class_ids(
+        data: np.ndarray, return_counts: bool = False
+    ) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
+        if return_counts:
+            return np.unique(data, return_counts=True)
         return np.unique(data)
 
     @staticmethod
     def class_count(data: np.ndarray, class_id: int) -> int:
         return int(np.sum(data == class_id))
 
     @staticmethod
-    def class_location(data: np.ndarray, class_id: int) -> tuple[np.ndarray, ...]:
+    def class_location(
+        data: np.ndarray, class_id: int, return_mask: bool = False
+    ) -> tuple[np.ndarray, ...] | np.ndarray:
+        if return_mask:
+            return np.asarray(data == class_id)  # data == class_id
         return np.where(data == class_id)
 
     @staticmethod

src/vidata/task_manager/task_manager.py

@@ -20,7 +20,9 @@ def empty(size: tuple[int, ...], num_classes: int) -> np.ndarray:
 
     @staticmethod
     @abstractmethod
-    def class_ids(data: np.ndarray) -> np.ndarray:
+    def class_ids(
+        data: np.ndarray, return_counts: bool = False
+    ) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
         """Return a sorted array of unique class IDs present in the data."""
 
     @staticmethod
@@ -30,7 +32,9 @@ def class_count(data: np.ndarray, class_id: int) -> int:
 
     @staticmethod
     @abstractmethod
-    def class_location(data: np.ndarray, class_id: int) -> tuple[np.ndarray, ...]:
+    def class_location(
+        data: np.ndarray, class_id: int, return_mask: bool = False
+    ) -> tuple[np.ndarray, ...] | np.ndarray:
         """Return the indices where the given class ID occurs."""
 
     @staticmethod