diff --git a/skimage/future/graph/rag.py b/skimage/future/graph/rag.py index dd45f2fc..ab31171e 100644 --- a/skimage/future/graph/rag.py +++ b/skimage/future/graph/rag.py @@ -12,6 +12,45 @@ except ImportError: pass +def _edge_generator_from_csr(csr_matrix): + """Yield weighted edge triples for use by NetworkX from a CSR matrix. + + This function is a straight rewrite of + `networkx.convert_matrix._csr_gen_triples`. Since that is a private + function, it is safer to include our own here. + + Parameters + ---------- + csr_matrix : scipy.sparse.csr_matrix + The input matrix. An edge (i, j, w) will be yielded if there is a + data value for coordinates (i, j) in the matrix, even if that value + is 0. + + Yields + ------ + i, j, w : (int, int, float) tuples + Each value `w` in the matrix along with its coordinates (i, j). + + Examples + -------- + + >>> dense = np.eye(2, dtype=np.float) + >>> csr = sparse.csr_matrix(dense) + >>> edges = _edge_generator_from_csr(csr) + >>> type(edges) + generator + >>> list(edges) + [(0, 0, 1.0), (1, 1, 1.0)] + """ + nrows = csr_matrix.shape[0] + values = csr_matrix.data + indptr = csr_matrix.indptr + col_indices = csr_matrix.indices + for i in range(nrows): + for j in range(indptr[i], indptr[i + 1]): + yield i, col_indices[j], data[j] + + def min_weight(graph, src, dst, n): """Callback to handle merging nodes by choosing minimum weight. @@ -365,9 +404,8 @@ def rag_boundary(labels, edge_map, connectivity=2): graph_matrix.data /= count_matrix.data rag = RAG() - rows, cols = graph_matrix.nonzero() - graph_data = zip(rows, cols, graph_matrix.data) - rag.add_weighted_edges_from(graph_data, attr='weight') + rag.add_weighted_edges_from(_edge_generator_from_csr(graph_matrix), + weight='weight') for n in rag.nodes(): rag.node[n].update({'labels': [n]})