From 58237a558a457ef7a32c67b99c2f81aa2b04e89a Mon Sep 17 00:00:00 2001
From: Andreas Mueller <amueller@ais.uni-bonn.de>
Date: Sun, 17 Jun 2012 20:52:45 +0200
Subject: [PATCH] ENH dirty fix, works though. Starting profiling.

---
 skimage/segmentation/__init__.py      |  4 +--
 skimage/segmentation/felzenszwalb.pyx | 52 +++++++++++++++++++++++++--
 2 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/skimage/segmentation/__init__.py b/skimage/segmentation/__init__.py
index 0ea91444..8fa8dfb8 100644
--- a/skimage/segmentation/__init__.py
+++ b/skimage/segmentation/__init__.py
@@ -1,5 +1,5 @@
 from .random_walker_segmentation import random_walker
-#from .felzenszwalb import felzenszwalb_segmentation
+from .felzenszwalb import felzenszwalb_segmentation
 from .quickshift import quickshift
 
-__all__ = [random_walker, quickshift]
+__all__ = [random_walker, quickshift, felzenszwalb_segmentation]
diff --git a/skimage/segmentation/felzenszwalb.pyx b/skimage/segmentation/felzenszwalb.pyx
index a53bf70e..11f3e93b 100644
--- a/skimage/segmentation/felzenszwalb.pyx
+++ b/skimage/segmentation/felzenszwalb.pyx
@@ -5,7 +5,49 @@ import scipy
 
 #from ..util import img_as_float
 #from ..color import rgb2grey
-from skimage.morphology.ccomp cimport find_root, join_trees
+#from skimage.morphology.ccomp cimport find_root, join_trees
+
+DTYPE = np.int
+ctypedef np.int_t DTYPE_t
+
+cdef DTYPE_t find_root(np.int_t *forest, np.int_t n):
+    """Find the root of node n.
+
+    """
+    cdef np.int_t root = n
+    while (forest[root] < root):
+        root = forest[root]
+    return root
+
+cdef set_root(np.int_t *forest, np.int_t n, np.int_t root):
+    """
+    Set all nodes on a path to point to new_root.
+
+    """
+    cdef np.int_t j
+    while (forest[n] < n):
+        j = forest[n]
+        forest[n] = root
+        n = j
+
+    forest[n] = root
+
+
+cdef join_trees(np.int_t *forest, np.int_t n, np.int_t m):
+    """Join two trees containing nodes n and m.
+
+    """
+    cdef np.int_t root = find_root(forest, n)
+    cdef np.int_t root_m
+
+    if (n != m):
+        root_m = find_root(forest, m)
+
+        if (root > root_m):
+            root = root_m
+
+        set_root(forest, n, root)
+        set_root(forest, m, root)
 
 
 def felzenszwalb_segmentation(image, k, sigma=0.8):
@@ -37,7 +79,7 @@ def felzenszwalb_segmentation(image, k, sigma=0.8):
     edge_queue = np.argsort(costs)
     cdef np.int_t *segments_p = <np.int_t*>segments.data
     cdef np.int_t seg_new
-    segment_size = defaultdict(lambda: 1)
+    cdef np.ndarray[np.int_t, ndim=1] segment_size = np.ones(width * height, dtype=np.int)
     # inner cost of segments
     cint = defaultdict(lambda: 0)
     for edge, cost in zip(edges[edge_queue], costs[edge_queue]):
@@ -53,4 +95,8 @@ def felzenszwalb_segmentation(image, k, sigma=0.8):
             seg_new = find_root(segments_p, seg0)
             segment_size[seg_new] = segment_size[seg0] + segment_size[seg1]
             cint[seg_new] = cost
-    return segments
+    # unravel the union find tree
+    old = np.zeros_like(flat)
+    while (old != flat).any():
+        old = flat
+        flat = flat[flat]