Add support for 4D Overlays (#100)

This PR restricts the overlay types to (U)Int8 and changes how the
segments are determined for 4D segmentations. Maximum memory usage is
reduced by 30% for 4D mask volumes.

See DIAGNijmegen/rse-roadmap#246
See DIAGNijmegen/rse-grand-challenge-admin#186

panimg/models.py

@@ -19,13 +19,7 @@
 
 MASK_TYPE_PIXEL_IDS = [
     SimpleITK.sitkInt8,
-    SimpleITK.sitkInt16,
-    SimpleITK.sitkInt32,
-    SimpleITK.sitkInt64,
     SimpleITK.sitkUInt8,
-    SimpleITK.sitkUInt16,
-    SimpleITK.sitkUInt32,
-    SimpleITK.sitkUInt64,
 ]
 
 
@@ -75,7 +69,9 @@ class PatientSex(str, Enum):
     "DA": lambda v: datetime.date(int(v[:4]), int(v[4:6]), int(v[6:8]))
 }
 
-MAXIMUM_SEGMENTS_LENGTH = 32
+# NOTE: Only int8 or uint8 data types are checked for segments
+# so the true maximum is 256
+MAXIMUM_SEGMENTS_LENGTH = 64
 
 
 class ExtraMetaData(NamedTuple):
@@ -207,6 +203,10 @@ def depth(self) -> Optional[int]:
 
         return depth or None
 
+    @property
+    def is_4d(self):
+        return len(self.image.GetSize()) == 4
+
     @staticmethod
     def _extract_first_float(value: str) -> float:
         if value.startswith("["):
@@ -271,7 +271,26 @@ def segments(self) -> Optional[FrozenSet[int]]:
         if self.image.GetPixelIDValue() not in MASK_TYPE_PIXEL_IDS:
             return None
 
-        segments = np.unique(GetArrayViewFromImage(self.image))
+        im_arr = GetArrayViewFromImage(self.image)
+
+        if self.is_4d:
+            segments = set()
+            n_volumes = self.image.GetSize()[3]
+
+            for volume in range(n_volumes):
+                # Calculate the segments for each volume for memory efficiency
+                volume_segments = np.unique(im_arr[volume, :, :, :])
+                segments.update({*volume_segments})
+
+                if not segments.issubset({0, 1}):
+                    # 4D Segmentations must only have values 0 and 1
+                    # as the 4th dimension encodes the overlay type
+                    return None
+
+            # Use 1-indexing for each segmentation
+            segments = {idx + 1 for idx in range(n_volumes)}
+        else:
+            segments = np.unique(im_arr)
 
         if len(segments) <= MAXIMUM_SEGMENTS_LENGTH:
             return frozenset(segments)

tests/test_models.py

@@ -166,33 +166,31 @@ def test_sitk_image_value_range(
     "src_image,builder,segments",
     [
         (
+            "segments/-10_10_Int8.mha",
+            image_builders.image_builder_mhd,
+            frozenset(range(-10, 10)),
+        ),
+        (
+            "segments/0_10_UInt8.mha",
+            image_builders.image_builder_mhd,
+            frozenset(range(10)),
+        ),
+        (
+            "segments/4D_1_1_1_5_UInt8.mha",
+            image_builders.image_builder_mhd,
+            frozenset({1, 2, 3, 4, 5}),
+        ),
+        (
+            # Contains non-zero or ones as values
+            "segments/4D_1_1_1_5_UInt8_threes.mha",
+            image_builders.image_builder_mhd,
+            None,
+        ),
+        (
+            # Datatype is not sitkUInt8 or sitkInt8
             "image_min10_max10.mha",
             image_builders.image_builder_mhd,
-            frozenset(
-                (
-                    -10,
-                    -9,
-                    -8,
-                    -7,
-                    -6,
-                    -5,
-                    -4,
-                    -3,
-                    -2,
-                    -1,
-                    0,
-                    1,
-                    2,
-                    3,
-                    4,
-                    5,
-                    6,
-                    7,
-                    8,
-                    9,
-                    10,
-                )
-            ),
+            None,
         ),
         (  # Too many values
             "dicom_2d/cxr.dcm",