Used mapping in graph_cut.py

skimage/graph/graph_cut.py

@@ -2,7 +2,7 @@ import networkx as nx
 import numpy as np
 
 
-def threshold_cut(label, rag, thresh):
+def threshold_cut(label_image, rag, thresh):
     """Combines regions seperated by weight less than threshold.
 
     Given an image's labels and its RAG, outputs new labels by
@@ -11,7 +11,7 @@ def threshold_cut(label, rag, thresh):
 
     Parameters
     ----------
-    label : (width, height) or (width, height, 3) ndarray
+    label_image : (width, height) or (width, height, 3) ndarray
         The array of labels.
     rag : RAG
         The region adjacency graph.
@@ -45,11 +45,11 @@ def threshold_cut(label, rag, thresh):
     rag.remove_edges_from(to_remove)
 
     comps = nx.connected_components(rag)
-    out = np.copy(label)
 
+    map_array = np.arange(label_image.max()+1, dtype = np.int)
     for i, nodes in enumerate(comps):
         for node in nodes:
-            for l in rag.node[node]['labels']:
-                out[label == l] = i
+            for label in rag.node[node]['labels']:
+                map_array[label] = i
 
-    return out
+    return map_array[label_image]

skimage/graph/rag.py

@@ -13,12 +13,12 @@ class RAG(nx.Graph):
     between their corresponding nodes.
     """
 
-    def merge_nodes(self, i, j):
+    def merge_nodes(self, i, j, function=max):
         """Merge node `i` into `j`.
 
         The new combined node is adjacent to all the neighbors of `i`
-        and `j`. In case of conflicting edges, edge with higher weight
-        is chosen.
+        and `j`. In case of conflicting edges the given function is
+        called.
 
         Parameters
         ----------
@@ -26,6 +26,9 @@ class RAG(nx.Graph):
             Node to be merged.
         j : int
             Node to be merged.
+        function : callable, optional
+            Function to decide which edge weight to keep when a node is
+            adjacent to both `i` and `j`.
 
         """
         for x in self.neighbors(i):