Repository: DIAGNijmegen/rse-panimg (language: python)

Issue:
`nifti` and `nrrd` readers are not doing any header validation
See https://github.com/comic/grand-challenge.org/issues/1376#issuecomment-747460066


Discussion / hints:
This issue seems to be resolved, at least I haven't come across this problem anymore for a while.

I noticed though that the nifti and nrrd readers are not doing any header validation at the moment, so the issue of having that not harmonized across different builders remains.

Fix patch:
diff --git a/panimg/image_builders/__init__.py b/panimg/image_builders/__init__.py
index 2462c5e..dd6d479 100644
--- a/panimg/image_builders/__init__.py
+++ b/panimg/image_builders/__init__.py
@@ -3,8 +3,8 @@
 from panimg.image_builders.dicom import image_builder_dicom
 from panimg.image_builders.fallback import image_builder_fallback
 from panimg.image_builders.metaio_mhd_mha import image_builder_mhd
-from panimg.image_builders.nifti import image_builder_nifti
-from panimg.image_builders.nrrd import image_builder_nrrd
+from panimg.image_builders.metaio_nifti import image_builder_nifti
+from panimg.image_builders.metaio_nrrd import image_builder_nrrd
 from panimg.image_builders.oct import image_builder_oct
 from panimg.image_builders.tiff import image_builder_tiff
 from panimg.types import ImageBuilder
diff --git a/panimg/image_builders/nifti.py b/panimg/image_builders/metaio_nifti.py
similarity index 71%
rename from panimg/image_builders/nifti.py
rename to panimg/image_builders/metaio_nifti.py
index 6e74e3d..defc9ae 100644
--- a/panimg/image_builders/nifti.py
+++ b/panimg/image_builders/metaio_nifti.py
@@ -2,9 +2,8 @@
 from pathlib import Path
 from typing import DefaultDict, Iterator, List, Set
 
-import SimpleITK
-
-from panimg.exceptions import UnconsumedFilesException
+from panimg.exceptions import UnconsumedFilesException, ValidationError
+from panimg.image_builders.metaio_utils import load_sitk_image
 from panimg.models import SimpleITKImage
 
 
@@ -37,22 +36,20 @@ def image_builder_nifti(*, files: Set[Path]) -> Iterator[SimpleITKImage]:
             continue
 
         try:
-            reader = SimpleITK.ImageFileReader()
-            reader.SetImageIO("NiftiImageIO")
-            reader.SetFileName(str(file.absolute()))
-            img: SimpleITK.Image = reader.Execute()
+            img = load_sitk_image(file=file, imageio="NiftiImageIO")
+        except (ValidationError, NotImplementedError) as e:
+            file_errors[file].append(format_error(str(e)))
         except RuntimeError:
             file_errors[file].append(
                 format_error("Not a valid NifTI image file")
             )
-            continue
-
-        yield SimpleITKImage(
-            image=img,
-            name=file.name,
-            consumed_files={file},
-            spacing_valid=True,
-        )
+        else:
+            yield SimpleITKImage(
+                image=img,
+                name=file.name,
+                consumed_files={file},
+                spacing_valid=True,
+            )
 
     if file_errors:
         raise UnconsumedFilesException(file_errors=file_errors)
diff --git a/panimg/image_builders/nrrd.py b/panimg/image_builders/metaio_nrrd.py
similarity index 92%
rename from panimg/image_builders/nrrd.py
rename to panimg/image_builders/metaio_nrrd.py
index 6712d01..a6b7550 100644
--- a/panimg/image_builders/nrrd.py
+++ b/panimg/image_builders/metaio_nrrd.py
@@ -3,9 +3,8 @@
 from pathlib import Path
 from typing import DefaultDict, Iterator, List, Set
 
-import SimpleITK
-
-from panimg.exceptions import UnconsumedFilesException
+from panimg.exceptions import UnconsumedFilesException, ValidationError
+from panimg.image_builders.metaio_utils import load_sitk_image
 from panimg.models import SimpleITKImage
 
 MAGIC_REGEX = re.compile("NRRD([0-9]{4})")
@@ -98,15 +97,16 @@ def image_builder_nrrd(*, files: Set[Path]) -> Iterator[SimpleITKImage]:
     for file in files:
         try:
             if verify_single_file_nrrd(file):
-                reader = SimpleITK.ImageFileReader()
-                reader.SetImageIO("NrrdImageIO")
-                reader.SetFileName(str(file.absolute()))
-                img: SimpleITK.Image = reader.Execute()
+                img = load_sitk_image(file=file, imageio="NrrdImageIO")
             else:
                 raise InvalidNrrdFileError(
                     "NRRD files with detached headers are not supported"
                 )
-        except InvalidNrrdFileError as e:
+        except (
+            InvalidNrrdFileError,
+            ValidationError,
+            NotImplementedError,
+        ) as e:
             file_errors[file].append(format_error(str(e)))
         except RuntimeError:
             file_errors[file].append(
diff --git a/panimg/image_builders/metaio_utils.py b/panimg/image_builders/metaio_utils.py
index 2dddc98..2e07ced 100644
--- a/panimg/image_builders/metaio_utils.py
+++ b/panimg/image_builders/metaio_utils.py
@@ -183,10 +183,11 @@ def resolve_mh_data_file_path(
 
 
 def validate_and_clean_additional_mh_headers(
-    headers: Dict[str, str]
+    reader: SimpleITK.ImageFileReader,
 ) -> Dict[str, str]:
     cleaned_headers = {}
-    for key, value in headers.items():
+    for key in reader.GetMetaDataKeys():
+        value = reader.GetMetaData(key)
         if key in EXPECTED_HEADERS:
             cleaned_headers[key] = value
         elif key in ADDITIONAL_HEADERS:
@@ -198,19 +199,22 @@ def validate_and_clean_additional_mh_headers(
                     f"pattern {match_pattern.pattern}"
                 )
             if key in HEADERS_MATCHING_NUM_TIMEPOINTS:
+                timepoints = (
+                    reader.GetSize()[3] if reader.GetDimension() == 4 else 1
+                )
                 validate_list_data_matches_num_timepoints(
-                    headers=headers, key=key, value=value
+                    key=key, value=value, expected_timepoints=timepoints
                 )
             if key in HEADERS_MATCHING_WINDOW_SETTINGS:
                 validate_center_matches_width_setting(
-                    headers=headers, key=key, value=value
+                    key=key, value=value, reader=reader
                 )
             cleaned_headers[key] = value
     return cleaned_headers
 
 
 def validate_center_matches_width_setting(
-    headers: Dict[str, str], key: str, value: str
+    key: str, value: str, reader: SimpleITK.ImageFileReader
 ):
     window_keys = ("WindowWidth", "WindowCenter")
     if key not in window_keys:
@@ -219,11 +223,18 @@ def validate_center_matches_width_setting(
         return
 
     counter_key = window_keys[0] if key == window_keys[1] else window_keys[1]
-    if not re.match(FLOAT_ARRAY_MATCH_REGEXP, headers[counter_key]):
+    if not reader.HasMetaDataKey(counter_key):
+        raise ValidationError(
+            f"Headers '{key}' and '{counter_key}' should both be present "
+            f"but '{counter_key}' is missing"
+        )
+
+    counter_value = reader.GetMetaData(counter_key)
+    if not re.match(FLOAT_ARRAY_MATCH_REGEXP, counter_value):
         raise ValidationError(
             f"Header '{key}' is of a different format than '{counter_key}'"
         )
-    if len(value.split(",")) != len(headers[counter_key].split(",")):
+    if len(value.split(",")) != len(counter_value.split(",")):
         raise ValidationError(
             f"Headers '{key}' and '{counter_key}' should "
             f"contain an equal number of values"
@@ -231,14 +242,9 @@ def validate_center_matches_width_setting(
 
 
 def validate_list_data_matches_num_timepoints(
-    headers: Dict[str, str], key: str, value: str
+    key: str, value: str, expected_timepoints: int
 ):
     num_timepoints = len(value.split(" "))
-    expected_timepoints = (
-        int(headers["DimSize"].split(" ")[3])
-        if int(headers["NDims"]) >= 4
-        else 1
-    )
     if num_timepoints != expected_timepoints:
         raise ValidationError(
             f"Found {num_timepoints} values for {key}, "
@@ -247,9 +253,8 @@ def validate_list_data_matches_num_timepoints(
 
 
 def add_additional_mh_headers_to_sitk_image(
-    sitk_image: SimpleITK.Image, headers: Dict[str, str]
+    sitk_image: SimpleITK.Image, cleaned_headers: Dict[str, str]
 ):
-    cleaned_headers = validate_and_clean_additional_mh_headers(headers)
     for header in ADDITIONAL_HEADERS:
         if header in cleaned_headers:
             value = cleaned_headers[header]
@@ -260,21 +265,30 @@ def add_additional_mh_headers_to_sitk_image(
             sitk_image.SetMetaData(header, value)
 
 
-def load_sitk_image(mhd_file: Path) -> SimpleITK.Image:
-    headers = parse_mh_header(mhd_file)
-    headers = validate_and_clean_additional_mh_headers(headers=headers)
-    ndims = int(headers["NDims"])
-    if ndims <= 4:
-        sitk_image = SimpleITK.ReadImage(str(mhd_file))
-        for key in sitk_image.GetMetaDataKeys():
-            if key not in ADDITIONAL_HEADERS:
-                sitk_image.EraseMetaData(key)
+def load_sitk_image(
+    file: Path, imageio: str = "MetaImageIO"
+) -> SimpleITK.Image:
+    # Read and validate only the header first
+    reader = SimpleITK.ImageFileReader()
+    reader.SetImageIO(imageio)
+    reader.SetFileName(str(file.absolute()))
+    reader.ReadImageInformation()
+
+    headers = validate_and_clean_additional_mh_headers(reader=reader)
+
+    # Header has been validated, read the pixel data
+    if reader.GetDimension() in (2, 3, 4):
+        sitk_image = reader.Execute()
     else:
-        error_msg = (
+        raise NotImplementedError(
             "SimpleITK images with more than 4 dimensions are not supported"
         )
-        raise NotImplementedError(error_msg)
+
+    # Remove all headers, add back cleaned versions of supported additional headers
+    for key in sitk_image.GetMetaDataKeys():
+        sitk_image.EraseMetaData(key)
     add_additional_mh_headers_to_sitk_image(
-        sitk_image=sitk_image, headers=headers
+        sitk_image=sitk_image, cleaned_headers=headers
     )
+
     return sitk_image


Test patch:
diff --git a/tests/resources/image4x4x4x4x4.mha b/tests/resources/image4x4x4x4x4.mha
new file mode 100644
index 0000000..aedca48
Binary files /dev/null and b/tests/resources/image4x4x4x4x4.mha differ
diff --git a/tests/test_dicom.py b/tests/test_dicom.py
index 7d60ae1..f8075e5 100644
--- a/tests/test_dicom.py
+++ b/tests/test_dicom.py
@@ -25,8 +25,8 @@ def test_get_headers_by_study():
     studies = _get_headers_by_study(files, defaultdict(list))
     assert len(studies) == 1
     for key in studies:
-        assert [str(x["file"]) for x in studies[key]["headers"]] == [
-            f"{DICOM_DIR}/{x}.dcm" for x in range(1, 77)
+        assert [x["file"] for x in studies[key]["headers"]] == [
+            DICOM_DIR / f"{x}.dcm" for x in range(1, 77)
         ]
 
     for root, _, files in os.walk(RESOURCE_PATH):
diff --git a/tests/test_mhd.py b/tests/test_mhd.py
index aea5ecb..6bf6e16 100644
--- a/tests/test_mhd.py
+++ b/tests/test_mhd.py
@@ -170,16 +170,9 @@ def test_4d_mh_loader_with_uncompressed_data(tmpdir):
     load_sitk_image(tmp_header_file)
 
 
-def test_4d_mh_loader_with_more_than_4_dimensions_fails(tmpdir):
-    src = RESOURCE_PATH / "image10x11x12x13.mhd"
-    dest = Path(tmpdir) / src.name
-    shutil.copy(str(src), str(dest))
-    with open(str(dest)) as f:
-        modified_header = f.read().replace("NDims = 4", "NDims = 5")
-    with open(str(dest), "w") as f:
-        f.write(modified_header)
+def test_4d_mh_loader_with_more_than_4_dimensions_fails():
     with pytest.raises(NotImplementedError):
-        load_sitk_image(dest)
+        load_sitk_image(RESOURCE_PATH / "image4x4x4x4x4.mha")
 
 
 @pytest.mark.parametrize(
diff --git a/tests/test_nifti.py b/tests/test_nifti.py
index a2b192e..e02ebeb 100644
--- a/tests/test_nifti.py
+++ b/tests/test_nifti.py
@@ -5,7 +5,7 @@
 
 import pytest
 
-from panimg.image_builders.nifti import image_builder_nifti
+from panimg.image_builders.metaio_nifti import image_builder_nifti
 from panimg.models import ColorSpace
 from panimg.panimg import _build_files
 from tests import RESOURCE_PATH
diff --git a/tests/test_nrrd.py b/tests/test_nrrd.py
index d17694e..951c67a 100644
--- a/tests/test_nrrd.py
+++ b/tests/test_nrrd.py
@@ -6,7 +6,7 @@
 
 import pytest
 
-from panimg.image_builders.nrrd import image_builder_nrrd
+from panimg.image_builders.metaio_nrrd import image_builder_nrrd
 from panimg.models import ColorSpace
 from panimg.panimg import _build_files
 from tests import RESOURCE_PATH