nrrd_io

Author: Lars-Kraemer

pyproject.toml

@@ -37,7 +37,8 @@ dependencies = [
     "pyyaml",
     "seaborn",
     "scikit-learn",
-    "opencv-python-headless"
+    "opencv-python-headless",
+    "pynrrd",
 ]
 
 [project.optional-dependencies]

src/vidata/io/__init__.py

@@ -7,6 +7,7 @@
 from .tif_io import load_tif, save_tif
 from .blosc2_io import load_blosc2, load_blosc2pkl, save_blosc2, save_blosc2pkl
 from .numpy_io import load_npy, load_npz, save_npy, save_npz
+from .nrrd_io import load_nrrd, save_nrrd
 from .json_io import load_json, save_json, load_jsongz, save_jsongz
 from .pickle_io import load_pickle, save_pickle
 from .txt_io import load_txt, save_txt
@@ -19,6 +20,8 @@
     "save_nib",
     "load_nibRO",
     "save_nibRO",
+    "load_nrrd",
+    "save_nrrd",
     "load_blosc2",
     "save_blosc2",
     "load_blosc2pkl",

src/vidata/io/nrrd_io.py

@@ -0,0 +1,116 @@
+from pathlib import Path
+
+import nrrd
+import numpy as np
+
+from vidata.registry import register_loader, register_writer
+from vidata.utils.affine import build_affine
+
+
+@register_writer("image", ".nrrd", backend="nrrd")
+@register_writer("mask", ".nrrd", backend="nrrd")
+def save_nrrd(data: np.ndarray, file: str | Path, metadata: dict | None = None) -> list[str]:
+    """Save a NumPy array as a nrrd using pynrrd.
+
+    Args:
+        data (np.ndarray): Image data array (z, y, x) or (y, x).
+        file (Union[str, Path]): Output file path (.nrrd).
+        metadata (Optional[dict]): Optional metadata dictionary containing:
+            - "spacing" (list or np.ndarray): Physical spacing per axis.
+            - "origin" (list or np.ndarray): Origin of the image.
+            - "direction" (list or np.ndarray): Flattened or matrix-form direction.
+    """
+    ndims = data.ndim
+
+    if ndims == 3:
+        data = data.transpose(2, 1, 0)  # (X,Y,Z) → (Z,Y,X)
+    elif ndims == 2:
+        data = data.transpose(1, 0)
+
+    # if metadata is not None:
+    if metadata is not None and "spacing" in metadata:
+        spacing = np.array(metadata["spacing"][::-1])
+    else:
+        spacing = np.ones(ndims)
+
+    if metadata is not None and "origin" in metadata:
+        origin = np.array(metadata["origin"][::-1])
+    else:
+        origin = np.zeros(ndims)
+
+    if metadata is not None and "direction" in metadata:
+        direction = np.array(metadata["direction"]).flatten()[::-1].reshape(ndims, ndims)
+    else:
+        direction = np.eye(ndims)
+
+    space_dirs = direction * spacing[:, None]
+
+    header = {
+        "type": "short",
+        "dimension": ndims,
+        "space origin": origin.tolist(),
+        "space directions": space_dirs.tolist(),
+        "encoding": "gzip",  # compressed NRRD
+    }
+    nrrd.write(str(file), data, header)
+    return [str(file)]
+
+
+@register_loader("image", ".nrrd", backend="nrrd")
+@register_loader("mask", ".nrrd", backend="nrrd")
+def load_nrrd(file: str | Path) -> tuple[np.ndarray, dict]:
+    """Load a nrrd file using pynrrd and return data and metadata. This function is consistent with
+    vidata.io.sitk_io.load_sitk
+
+    Args:
+        file (Union[str, Path]): Path to the medical image file (.nrrd).
+
+    Returns:
+        tuple[np.ndarray, dict[str, np.ndarray]]: A tuple containing:
+            - Image data as a NumPy array (z, y, x or y, x).
+            - Metadata dictionary with:
+                - "spacing": voxel spacing (np.ndarray)
+                - "origin": image origin (np.ndarray)
+                - "direction": orientation matrix (np.ndarray)
+                - "affine": computed affine matrix (np.ndarray)
+    """
+    array, header = nrrd.read(str(file))
+    ndims = array.ndim
+
+    # Convert the array to z,y,x axis ordering (numpy)
+    if ndims == 3:
+        array = array.transpose(2, 1, 0)  # (X,Y,Z) → (Z,Y,X)
+    elif ndims == 2:
+        array = array.transpose(1, 0)  # (X,Y) → (Y,X)
+
+    # Extract the metadata
+    origin = np.array(header["space origin"]) if "space origin" in header else np.zeros(ndims)
+    spacing = (
+        np.linalg.norm(np.array(header["space directions"]), axis=1)
+        if "space directions" in header
+        else np.ones(ndims)
+    )
+
+    if "space directions" in header:
+        norms = np.linalg.norm(header["space directions"], axis=1, keepdims=True)
+        direction = header["space directions"] / (norms + 1e-12)
+    else:
+        direction = np.eye(ndims)
+    direction = direction.flatten()  # Flattened SITK style
+
+    # Convert Metadata from (X,Y,Z) → (Z,Y,X)
+    spacing = spacing[::-1]
+    origin = origin[::-1]
+    direction = np.array(direction[::-1]).reshape(ndims, ndims)
+
+    affine = build_affine(ndims, spacing, origin, direction)
+
+    metadata = {
+        "spacing": spacing,
+        "origin": origin,
+        "direction": direction,
+        "affine": affine,
+        "header": header,
+    }
+
+    return array, metadata

tests/io/test_complex_io.py

@@ -164,15 +164,68 @@ def test_nibRO(dtype, file_ending, size, metadata, tmp_path):
 
     save_nibRO(data, file, metadata)
     data_l, metadata_l = load_nibRO(file)
-    print("D1", data)
-    print("D2", data_l)
-    print(metadata_l)
     assert np.array_equal(data, data_l)
     if metadata is not None:
         for key in metadata:
             assert np.all(metadata_l[key] == metadata[key])
 
 
+@pytest.mark.parametrize("dtype", ["float", "int"])
+@pytest.mark.parametrize(
+    "size, metadata",
+    [
+        ((100, 100, 100), None),
+        ((100, 100), None),
+        (
+            (110, 90, 100),
+            {
+                "spacing": [1, 0.5, 0.25],
+                "origin": [10, 20, 30],
+                "direction": [
+                    [1.0, 0.0, 0.0],
+                    [0.0, -1.0, 0.0],
+                    [0.0, 0.0, -1.0],
+                ],
+            },
+        ),
+        (
+            (110, 90, 100),
+            {
+                "spacing": np.array([1, 0.5, 0.25]),
+                "origin": np.array([10, 20, 30]),
+                "direction": np.array([[1.0, 0.0, 0.0], [0.0, -1.0, 0.0], [0.0, 0.0, -1.0]]),
+            },
+        ),
+        (
+            (110, 90),
+            {
+                "spacing": [1, 0.5],
+                "origin": [10, 20],
+                "direction": [[1.0, 0.0], [0.0, -1.0]],
+            },
+        ),
+        ((120, 100, 80), {"spacing": [1, 0.5, 0.25]}),
+        ((120, 100), {"spacing": [1, 0.5]}),
+    ],
+)
+def test_nrrd(dtype, size, metadata, tmp_path):
+    from vidata.io import load_nrrd, save_nrrd
+
+    file_ending = ".nrrd"
+
+    file = Path(tmp_path).joinpath(f"nrrd_{dtype}{file_ending}")
+
+    data = dummy_data(size, dtype)
+
+    save_nrrd(data, file, metadata)
+    data_l, metadata_l = load_nrrd(file)
+
+    assert np.array_equal(data, data_l)
+    if metadata is not None:
+        for key in metadata:
+            assert np.allclose(metadata_l[key], metadata[key])
+
+
 @pytest.mark.parametrize("dtype", ["float", "int"])
 @pytest.mark.parametrize(
     "size, metadata",