Merge pull request #1988 from larrybradley/deblend-map

Add map of deblended labels

CHANGES.rst

@@ -28,6 +28,12 @@ New Features
   - An optional ``mask`` keyword was added to the ``gini`` function.
     [#1979]
 
+- ``photutils.segmentation``
+
+  - Added ``deblended_labels``, ``deblended_labels_map``, and
+    ``deblended_labels_inverse_map`` properties to ``SegmentationImage``
+    to identify and map any deblended labels. [#1988]
+
 Bug Fixes
 ^^^^^^^^^
 
@@ -41,6 +47,9 @@ Bug Fixes
   - Fixed a bug to ensure that the dtype of the ``SegmentationImage``
     ``labels`` always matches the image dtype. [#1986]
 
+  - Fixed a issue with the source labels after source deblending when
+    using ``relabel=False``. [#1988]
+
 API Changes
 ^^^^^^^^^^^
 

photutils/segmentation/core.py

@@ -47,6 +47,7 @@ def __init__(self, data):
         if not isinstance(data, np.ndarray):
             raise TypeError('Input data must be a numpy array')
         self.data = data
+        self._deblend_label_map = {}  # set by source deblender
 
     def __str__(self):
         cls_name = f'<{self.__class__.__module__}.{self.__class__.__name__}>'
@@ -136,6 +137,52 @@ def segments(self):
 
         return segments
 
+    @lazyproperty
+    def deblended_labels(self):
+        """
+        A sorted 1D array of deblended label numbers.
+
+        The list will be empty if deblending has not been performed or
+        if no sources were deblended.
+        """
+        if len(self._deblend_label_map) == 0:
+            return np.array([], dtype=self._data.dtype)
+        return np.sort(np.concatenate(list(self._deblend_label_map.values())))
+
+    @lazyproperty
+    def deblended_labels_map(self):
+        """
+        A dictionary mapping deblended label numbers to the original
+        parent label numbers.
+
+        The keys are the deblended label numbers and the values are the
+        original parent label numbers. Only deblended sources are
+        included in the dictionary.
+
+        The dictionary will be empty if deblending has not been
+        performed or if no sources were deblended.
+        """
+        inverse_map = {}
+        for key, values in self._deblend_label_map.items():
+            for value in values:
+                inverse_map[value] = key
+        return inverse_map
+
+    @lazyproperty
+    def deblended_labels_inverse_map(self):
+        """
+        A dictionary mapping the original parent label numbers to the
+        deblended label numbers.
+
+        The keys are the original parent label numbers and the values
+        are the deblended label numbers. Only deblended sources are
+        included in the dictionary.
+
+        The dictionary will be empty if deblending has not been
+        performed or if no sources were deblended.
+        """
+        return self._deblend_label_map
+
     @property
     def data(self):
         """
@@ -174,6 +221,7 @@ def data(self, value):
 
         self._data = value  # pylint: disable=attribute-defined-outside-init
         self.__dict__['labels'] = labels
+        self.__dict__['_deblend_label_map'] = {}  # reset deblended labels
 
     @lazyproperty
     def data_ma(self):
@@ -543,6 +591,20 @@ def cmap(self):
         """
         return self.make_cmap(background_color='#000000ff', seed=0)
 
+    def _update_deblend_label_map(self, relabel_map):
+        """
+        Update the deblended label map based on a relabel map.
+
+        Parameters
+        ----------
+        relabel_map : `~numpy.ndarray`
+            An array mapping the original label numbers to the new label
+            numbers.
+        """
+        # child_labels are the deblended labels
+        for parent_label, child_labels in self._deblend_label_map.items():
+            self._deblend_label_map[parent_label] = relabel_map[child_labels]
+
     def reassign_label(self, label, new_label, relabel=False):
         """
         Reassign a label number to a new number.
@@ -698,20 +760,22 @@ def reassign_labels(self, labels, new_label, relabel=False):
         if labels.size == 0:
             return
 
-        idx = np.zeros(self.max_label + 1, dtype=self.data.dtype)
-        idx[self.labels] = self.labels
-        idx[labels] = new_label  # reassign labels
+        dtype = self.data.dtype  # keep the original dtype
+        relabel_map = np.zeros(self.max_label + 1, dtype=dtype)
+        relabel_map[self.labels] = self.labels
+        relabel_map[labels] = new_label  # reassign labels
 
         if relabel:
-            labels = np.unique(idx[idx != 0])
+            labels = np.unique(relabel_map[relabel_map != 0])
             if len(labels) != 0:
-                idx2 = np.zeros(max(labels) + 1, dtype=self.data.dtype)
-                idx2[labels] = np.arange(len(labels)) + 1
-                idx = idx2[idx]
+                map2 = np.zeros(max(labels) + 1, dtype=dtype)
+                map2[labels] = np.arange(len(labels), dtype=dtype) + 1
+                relabel_map = map2[relabel_map]
 
-        data_new = idx[self.data]
+        data_new = relabel_map[self.data]
         self._reset_lazyproperties()  # reset all cached properties
         self._data = data_new  # use _data to avoid validation
+        self._update_deblend_label_map(relabel_map)
 
     def relabel_consecutive(self, start_label=1):
         """
@@ -767,6 +831,7 @@ def relabel_consecutive(self, start_label=1):
         self.__dict__['labels'] = new_labels
         if old_slices is not None:
             self.__dict__['slices'] = old_slices  # slice order is unchanged
+        self._update_deblend_label_map(new_label_map)
 
     def keep_label(self, label, relabel=False):
         """

photutils/segmentation/deblend.py

@@ -186,12 +186,13 @@ def deblend_sources(data, segment_img, npixels, *, labels=None, nlevels=32,
     if nproc is None:
         nproc = cpu_count()  # pragma: no cover
 
+    deblend_label_map = {}
+    max_label = segment_img.max_label
     if nproc == 1:
         if progress_bar:  # pragma: no cover
             desc = 'Deblending'
             label_indices = add_progress_bar(label_indices, desc=desc)
 
-        max_label = segment_img.max_label + 1
         nonposmin_labels = []
         nmarkers_labels = []
         for label, label_idx in zip(labels, label_indices, strict=True):
@@ -213,10 +214,12 @@ def deblend_sources(data, segment_img, npixels, *, labels=None, nlevels=32,
 
             if source_deblended is not None:
                 source_mask = source_deblended > 0
+                new_segm = source_deblended[source_mask]  # min label = 1
                 segm_deblended[source_slice][source_mask] = (
-                    source_deblended[source_mask] + max_label)
-                nlabels = len(_get_labels(source_deblended))
-                max_label += nlabels
+                    new_segm + max_label)
+                new_labels = _get_labels(new_segm) + max_label
+                deblend_label_map[label] = new_labels
+                max_label += len(new_labels)
 
     else:
         # Use multiprocessing to deblend sources
@@ -263,7 +266,6 @@ def deblend_sources(data, segment_img, npixels, *, labels=None, nlevels=32,
                     results[idx] = future.result()
 
         # Process the results
-        max_label = segment_img.max_label + 1
         nonposmin_labels = []
         nmarkers_labels = []
         for label, source_slice, source_deblended in zip(labels,
@@ -279,10 +281,12 @@ def deblend_sources(data, segment_img, npixels, *, labels=None, nlevels=32,
 
             if source_deblended is not None:
                 source_mask = source_deblended > 0
+                new_segm = source_deblended[source_mask]  # min label = 1
                 segm_deblended[source_slice][source_mask] = (
-                    source_deblended[source_mask] + max_label)
-                nlabels = len(_get_labels(source_deblended))
-                max_label += nlabels
+                    new_segm + max_label)
+                new_labels = _get_labels(new_segm) + max_label
+                deblend_label_map[label] = new_labels
+                max_label += len(new_labels)
 
     # process any warnings during deblending
     warning_info = {}
@@ -308,10 +312,15 @@ def deblend_sources(data, segment_img, npixels, *, labels=None, nlevels=32,
             warning_info['nmarkers'] = warn
 
     if relabel:
-        segm_deblended = _relabel_array(segm_deblended, start_label=1)
+        relabel_map = _create_relabel_map(segm_deblended, start_label=1)
+        if relabel_map is not None:
+            segm_deblended = relabel_map[segm_deblended]
+            deblend_label_map = _update_deblend_label_map(deblend_label_map,
+                                                          relabel_map)
 
     segm_img = object.__new__(SegmentationImage)
     segm_img._data = segm_deblended
+    segm_img._deblend_label_map = deblend_label_map
 
     # store the warnings in the output SegmentationImage info attribute
     if warning_info:
@@ -662,7 +671,10 @@ def deblend_source(self):
 
         # markers may not be consecutive if a label was removed due to
         # the contrast criterion
-        return _relabel_array(markers, start_label=1)
+        relabel_map = _create_relabel_map(markers, start_label=1)
+        if relabel_map is not None:
+            markers = relabel_map[markers]
+        return markers
 
 
 def _get_labels(array):
@@ -683,10 +695,12 @@ def _get_labels(array):
     return labels[labels != 0]
 
 
-def _relabel_array(array, start_label=1):
+def _create_relabel_map(array, start_label=1):
     """
-    Relabel an array such that the labels are consecutive integers
-    starting from 1.
+    Create a mapping of original labels to new labels that are
+    consecutive integers.
+
+    By default, the new labels start from 1.
 
     Parameters
     ----------
@@ -698,19 +712,45 @@ def _relabel_array(array, start_label=1):
 
     Returns
     -------
-    relabeled_array : 2D `~numpy.ndarray`
-        The relabeled array.
+    relabel_map : 1D `~numpy.ndarray` or None
+        The array mapping the original labels to the new labels. If the
+        labels are already consecutive starting from ``start_label``,
+        then `None` is returned.
     """
     labels = _get_labels(array)
 
     # check if the labels are already consecutive starting from
     # start_label
     if (labels[0] == start_label
             and (labels[-1] - start_label + 1) == len(labels)):
-        return array
+        return None
 
     # Create an array to map old labels to new labels
     relabel_map = np.zeros(labels.max() + 1, dtype=array.dtype)
     relabel_map[labels] = np.arange(len(labels)) + start_label
 
-    return relabel_map[array]
+    return relabel_map
+
+
+def _update_deblend_label_map(deblend_label_map, relabel_map):
+    """
+    Update the deblend_label_map to reflect the new labels that are
+    consecutive integers.
+
+    Parameters
+    ----------
+    deblend_label_map : dict
+        A dictionary mapping the original labels to the new deblended
+        labels.
+
+    relabel_map : 1D `~numpy.ndarray`
+        The array mapping the original labels to the new labels.
+
+    Returns
+    -------
+    deblend_label_map : dict
+        The updated deblend_label_map.
+    """
+    for old_label, new_labels in deblend_label_map.items():
+        deblend_label_map[old_label] = relabel_map[new_labels]
+    return deblend_label_map

photutils/segmentation/detect.py

@@ -242,6 +242,7 @@ def _detect_sources(data, threshold, npixels, footprint, inverse_mask, *,
         segm._data = segment_img
         segm.__dict__['labels'] = labels
         segm.__dict__['slices'] = segm_slices
+        segm.__dict__['_deblend_label_map'] = {}
         return segm
 
     # this is used by deblend_sources

photutils/segmentation/tests/test_core.py

@@ -477,6 +477,61 @@ def test_patches(self):
         assert isinstance(patches, list)
         assert isinstance(patches[0], Polygon)
 
+    def test_deblended_labels(self):
+        data = np.array([[1, 1, 0, 0, 4, 4],
+                         [0, 0, 0, 0, 0, 4],
+                         [0, 0, 7, 8, 0, 0],
+                         [6, 0, 0, 0, 0, 5],
+                         [6, 6, 0, 5, 5, 5],
+                         [6, 6, 0, 0, 5, 5]])
+        segm = SegmentationImage(data)
+
+        segm0 = segm.copy()
+        assert segm0._deblend_label_map == {}
+        assert segm0.deblended_labels.size == 0
+        assert segm0.deblended_labels_map == {}
+        assert segm0.deblended_labels_inverse_map == {}
+
+        deblend_map = {2: np.array([5, 6]), 3: np.array([7, 8])}
+        segm._deblend_label_map = deblend_map
+        assert_equal(segm._deblend_label_map, deblend_map)
+        assert_equal(segm.deblended_labels, [5, 6, 7, 8])
+        assert segm.deblended_labels_map == {5: 2, 6: 2, 7: 3, 8: 3}
+        assert segm.deblended_labels_inverse_map == deblend_map
+
+        segm2 = segm.copy()
+        segm2.relabel_consecutive()
+        deblend_map = {2: [3, 4], 3: [5, 6]}
+        assert_equal(segm2._deblend_label_map, deblend_map)
+        assert_equal(segm2.deblended_labels, [3, 4, 5, 6])
+        assert segm2.deblended_labels_map == {3: 2, 4: 2, 5: 3, 6: 3}
+        assert_equal(segm2.deblended_labels_inverse_map, deblend_map)
+
+        segm3 = segm.copy()
+        segm3.relabel_consecutive(start_label=10)
+        deblend_map = {2: [12, 13], 3: [14, 15]}
+        assert_equal(segm3._deblend_label_map, deblend_map)
+        assert_equal(segm3.deblended_labels, [12, 13, 14, 15])
+        assert segm3.deblended_labels_map == {12: 2, 13: 2, 14: 3, 15: 3}
+        assert_equal(segm3.deblended_labels_inverse_map, deblend_map)
+
+        segm4 = segm.copy()
+        segm4.reassign_label(5, 50)
+        segm4.reassign_label(7, 70)
+        deblend_map = {2: [50, 6], 3: [70, 8]}
+        assert_equal(segm4._deblend_label_map, deblend_map)
+        assert_equal(segm4.deblended_labels, [6, 8, 50, 70])
+        assert segm4.deblended_labels_map == {50: 2, 6: 2, 70: 3, 8: 3}
+        assert_equal(segm4.deblended_labels_inverse_map, deblend_map)
+
+        segm5 = segm.copy()
+        segm5.reassign_label(5, 50, relabel=True)
+        deblend_map = {2: [6, 3], 3: [4, 5]}
+        assert_equal(segm5._deblend_label_map, deblend_map)
+        assert_equal(segm5.deblended_labels, [3, 4, 5, 6])
+        assert segm5.deblended_labels_map == {6: 2, 3: 2, 4: 3, 5: 3}
+        assert_equal(segm5.deblended_labels_inverse_map, deblend_map)
+
 
 class CustomSegm(SegmentationImage):
     @lazyproperty
@@ -503,5 +558,5 @@ def test_subclass():
                       [70, 70, 0, 0],
                       [70, 70, 0, 1]])
     segm.data = data2
-    assert len(segm.__dict__) == 2
+    assert len(segm.__dict__) == 3
     assert_equal(segm.areas, [1, 2, 2, 4])

photutils/segmentation/tests/test_deblend.py

@@ -52,6 +52,7 @@ def test_deblend_sources(self, mode):
         assert_allclose(len(result.data[mask1]), len(result.data[mask2]))
         assert_allclose(np.sum(self.data[mask1]), np.sum(self.data[mask2]))
         assert_allclose(np.nonzero(self.segm), np.nonzero(result))
+        assert_equal(result.deblended_labels_inverse_map, {1: [1, 2]})
 
     def test_deblend_multiple_sources(self):
         g4 = Gaussian2D(100, 50, 15, 5, 5)
@@ -193,6 +194,8 @@ def test_deblend_sources_norelabel(self, mode):
         result = deblend_sources(self.data, self.segm, self.npixels,
                                  mode=mode, relabel=False, progress_bar=False)
         assert result.nlabels == 2
+        assert_equal(result.labels, [2, 3])
+        assert_equal(result.deblended_labels_inverse_map, {1: [2, 3]})
         assert len(result.slices) <= result.max_label
         assert len(result.slices) == result.nlabels
         assert_allclose(np.nonzero(self.segm), np.nonzero(result))