diff --git a/skimage/graph/graph_cut.py b/skimage/graph/graph_cut.py index a255d01e..27e4400d 100644 --- a/skimage/graph/graph_cut.py +++ b/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] diff --git a/skimage/graph/rag.py b/skimage/graph/rag.py index 8bcb06b9..e2f60195 100644 --- a/skimage/graph/rag.py +++ b/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):