Repository: comic/grand-challenge.org

Issue:
Enable additional header information for 4DCT
For the implementation of the colormapping pipeline, additional header information must be read out of the mhd file that is uploaded. These are the contenttimes and mAs of the CTP acquisition, t0 and t1 values for the visualization.




Discussion:
You should be able to do that within your container, I don't think that this needs anything on the grand challenge side. Please re-open if I'm mistaken.
@jmsmkn I think that we need to be able to do this on the grand challenge side as well. If we are at least uploading our images through grand-challenge and we pass our 4D CTP image to the algorithm, then it needs to contain the additional headers. However, these headers are discarded on upload currently. We should properly support these specific headers and test them for valid values.
Ah OK, I thought that they weren't cleaned. Maybe add an allow list of headers that we keep, and only use ones where there cannot be any identifiable information. 
Ok, the minimal header information we currently require is:

```
Exposures:List[float] (in mAs / time point)
ContentTimes:List[str]  (where str is timestamp in "HHMMSS.fff" / time point)
t0:float (begin timepoint in flow window, specific for colormapping algorithm)
t1:float (end timepoint in flow window, specific for colormapping algorithm)
```

I'll create a whitelist approach for this.


Resolving patch:
diff --git a/app/grandchallenge/cases/image_builders/metaio_mhd_mha.py b/app/grandchallenge/cases/image_builders/metaio_mhd_mha.py
--- a/app/grandchallenge/cases/image_builders/metaio_mhd_mha.py
+++ b/app/grandchallenge/cases/image_builders/metaio_mhd_mha.py
@@ -27,7 +27,7 @@ def image_builder_mhd(path: Path) -> ImageBuilderResult:  # noqa: C901
     Parameters
     ----------
     path: Path
-        Path to a directory that contains all images that were uploaded duing
+        Path to a directory that contains all images that were uploaded during
         an upload session.
 
     Returns
diff --git a/app/grandchallenge/cases/image_builders/metaio_utils.py b/app/grandchallenge/cases/image_builders/metaio_utils.py
--- a/app/grandchallenge/cases/image_builders/metaio_utils.py
+++ b/app/grandchallenge/cases/image_builders/metaio_utils.py
@@ -1,6 +1,7 @@
+import re
 import zlib
 from pathlib import Path
-from typing import Any, Dict, List, Mapping, Tuple, Union
+from typing import Any, Dict, List, Mapping, Pattern, Tuple, Union
 
 import SimpleITK
 import SimpleITK._SimpleITK as _SimpleITK
@@ -37,6 +38,45 @@
     "MET_OTHER": None,
 }
 
+FLOAT_MATCH_REGEXP: Pattern = re.compile(
+    r"^[-+]?[0-9]*\.?[0-9]+([eE][-+]?[0-9]+)?$"
+)
+FLOAT_LIST_MATCH_REGEXP: Pattern = re.compile(
+    r"^([-+]?[0-9]*\.?[0-9]+([eE][-+]?[0-9]+)?)"
+    r"(\s[-+]?[0-9]*\.?[0-9]+([eE][-+]?[0-9]+)?)*$"
+)
+CONTENT_TIMES_LIST_MATCH_REGEXP: Pattern = re.compile(
+    r"^((2[0-3]|[0-1]\d)[0-5]\d[0-5]\d(\.\d\d\d)?)"
+    r"(\s(2[0-3]|[0-1]\d)[0-5]\d[0-5]\d(\.\d\d\d)?)*$"
+)
+
+ADDITIONAL_HEADERS: Dict[str, Pattern] = {
+    "Exposures": FLOAT_LIST_MATCH_REGEXP,
+    "ContentTimes": CONTENT_TIMES_LIST_MATCH_REGEXP,
+    "t0": FLOAT_MATCH_REGEXP,
+    "t1": FLOAT_MATCH_REGEXP,
+}
+
+HEADERS_MATCHING_NUM_TIMEPOINTS: List[str] = ["Exposures", "ContentTimes"]
+
+EXPECTED_HEADERS: List[str] = [
+    "ObjectType",
+    "NDims",
+    "BinaryData",
+    "BinaryDataByteOrderMSB",
+    "CompressedData",
+    "CompressedDataSize",
+    "TransformMatrix",
+    "Offset",
+    "CenterOfRotation",
+    "AnatomicalOrientation",
+    "ElementSpacing",
+    "ElementNumberOfChannels",
+    "DimSize",
+    "ElementType",
+    "ElementDataFile",
+]
+
 
 def parse_mh_header(filename: Path) -> Mapping[str, Union[str, None]]:
     """
@@ -183,16 +223,77 @@ def create_sitk_img_from_mh_data(
     return sitk_image
 
 
+def validate_and_clean_additional_mh_headers(
+    headers: Mapping[str, Union[str, None]]
+) -> Mapping[str, Union[str, None]]:
+    cleaned_headers = {}
+    for key, value in headers.items():
+        if key in EXPECTED_HEADERS:
+            cleaned_headers[key] = value
+        else:
+            if key in ADDITIONAL_HEADERS:
+                match_pattern = ADDITIONAL_HEADERS[key]
+                if not re.match(match_pattern, value):
+                    raise ValueError(
+                        f"Invalid data type found for "
+                        f"additional header key: {key}"
+                    )
+                cleaned_headers[key] = value
+        if key in HEADERS_MATCHING_NUM_TIMEPOINTS:
+            validate_list_data_matches_num_timepoints(
+                headers=headers, key=key, value=value
+            )
+
+    return cleaned_headers
+
+
+def validate_list_data_matches_num_timepoints(
+    headers: Mapping[str, Union[str, None]], key: str, value: str
+):
+    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 ValueError(
+            f"Found {num_timepoints} values for {key}, "
+            f"but expected {expected_timepoints} (1/timepoint)"
+        )
+
+
+def add_additional_mh_headers_to_sitk_image(
+    sitk_image: SimpleITK.Image, headers: Mapping[str, Union[str, None]]
+):
+    cleaned_headers = validate_and_clean_additional_mh_headers(headers)
+    for header in ADDITIONAL_HEADERS:
+        if header in cleaned_headers:
+            value = cleaned_headers[header]
+            if isinstance(value, (list, tuple)):
+                value = " ".format([str(v) for v in value])
+            else:
+                value = str(value)
+            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))
-    elif ndims == 4:
+        for key in sitk_image.GetMetaDataKeys():
+            if key not in ADDITIONAL_HEADERS:
+                sitk_image.EraseMetaData(key)
+    elif ndims <= 4:
         sitk_image = load_sitk_image_with_nd_support(mhd_file=mhd_file)
     else:
         error_msg = (
             "SimpleITK images with more than 4 dimensions are not supported"
         )
         raise NotImplementedError(error_msg)
+    add_additional_mh_headers_to_sitk_image(
+        sitk_image=sitk_image, headers=headers
+    )
     return sitk_image


Test patch:
diff --git a/app/tests/cases_tests/resources/image10x11x12x13-extra-stuff.mhd b/app/tests/cases_tests/resources/image10x11x12x13-extra-stuff.mhd
new file mode 100644
--- /dev/null
+++ b/app/tests/cases_tests/resources/image10x11x12x13-extra-stuff.mhd
@@ -0,0 +1,19 @@
+ObjectType = Image
+NDims = 4
+BinaryData = True
+BinaryDataByteOrderMSB = False
+CompressedData = True
+CompressedDataSize = 39
+TransformMatrix = 1 0 0 0 0 1 0 0 0 0 1 0 0 0 0 1
+Offset = -131 -99 -917 0
+CenterOfRotation = 0 0 0 0
+ElementSpacing = 0.429 0.429 0.5 1
+DimSize = 10 11 12 13
+AnatomicalOrientation = RAI
+ElementType = MET_UCHAR
+ContentTimes = 000000.000 111111.111 121212.333 121213 121214 121215 121216 121217 121218 121219 121220 121221 235959.999
+Exposures = 1 2 3e-5 4 5.2 6 7 8 9 10 11 12 13.0
+t0 = 12.5
+t1 = 1.2e-6
+Bogus = nonsense 1234
+ElementDataFile = image10x11x12x13.zraw
diff --git a/app/tests/cases_tests/resources/image3x4-extra-stuff.mhd b/app/tests/cases_tests/resources/image3x4-extra-stuff.mhd
new file mode 100644
--- /dev/null
+++ b/app/tests/cases_tests/resources/image3x4-extra-stuff.mhd
@@ -0,0 +1,20 @@
+ObjectType = Image
+NDims = 2
+BinaryData = True
+BinaryDataByteOrderMSB = False
+CompressedData = True
+CompressedDataSize = 11
+TransformMatrix = 1 0 0 1
+Offset = 0 0
+CenterOfRotation = 0 0
+ElementSpacing = 1 1
+DimSize = 3 4
+AnatomicalOrientation = ??
+ElementNumberOfChannels = 3
+ElementType = MET_UCHAR
+ContentTimes = 235959.999
+Exposures = 100.2
+t0 = 12.5
+t1 = 1.2e-6
+Bogus = nonsense 1234
+ElementDataFile = image3x4.zraw
diff --git a/app/tests/cases_tests/test_background_tasks.py b/app/tests/cases_tests/test_background_tasks.py
--- a/app/tests/cases_tests/test_background_tasks.py
+++ b/app/tests/cases_tests/test_background_tasks.py
@@ -1,7 +1,14 @@
 from typing import Dict, List, Tuple
 
+import SimpleITK
 import pytest
 
+from grandchallenge.cases.image_builders.metaio_utils import (
+    ADDITIONAL_HEADERS,
+    EXPECTED_HEADERS,
+    HEADERS_MATCHING_NUM_TIMEPOINTS,
+    parse_mh_header,
+)
 from grandchallenge.cases.models import (
     Image,
     RawImageFile,
@@ -151,6 +158,59 @@ def test_staged_4d_mha_and_4d_mhd_upload(settings, images: List):
     assert [e for e in reversed(sitk_image.GetSize())] == image.shape
 
 
+@pytest.mark.parametrize(
+    "images",
+    (
+        ["image10x11x12x13-extra-stuff.mhd", "image10x11x12x13.zraw"],
+        ["image3x4-extra-stuff.mhd", "image3x4.zraw"],
+    ),
+)
+@pytest.mark.django_db
+def test_staged_mhd_upload_with_additional_headers(
+    settings, tmp_path, images: List[str]
+):
+    # Override the celery settings
+    settings.task_eager_propagates = (True,)
+    settings.task_always_eager = (True,)
+
+    session, uploaded_images = create_raw_upload_image_session(images)
+
+    session.refresh_from_db()
+    assert session.session_state == UploadSessionState.stopped
+    assert session.error_message is None
+
+    images = Image.objects.filter(origin=session).all()
+    assert len(images) == 1
+
+    raw_image_file: RawImageFile = list(uploaded_images.values())[0]
+    raw_image_file.refresh_from_db()
+    assert raw_image_file.staged_file_id is None
+
+    image: Image = images[0]
+    tmp_header_filename = tmp_path / "tmp_header.mhd"
+    with image.files.get(file__endswith=".mhd").file.open(
+        "rb"
+    ) as in_file, open(tmp_header_filename, "wb") as out_file:
+        out_file.write(in_file.read())
+
+    headers = parse_mh_header(tmp_header_filename)
+    for key in headers.keys():
+        assert (key in ADDITIONAL_HEADERS) or (key in EXPECTED_HEADERS)
+
+    sitk_image: SimpleITK.Image = image.get_sitk_image()
+    for key in ADDITIONAL_HEADERS:
+        assert key in sitk_image.GetMetaDataKeys()
+        if key in HEADERS_MATCHING_NUM_TIMEPOINTS:
+            if sitk_image.GetDimension() >= 4:
+                assert (
+                    len(sitk_image.GetMetaData(key).split(" "))
+                    == sitk_image.GetSize()[3]
+                )
+            else:
+                assert len(sitk_image.GetMetaData(key).split(" ")) == 1
+    assert "Bogus" not in sitk_image.GetMetaDataKeys()
+
+
 @pytest.mark.django_db
 def test_no_convertible_file(settings):
     # Override the celery settings
diff --git a/app/tests/cases_tests/test_mhd.py b/app/tests/cases_tests/test_mhd.py
--- a/app/tests/cases_tests/test_mhd.py
+++ b/app/tests/cases_tests/test_mhd.py
@@ -6,6 +6,7 @@
 import pytest
 
 from grandchallenge.cases.image_builders.metaio_utils import (
+    ADDITIONAL_HEADERS,
     load_sitk_image,
     load_sitk_image_with_nd_support,
     parse_mh_header,
@@ -178,6 +179,49 @@ def test_4d_mh_loader_with_more_than_4_dimensions_fails(tmpdir):
         load_sitk_image(dest)
 
 
+@pytest.mark.parametrize(
+    "test_img",
+    ["image10x11x12x13-extra-stuff.mhd", "image3x4-extra-stuff.mhd"],
+)
+def test_load_sitk_image_with_additional_meta_data(tmpdir, test_img: str):
+    src = RESOURCE_PATH / test_img
+    sitk_image = load_sitk_image(src)
+    for key in sitk_image.GetMetaDataKeys():
+        assert key in ADDITIONAL_HEADERS
+        assert ADDITIONAL_HEADERS[key].match(sitk_image.GetMetaData(key))
+    assert "Bogus" not in sitk_image.GetMetaDataKeys()
+
+
+@pytest.mark.parametrize("test_img", ["image10x11x12x13.mhd", "image3x4.mhd"])
+@pytest.mark.parametrize(
+    ["key", "value"],
+    [
+        ("Exposures", "1 2 3e-5 4 5.2 6 7 8 9 10 11 12 13.0 e"),
+        ("Exposures", "1 2 3e-5 4 5.2 6 7 8 9 10 11 12"),
+        ("t0", "1 2"),
+        ("t1", "string"),
+        (
+            "ContentTimes",
+            ("245959.999 000000.000 111111.111 " * 4) + "121212.333",
+        ),
+        ("ContentTimes", "1 2 3e-5 4 5.2 6 7 8 9 10 11 12 13.0 14.0"),
+    ],
+)
+def test_load_sitk_image_with_corrupt_additional_meta_data_fails(
+    tmpdir, test_img: str, key: str, value: str
+):
+    src = RESOURCE_PATH / "image10x11x12x13.mhd"
+    dest = Path(tmpdir) / src.name
+    shutil.copy(str(src), str(dest))
+    with open(str(dest), "r") as f:
+        lines = f.readlines()
+    lines.insert(-1, f"{key} = {value}\n")
+    with open(str(dest), "w") as f:
+        f.writelines(lines)
+    with pytest.raises(ValueError):
+        load_sitk_image(dest)
+
+
 @pytest.mark.parametrize(
     "image",
     (