Merge pull request #496 from ahojnnes/label-performance

ENH: Inline helper functions of label and add number of regions as an optional return.

skimage/morphology/ccomp.pxd

@@ -5,6 +5,6 @@ DTYPE = cnp.intp
 ctypedef cnp.intp_t DTYPE_t
 
 cdef DTYPE_t find_root(DTYPE_t *forest, DTYPE_t n)
-cdef set_root(DTYPE_t *forest, DTYPE_t n, DTYPE_t root)
-cdef join_trees(DTYPE_t *forest, DTYPE_t n, DTYPE_t m)
-cdef link_bg(DTYPE_t *forest, DTYPE_t n, DTYPE_t *background_node)
+cdef void set_root(DTYPE_t *forest, DTYPE_t n, DTYPE_t root)
+cdef void join_trees(DTYPE_t *forest, DTYPE_t n, DTYPE_t m)
+cdef void link_bg(DTYPE_t *forest, DTYPE_t n, DTYPE_t *background_node)

skimage/morphology/ccomp.pyx

@@ -39,7 +39,7 @@ cdef DTYPE_t find_root(DTYPE_t *forest, DTYPE_t n):
     return root
 
 
-cdef set_root(DTYPE_t *forest, DTYPE_t n, DTYPE_t root):
+cdef inline void set_root(DTYPE_t *forest, DTYPE_t n, DTYPE_t root):
     """
     Set all nodes on a path to point to new_root.
 
@@ -53,7 +53,7 @@ cdef set_root(DTYPE_t *forest, DTYPE_t n, DTYPE_t root):
     forest[n] = root
 
 
-cdef join_trees(DTYPE_t *forest, DTYPE_t n, DTYPE_t m):
+cdef inline void join_trees(DTYPE_t *forest, DTYPE_t n, DTYPE_t m):
     """Join two trees containing nodes n and m.
 
     """
@@ -70,7 +70,7 @@ cdef join_trees(DTYPE_t *forest, DTYPE_t n, DTYPE_t m):
         set_root(forest, m, root)
 
 
-cdef link_bg(DTYPE_t *forest, DTYPE_t n, DTYPE_t *background_node):
+cdef inline void link_bg(DTYPE_t *forest, DTYPE_t n, DTYPE_t *background_node):
     """
     Link a node to the background node.
 
@@ -80,9 +80,9 @@ cdef link_bg(DTYPE_t *forest, DTYPE_t n, DTYPE_t *background_node):
 
     join_trees(forest, n, background_node[0])
 
-# Connected components search as described in Fiorio et al.
 
-def label(input, DTYPE_t neighbors=8, DTYPE_t background=-1):
+# Connected components search as described in Fiorio et al.
+def label(input, DTYPE_t neighbors=8, DTYPE_t background=-1, return_num=False):
     """Label connected regions of an integer array.
 
     Two pixels are connected when they are neighbors and have the same value.
@@ -111,6 +111,9 @@ def label(input, DTYPE_t neighbors=8, DTYPE_t background=-1):
     labels : ndarray of dtype int
         Labeled array, where all connected regions are assigned the
         same integer value.
+    num : int, optional
+        Number of labels, which equals the maximum label index and is only
+        returned if return_num is `True`.
 
     Examples
     --------
@@ -202,7 +205,6 @@ def label(input, DTYPE_t neighbors=8, DTYPE_t background=-1):
                 join_trees(forest_p, i*cols + j, i*cols + j - 1)
 
     # Label output
-
     cdef DTYPE_t ctr = 0
     for i in range(rows):
         for j in range(cols):
@@ -216,6 +218,9 @@ def label(input, DTYPE_t neighbors=8, DTYPE_t background=-1):
 
     # Work around a bug in ndimage's type checking on 32-bit platforms
     if data.dtype == np.int32:
-        return data.view(np.int32)
+        data = data.view(np.int32)
+
+    if return_num:
+        return data, ctr
     else:
         return data

skimage/morphology/tests/test_ccomp.py

@@ -82,6 +82,14 @@ class TestConnectedComponents:
                             [-1,  0, -1],
                             [-1, -1, -1]])
 
+    def test_return_num(self):
+        x = np.array([[1, 0, 6],
+                      [0, 0, 6],
+                      [5, 5, 5]])
+
+        assert_array_equal(label(x, return_num=True)[1], 4)
+        assert_array_equal(label(x, background=0, return_num=True)[1], 3)
+
 
 if __name__ == "__main__":
     run_module_suite()