Merge pull request #1986 from larrybradley/segm_labels-dtype

Ensure labels dtype matches the segmentation data

CHANGES.rst

@@ -38,6 +38,9 @@ Bug Fixes
     representation instead of ``str`` representation when using NumPy
     2.0+. [#1956]
 
+  - Fixed a bug to ensure that the dtype of the ``SegmentationImage``
+    ``labels`` always matches the image dtype. [#1986]
+
 API Changes
 ^^^^^^^^^^^
 

photutils/segmentation/core.py

@@ -105,7 +105,7 @@ def _get_labels(data):
         :meth:`remove_masked_labels` on a masked version of the
         segmentation array.
         """
-        # np.unique also sorts elements
+        # np.unique preserves dtype and also sorts elements
         return np.unique(data[data != 0])
 
     @lazyproperty
@@ -209,7 +209,7 @@ def labels(self):
             for label, slc in zip(labels_all, self._raw_slices, strict=True):
                 if slc is not None:
                     labels.append(label)
-            return np.array(labels)
+            return np.array(labels, dtype=self._data.dtype)
 
         return self._get_labels(self.data)
 

photutils/segmentation/detect.py

@@ -202,7 +202,7 @@ def _detect_sources(data, threshold, npixels, footprint, inverse_mask, *,
     # NOTE: recasting segment_img to int and using output=segment_img
     # gives similar performance
     segment_img, nlabels = ndi_label(segment_img, structure=footprint)
-    labels = np.arange(nlabels) + 1
+    labels = np.arange(nlabels, dtype=segment_img.dtype) + 1
 
     # remove objects with less than npixels
     # NOTE: making cutout images and setting their pixels to 0 is

photutils/segmentation/tests/test_detect.py

@@ -146,6 +146,9 @@ def test_detection(self):
         segm = detect_sources(self.data, threshold=0.9, npixels=2)
         assert_equal(segm.data, self.refdata)
 
+        assert segm.data.dtype == np.int32
+        assert segm.labels.dtype == np.int32
+
         segm = detect_sources(self.data << u.uJy, threshold=0.9 * u.uJy,
                               npixels=2)
         assert_equal(segm.data, self.refdata)