Fractional and Multisegment Pyramids (#297)

* Allow fractional and multisegment pyramids; add tests

* Update docstring

* Fix for labelmap style masks

* Lint fix

Author: CPBridge

src/highdicom/seg/pyramid.py

@@ -77,7 +77,10 @@ def create_segmentation_pyramid(
         List of source images. If there are multiple source images, they should
         be from the same series and pyramid.
     pixel_arrays: Sequence[numpy.ndarray]
-        List of segmentation pixel arrays. Each should be a total pixel matrix.
+        List of segmentation pixel arrays. Each should be a total pixel matrix,
+        i.e. have shape (rows, columns), (1, rows, columns), or (1, rows,
+        columns, segments). Otherwise all options supported by the constructor
+        of :class:`highdicom.seg.Segmentation` are permitted.
     segmentation_type: Union[str, highdicom.seg.SegmentationTypeValues]
         Type of segmentation, either ``"BINARY"`` or ``"FRACTIONAL"``
     segment_descriptions: Sequence[highdicom.seg.SegmentDescription]
@@ -123,16 +126,17 @@ def create_segmentation_pyramid(
         A human readable label for the output pyramid.
     **kwargs: Any
         Any further parameters are passed directly to the constructor of the
-        :class:highdicom.seg.Segmentation object. However the following
+        :class:`highdicom.seg.Segmentation` object. However the following
         parameters are disallowed: ``instance_number``, ``sop_instance_uid``,
         ``plane_orientation``, ``plane_positions``, ``pixel_measures``,
         ``pixel_array``, ``tile_pixel_array``.
 
     Note
     ----
-    Downsampling is performed via simple nearest neighbor interpolation. If
-    more control is needed over the downsampling process (for example
-    anti-aliasing), explicitly pass the downsampled arrays.
+    Downsampling is performed via simple nearest neighbor interpolation (for
+    ``BINARY`` segmentations) or bi-linear interpolation (for ``FRACTIONAL``
+    segmentations). If more control is needed over the downsampling process
+    (for example anti-aliasing), explicitly pass the downsampled arrays.
 
     """
     # Disallow duplicate items in kwargs
@@ -149,9 +153,11 @@ def create_segmentation_pyramid(
     if len(error_keys) > 0:
         raise TypeError(
             f'kwargs supplied to the create_segmentation_pyramid function '
-            f'should not contain a value for parameter {error_keys[0]}.'
+            f'should not contain a value for parameter {list(error_keys)[0]}.'
         )
 
+    segmentation_type = SegmentationTypeValues(segmentation_type)
+
     if pyramid_uid is None:
         pyramid_uid = UID()
     if series_instance_uid is None:
@@ -263,6 +269,30 @@ def create_segmentation_pyramid(
             )
 
     # Check that pixel arrays have an appropriate shape
+    if len(set(p.ndim for p in pixel_arrays)) != 1:
+        raise ValueError(
+            'Each item of argument "pixel_arrays" must have the same number of '
+            'dimensions.'
+        )
+    if pixel_arrays[0].ndim == 4:
+        n_segment_channels = pixel_arrays[0].shape[3]
+    else:
+        n_segment_channels = None
+
+    dtype = pixel_arrays[0].dtype
+    if dtype in (np.bool_, np.uint8, np.uint16):
+        resampler = Image.Resampling.NEAREST
+    elif dtype in (np.float32, np.float64):
+        if segmentation_type == SegmentationTypeValues.FRACTIONAL:
+            resampler = Image.Resampling.BILINEAR
+        else:
+            # This is a floating point image that will ultimately be treated as
+            # binary
+            resampler = Image.Resampling.NEAREST
+    else:
+        raise TypeError('Pixel array has an invalid data type.')
+
+    # Checks on consistency of the pixel arrays
     for pixel_array in pixel_arrays:
         if pixel_array.ndim not in (2, 3, 4):
             raise ValueError(
@@ -274,6 +304,17 @@ def create_segmentation_pyramid(
                 'Each item of argument "pixel_arrays" must contain a single '
                 'frame, with a size of 1 along dimension 0.'
             )
+        if pixel_array.dtype != dtype:
+            raise TypeError(
+                'Each item of argument "pixel_arrays" must have '
+                'the same datatype.'
+            )
+        if pixel_array.ndim == 4:
+            if pixel_array.shape[3] != n_segment_channels:
+                raise ValueError(
+                    'Each item of argument "pixel_arrays" must have '
+                    'the same shape down axis 3.'
+                )
 
     # Check the pixel arrays are appropriately ordered
     for index in range(1, len(pixel_arrays)):
@@ -322,7 +363,17 @@ def create_segmentation_pyramid(
 
     if n_pix_arrays == 1:
         # Create a pillow image for use later with resizing
-        mask_image = Image.fromarray(pixel_arrays[0])
+        if pixel_arrays[0].ndim == 2:
+            mask_images = [Image.fromarray(pixel_arrays[0])]
+        elif pixel_arrays[0].ndim == 3:
+            # Remove frame dimension before casting
+            mask_images = [Image.fromarray(pixel_arrays[0][0])]
+        else:  # ndim = 4
+            # One "Image" for each segment
+            mask_images = [
+                Image.fromarray(pixel_arrays[0][0, :, :, i])
+                for i in range(pixel_arrays[0].shape[3])
+            ]
 
     all_segs = []
 
@@ -338,9 +389,7 @@ def create_segmentation_pyramid(
         else:
             if output_level == 0:
                 pixel_array = pixel_arrays[0]
-                need_resize = False
             else:
-                need_resize = True
                 if n_sources > 1:
                     output_size = (
                         source_image.TotalPixelMatrixColumns,
@@ -353,10 +402,14 @@ def create_segmentation_pyramid(
                         int(source_images[0].TotalPixelMatrixRows / f)
                     )
 
-            if need_resize:
-                pixel_array = np.array(
-                    mask_image.resize(output_size, Image.Resampling.NEAREST)
-                )
+                resized_masks = [
+                    np.array(im.resize(output_size, resampler))
+                    for im in mask_images
+                ]
+                if len(resized_masks) > 1:
+                    pixel_array = np.stack(resized_masks, axis=-1)[None]
+                else:
+                    pixel_array = resized_masks[0][None]
 
         if n_sources == 1:
             source_pixel_measures = (

tests/test_seg.py

@@ -3761,9 +3761,31 @@ def setUp(self):
         )
         self._seg_pix[5:15, 3:8] = 1
 
+        self._seg_pix_multisegment = np.zeros(
+            (1, *tpm_size, 3),
+            dtype=np.uint8,
+        )
+        self._seg_pix_multisegment[0, 5:15, 3:8, 0] = 1
+        self._seg_pix_multisegment[0, 23:42, 13:18, 1] = 1
+        self._seg_pix_multisegment[0, 45:48, 25:30, 1] = 1
+
+        self._seg_pix_fractional = np.zeros(
+            tpm_size,
+            dtype=np.float32,
+        )
+        self._seg_pix_fractional[5:15, 3:8] = 0.5
+        self._seg_pix_fractional[3:5, 10:12] = 1.0
+
         self._n_downsamples = 3
         self._downsampled_pix_arrays = [self._seg_pix]
+        self._downsampled_pix_arrays_multisegment = [self._seg_pix_multisegment]
+        self._downsampled_pix_arrays_fractional = [self._seg_pix_fractional]
         seg_pil = Image.fromarray(self._seg_pix)
+        seg_pil_fractional = Image.fromarray(self._seg_pix_fractional)
+        seg_pil_multisegment = [
+            Image.fromarray(self._seg_pix_multisegment[0, :, :, i])
+            for i in range(self._seg_pix_multisegment.shape[3])
+        ]
         pyramid_uid = UID()
         self._source_pyramid = [deepcopy(self._sm_image)]
         self._source_pyramid[0].PyramidUID = pyramid_uid
@@ -3780,6 +3802,19 @@ def setUp(self):
             )
             self._downsampled_pix_arrays.append(resized)
 
+            resized_fractional = np.array(
+                seg_pil_fractional.resize(out_size, Image.Resampling.BILINEAR)
+            )
+            self._downsampled_pix_arrays_fractional.append(resized_fractional)
+
+            resized_multisegment = np.stack(
+                [
+                    im.resize(out_size, Image.Resampling.NEAREST) for im in seg_pil_multisegment
+                ],
+                axis=-1
+            )[None]
+            self._downsampled_pix_arrays_multisegment.append(resized_multisegment)
+
             # Mock lower-resolution source images. No need to have their pixel
             # data correctly set as it isn't used. Just update the relevant
             # metadata
@@ -3820,6 +3855,21 @@ def setUp(self):
                 )
             ),
         ]
+        self._segment_descriptions_multi = [
+            SegmentDescription(
+                segment_number=i,
+                segment_label=f'Segment #{i}',
+                segmented_property_category=self._segmented_property_category,
+                segmented_property_type=self._segmented_property_type,
+                algorithm_type=SegmentAlgorithmTypeValues.AUTOMATIC.value,
+                algorithm_identification=AlgorithmIdentificationSequence(
+                    name='bla',
+                    family=codes.DCM.ArtificialIntelligence,
+                    version='v1'
+                )
+            )
+            for i in range(1, 4)
+        ]
 
     def test_pyramid_factors(self):
         downsample_factors = [2.0, 5.0]
@@ -3925,6 +3975,83 @@ def test_multiple_source_single_pixel_array(self):
                 pix
             )
 
+    def test_multiple_source_single_pixel_array_multisegment(self):
+        # Test construction when given multiple source images and a single
+        # segmentation image
+        segs = create_segmentation_pyramid(
+            source_images=self._source_pyramid,
+            pixel_arrays=[self._seg_pix_multisegment],
+            segmentation_type=SegmentationTypeValues.BINARY,
+            segment_descriptions=self._segment_descriptions_multi,
+            series_instance_uid=UID(),
+            series_number=1,
+            manufacturer='Foo',
+            manufacturer_model_name='Bar',
+            software_versions='1',
+            device_serial_number='123',
+        )
+
+        assert len(segs) == len(self._source_pyramid)
+        for pix, seg in zip(self._downsampled_pix_arrays_multisegment, segs):
+            assert hasattr(seg, 'PyramidUID')
+            seg_pix = seg.get_total_pixel_matrix()
+            print("pixel array", pix.shape, seg_pix.shape)
+            print("pixel array", pix.max(), seg_pix.max())
+            assert np.array_equal(
+                seg.get_total_pixel_matrix(),
+                pix[0]
+            )
+
+    def test_multiple_source_single_pixel_array_float(self):
+        # Test construction when given multiple source images and a single
+        # segmentation image
+        segs = create_segmentation_pyramid(
+            source_images=self._source_pyramid,
+            pixel_arrays=[self._seg_pix.astype(np.float32)],
+            segmentation_type=SegmentationTypeValues.BINARY,
+            segment_descriptions=self._segment_descriptions,
+            series_instance_uid=UID(),
+            series_number=1,
+            manufacturer='Foo',
+            manufacturer_model_name='Bar',
+            software_versions='1',
+            device_serial_number='123',
+        )
+
+        assert len(segs) == len(self._source_pyramid)
+        for pix, seg in zip(self._downsampled_pix_arrays, segs):
+            assert hasattr(seg, 'PyramidUID')
+            assert np.array_equal(
+                seg.get_total_pixel_matrix(combine_segments=True),
+                pix
+            )
+
+    def test_multiple_source_single_pixel_array_fractional(self):
+        # Test construction when given multiple source images and a single
+        # segmentation image
+        segs = create_segmentation_pyramid(
+            source_images=self._source_pyramid,
+            pixel_arrays=[self._seg_pix_fractional],
+            segmentation_type=SegmentationTypeValues.FRACTIONAL,
+            segment_descriptions=self._segment_descriptions,
+            series_instance_uid=UID(),
+            series_number=1,
+            manufacturer='Foo',
+            manufacturer_model_name='Bar',
+            software_versions='1',
+            device_serial_number='123',
+            max_fractional_value=200,
+        )
+
+        assert len(segs) == len(self._source_pyramid)
+        for pix, seg in zip(self._downsampled_pix_arrays_fractional, segs):
+            assert hasattr(seg, 'PyramidUID')
+            assert np.allclose(
+                seg.get_total_pixel_matrix(combine_segments=False)[:, :, 0],
+                pix,
+                atol=0.005,  # 1/200
+            )
+
     def test_multiple_source_multiple_pixel_arrays(self):
         # Test construction when given multiple source images and multiple
         # segmentation images
@@ -3948,3 +4075,53 @@ def test_multiple_source_multiple_pixel_arrays(self):
                 seg.get_total_pixel_matrix(combine_segments=True),
                 pix
             )
+
+    def test_multiple_source_multiple_pixel_arrays_multisegment(self):
+        # Test construction when given multiple source images and multiple
+        # segmentation images
+        segs = create_segmentation_pyramid(
+            source_images=self._source_pyramid,
+            pixel_arrays=self._downsampled_pix_arrays_multisegment,
+            segmentation_type=SegmentationTypeValues.BINARY,
+            segment_descriptions=self._segment_descriptions_multi,
+            series_instance_uid=UID(),
+            series_number=1,
+            manufacturer='Foo',
+            manufacturer_model_name='Bar',
+            software_versions='1',
+            device_serial_number='123',
+        )
+
+        assert len(segs) == len(self._source_pyramid)
+        for pix, seg in zip(self._downsampled_pix_arrays_multisegment, segs):
+            assert hasattr(seg, 'PyramidUID')
+            assert np.array_equal(
+                seg.get_total_pixel_matrix(),
+                pix[0]
+            )
+
+    def test_multiple_source_multiple_pixel_arrays_multisegment_labelmap(self):
+        # Test construction when given multiple source images and multiple
+        # segmentation images
+        mask = np.argmax(self._seg_pix_multisegment, axis=3).astype(np.uint8)
+        segs = create_segmentation_pyramid(
+            source_images=self._source_pyramid,
+            pixel_arrays=mask,
+            segmentation_type=SegmentationTypeValues.BINARY,
+            segment_descriptions=self._segment_descriptions_multi,
+            series_instance_uid=UID(),
+            series_number=1,
+            manufacturer='Foo',
+            manufacturer_model_name='Bar',
+            software_versions='1',
+            device_serial_number='123',
+        )
+
+        assert len(segs) == len(self._source_pyramid)
+        for pix, seg in zip(self._downsampled_pix_arrays_multisegment, segs):
+            mask = np.argmax(pix, axis=3).astype(np.uint8)
+            assert hasattr(seg, 'PyramidUID')
+            assert np.array_equal(
+                seg.get_total_pixel_matrix(combine_segments=True),
+                mask[0]
+            )