From 58ef059f114c2345cbf184b948394e89dabd81b2 Mon Sep 17 00:00:00 2001
From: "Peter H. Li" <phli@google.com>
Date: Wed, 10 Jan 2024 17:44:31 -0800
Subject: [PATCH] Move split_disconnected_components from FFN repo to
 connectomics.

PiperOrigin-RevId: 597400829
---
 connectomics/segmentation/labels.py      | 31 ++++++++++++++++++++++++
 connectomics/segmentation/labels_test.py | 15 ++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/connectomics/segmentation/labels.py b/connectomics/segmentation/labels.py
index eb94242..647886d 100644
--- a/connectomics/segmentation/labels.py
+++ b/connectomics/segmentation/labels.py
@@ -19,6 +19,7 @@
 
 import edt
 import numpy as np
+import skimage.measure
 import skimage.morphology
 import skimage.segmentation
 
@@ -202,3 +203,33 @@ def watershed_expand(seg: np.ndarray,
 
   orig_ids, low_ids = zip(*orig_to_low)
   return relabel(ws, np.array(low_ids), np.array(orig_ids)), dist_map
+
+
+def split_disconnected_components(labels: np.ndarray, connectivity=1):
+  """Relabels the connected components of a 3-D integer array.
+
+  Connected components are determined based on 6-connectivity, where two
+  neighboring positions are considering part of the same component if they have
+  identical labels.
+
+  The label 0 is treated specially: all positions labeled 0 in the input are
+  labeled 0 in the output, regardless of whether they are contiguous.
+
+  Connected components of the input array (other than segment id 0) are given
+  consecutive ids in the output, starting from 1.
+
+  Args:
+    labels: 3-D integer numpy array.
+    connectivity: 1, 2, or 3; for 6-, 18-, or 26-connectivity respectively.
+
+  Returns:
+    The relabeled numpy array, same dtype as `labels`.
+  """
+  has_zero = 0 in labels
+  fixed_labels = skimage.measure.label(
+      labels, connectivity=connectivity, background=0)
+  if has_zero or (not has_zero and 0 in fixed_labels):
+    if np.any((fixed_labels == 0) != (labels == 0)):
+      fixed_labels[...] += 1
+      fixed_labels[labels == 0] = 0
+  return np.cast[labels.dtype](fixed_labels)
diff --git a/connectomics/segmentation/labels_test.py b/connectomics/segmentation/labels_test.py
index fd450d1..50ee8c5 100644
--- a/connectomics/segmentation/labels_test.py
+++ b/connectomics/segmentation/labels_test.py
@@ -19,6 +19,21 @@
 import numpy as np
 
 
+class SplitDisconnectedComponentsTest(absltest.TestCase):
+
+  def test_disconnected_components(self):
+    old_labels = np.array([[[2, 2, 3, 4],
+                            [2, 3, 3, 2],
+                            [0, 0, 0, 2],
+                            [0, 3, 3, 3]]], dtype=np.uint64)
+    new_labels = labels.split_disconnected_components(old_labels)
+    expected = np.array([[[1, 1, 2, 3],
+                          [1, 2, 2, 4],
+                          [0, 0, 0, 4],
+                          [0, 5, 5, 5]]], dtype=np.uint64)
+    self.assertTrue(labels.are_equivalent(new_labels, expected))
+
+
 class UtilsTest(absltest.TestCase):
 
   def test_equivalence(self):