Add instance ID support to synthetic test images

monai/data/synthetic.py

@@ -28,7 +28,8 @@ def create_test_image_2d(
     num_seg_classes: int = 5,
     channel_dim: int | None = None,
     random_state: np.random.RandomState | None = None,
-) -> tuple[np.ndarray, np.ndarray]:
+    return_instance_id: bool = False,
+) -> tuple[np.ndarray, np.ndarray] | tuple[np.ndarray, np.ndarray, np.ndarray]:
     """
     Return a noisy 2D image with `num_objs` circles and a 2D mask image. The maximum and minimum radii of the circles
     are given as `rad_max` and `rad_min`. The mask will have `num_seg_classes` number of classes for segmentations labeled
@@ -48,9 +49,13 @@ def create_test_image_2d(
         channel_dim: if None, create an image without channel dimension, otherwise create
             an image with channel dimension as first dim or last dim. Defaults to `None`.
         random_state: the random generator to use. Defaults to `np.random`.
+        return_instance_id: if True, also return an instance ID mask where every generated
+            object is assigned a unique positive integer. Later objects overwrite earlier IDs
+            where they overlap, so fully covered objects may not appear in the mask. Defaults to `False`.
 
     Returns:
-        Randomised Numpy array with shape (`height`, `width`)
+        A tuple of image and segmentation label arrays. If `return_instance_id=True`, also returns
+        an instance ID array as the third element.
     """
 
     if rad_max <= rad_min:
@@ -62,9 +67,10 @@ def create_test_image_2d(
         raise ValueError(f"the minimal size {min_size} of the image should be larger than `2 * rad_max` 2x{rad_max}.")
 
     image = np.zeros((height, width))
+    instance_ids = np.zeros((height, width), dtype=np.int32) if return_instance_id else None
     rs: np.random.RandomState = np.random.random.__self__ if random_state is None else random_state  # type: ignore
 
-    for _ in range(num_objs):
+    for obj_id in range(1, num_objs + 1):
         x = rs.randint(rad_max, height - rad_max)
         y = rs.randint(rad_max, width - rad_max)
         rad = rs.randint(rad_min, rad_max)
@@ -75,6 +81,8 @@ def create_test_image_2d(
             image[circle] = np.ceil(rs.random() * num_seg_classes)
         else:
             image[circle] = rs.random() * 0.5 + 0.5
+        if instance_ids is not None:
+            instance_ids[circle] = obj_id
 
     labels = np.ceil(image).astype(np.int32, copy=False)
 
@@ -87,10 +95,16 @@ def create_test_image_2d(
         if channel_dim == 0:
             noisyimage = noisyimage[None]
             labels = labels[None]
+            if instance_ids is not None:
+                instance_ids = instance_ids[None]
         else:
             noisyimage = noisyimage[..., None]
             labels = labels[..., None]
+            if instance_ids is not None:
+                instance_ids = instance_ids[..., None]
 
+    if instance_ids is not None:
+        return noisyimage, labels, instance_ids
     return noisyimage, labels
 
 
@@ -105,7 +119,8 @@ def create_test_image_3d(
     num_seg_classes: int = 5,
     channel_dim: int | None = None,
     random_state: np.random.RandomState | None = None,
-) -> tuple[np.ndarray, np.ndarray]:
+    return_instance_id: bool = False,
+) -> tuple[np.ndarray, np.ndarray] | tuple[np.ndarray, np.ndarray, np.ndarray]:
     """
     Return a noisy 3D image and segmentation.
 
@@ -122,9 +137,13 @@ def create_test_image_3d(
         channel_dim: if None, create an image without channel dimension, otherwise create
             an image with channel dimension as first dim or last dim. Defaults to `None`.
         random_state: the random generator to use. Defaults to `np.random`.
+        return_instance_id: if True, also return an instance ID mask where every generated
+            object is assigned a unique positive integer. Later objects overwrite earlier IDs
+            where they overlap, so fully covered objects may not appear in the mask. Defaults to `False`.
 
     Returns:
-        Randomised Numpy array with shape (`height`, `width`, `depth`)
+        A tuple of image and segmentation label arrays. If `return_instance_id=True`, also returns
+        an instance ID array as the third element.
 
     See also:
         :py:meth:`~create_test_image_2d`
@@ -139,9 +158,10 @@ def create_test_image_3d(
         raise ValueError(f"the minimal size {min_size} of the image should be larger than `2 * rad_max` 2x{rad_max}.")
 
     image = np.zeros((height, width, depth))
+    instance_ids = np.zeros((height, width, depth), dtype=np.int32) if return_instance_id else None
     rs: np.random.RandomState = np.random.random.__self__ if random_state is None else random_state  # type: ignore
 
-    for _ in range(num_objs):
+    for obj_id in range(1, num_objs + 1):
         x = rs.randint(rad_max, height - rad_max)
         y = rs.randint(rad_max, width - rad_max)
         z = rs.randint(rad_max, depth - rad_max)
@@ -153,6 +173,8 @@ def create_test_image_3d(
             image[circle] = np.ceil(rs.random() * num_seg_classes)
         else:
             image[circle] = rs.random() * 0.5 + 0.5
+        if instance_ids is not None:
+            instance_ids[circle] = obj_id
 
     labels = np.ceil(image).astype(np.int32, copy=False)
 
@@ -162,8 +184,17 @@ def create_test_image_3d(
     if channel_dim is not None:
         if not (isinstance(channel_dim, int) and channel_dim in (-1, 0, 3)):
             raise AssertionError("invalid channel dim.")
-        noisyimage, labels = (
-            (noisyimage[None], labels[None]) if channel_dim == 0 else (noisyimage[..., None], labels[..., None])
-        )
+        if channel_dim == 0:
+            noisyimage = noisyimage[None]
+            labels = labels[None]
+            if instance_ids is not None:
+                instance_ids = instance_ids[None]
+        else:
+            noisyimage = noisyimage[..., None]
+            labels = labels[..., None]
+            if instance_ids is not None:
+                instance_ids = instance_ids[..., None]
 
+    if instance_ids is not None:
+        return noisyimage, labels, instance_ids
     return noisyimage, labels

tests/data/test_synthetic.py

@@ -12,6 +12,8 @@
 from __future__ import annotations
 
 import unittest
+from itertools import product
+from unittest.mock import Mock
 
 import numpy as np
 from parameterized import parameterized
@@ -39,11 +41,26 @@
     ],
 ]
 
+INSTANCE_ID_CASES = [
+    [2, {"width": 64, "height": 64, "num_objs": 5, "rad_max": 10, "rad_min": 4}],
+    [3, {"width": 40, "height": 40, "depth": 40, "num_objs": 4, "rad_max": 8, "rad_min": 3, "channel_dim": -1}],
+]
+
 
 class TestDiceCELoss(unittest.TestCase):
 
     @parameterized.expand(TEST_CASES)
     def test_create_test_image(self, dim, input_param, expected_img, expected_seg, expected_shape, expected_max_cls):
+        """Verify synthetic image shapes, label classes, and deterministic means.
+
+        Args:
+            dim: Spatial dimensionality of the generator.
+            input_param: Keyword arguments passed to the generator.
+            expected_img: Expected mean image intensity.
+            expected_seg: Expected mean segmentation label.
+            expected_shape: Expected image shape.
+            expected_max_cls: Expected maximum segmentation class.
+        """
         set_determinism(seed=0)
         if dim == 2:
             img, seg = create_test_image_2d(**input_param)
@@ -54,7 +71,62 @@ def test_create_test_image(self, dim, input_param, expected_img, expected_seg, e
         np.testing.assert_allclose(img.mean(), expected_img, atol=1e-7, rtol=1e-7)
         np.testing.assert_allclose(seg.mean(), expected_seg, atol=1e-7, rtol=1e-7)
 
+    @parameterized.expand(INSTANCE_ID_CASES)
+    def test_return_instance_id(self, dim, input_param):
+        """Verify instance mask shape, dtype, ID bounds, and foreground alignment.
+
+        Args:
+            dim: Spatial dimensionality of the generator.
+            input_param: Keyword arguments passed to the generator.
+        """
+        set_determinism(seed=0)
+        if dim == 2:
+            img, seg, instance_ids = create_test_image_2d(**input_param, return_instance_id=True)
+        else:  # dim == 3
+            img, seg, instance_ids = create_test_image_3d(**input_param, return_instance_id=True)
+
+        self.assertEqual(img.shape, seg.shape)
+        self.assertEqual(instance_ids.shape, seg.shape)
+        self.assertEqual(instance_ids.dtype, np.int32)
+        unique_ids = np.unique(instance_ids)
+        self.assertGreaterEqual(len(unique_ids), 2)
+        self.assertEqual(unique_ids[0], 0)
+        self.assertTrue(np.all(unique_ids <= input_param["num_objs"]))
+        np.testing.assert_array_equal(instance_ids > 0, seg > 0)
+
+    @parameterized.expand(product((2, 3), (0, 2, 12), (None, 0, -1)))
+    def test_instance_id_overlap(self, dim, offset, channel_dim):
+        """Check distinct IDs and later-object precedence at fixed object positions."""
+        generator = create_test_image_2d if dim == 2 else create_test_image_3d
+        centers = [(8,) * dim, (8 + offset,) + (8,) * (dim - 1)]
+        rs = Mock(spec=np.random.RandomState, wraps=np.random.RandomState(0))
+        rs.randint.side_effect = [value for center in centers for value in (*center, 3)]
+        image, labels, instance_ids = generator(
+            *((32,) * dim),
+            num_objs=2,
+            rad_min=3,
+            rad_max=4,
+            num_seg_classes=1,
+            channel_dim=channel_dim,
+            random_state=rs,
+            return_instance_id=True,
+        )
+
+        self.assertEqual(image.shape, labels.shape)
+        self.assertEqual(instance_ids.shape, labels.shape)
+        ids = instance_ids.squeeze()
+        self.assertEqual(ids[(0,) * dim], 0)
+        self.assertEqual(ids[centers[1]], 2)
+        first_only = (5,) + (8,) * (dim - 1)
+        self.assertEqual(ids[first_only], 2 if offset == 0 else 1)
+        if offset <= 2:
+            self.assertEqual(ids[centers[0]], 2)
+        expected_ids = [0, 2] if offset == 0 else [0, 1, 2]
+        np.testing.assert_array_equal(np.unique(ids), expected_ids)
+        np.testing.assert_array_equal(instance_ids > 0, labels > 0)
+
     def test_ill_radius(self):
+        """Verify invalid radius bounds and image sizes raise ValueError."""
         with self.assertRaisesRegex(ValueError, ""):
             img, seg = create_test_image_2d(32, 32, rad_max=20)
         with self.assertRaisesRegex(ValueError, ""):