From bcbfdb28f59170c2c96aea9444c093e9556f6df2 Mon Sep 17 00:00:00 2001
From: Vighnesh Birodkar <vighneshbirodkar@gmail.com>
Date: Thu, 14 Aug 2014 01:20:55 +0530
Subject: [PATCH] add desaturate option

---
 doc/examples/plot_rag_draw.py |  4 ++--
 skimage/graph/rag.py          | 38 ++++++++++++++++++++++-------------
 2 files changed, 26 insertions(+), 16 deletions(-)

diff --git a/doc/examples/plot_rag_draw.py b/doc/examples/plot_rag_draw.py
index 21dbbc28..61d7e134 100644
--- a/doc/examples/plot_rag_draw.py
+++ b/doc/examples/plot_rag_draw.py
@@ -19,8 +19,8 @@ plt.figure()
 plt.title("RAG with all edges shown in green.")
 plt.imshow(out)
 
-cmap = colors.ListedColormap(['cyan', 'red'])
-out = graph.draw_rag(labels, g, img, colormap=cmap, thresh=30)
+cmap = colors.ListedColormap(['#00ff00', '#ff0000'])
+out = graph.draw_rag(labels, g, img,colormap=cmap, thresh=30, desaturate=True)
 plt.figure()
 plt.title("RAG with edge weights less than 30, color "
           "mapped between green and red.")
diff --git a/skimage/graph/rag.py b/skimage/graph/rag.py
index e3849184..e1ab79ca 100644
--- a/skimage/graph/rag.py
+++ b/skimage/graph/rag.py
@@ -14,7 +14,7 @@ import numpy as np
 from scipy.ndimage import filters
 from scipy import ndimage as nd
 import math
-from .. import draw, measure, segmentation, util
+from .. import draw, measure, segmentation, util, color
 try:
     from matplotlib import colors
     from matplotlib import cm
@@ -246,8 +246,9 @@ def rag_mean_color(image, labels, connectivity=2, mode='distance',
     return graph
 
 
-def draw_rag(labels, rag, img, border_color=None, node_color='yellow',
-             edge_color='green', colormap=None, thresh=np.inf):
+def draw_rag(labels, rag, img, border_color=None, node_color='#ffff00',
+             edge_color='#00ff00', colormap=None, thresh=np.inf,
+             desaturate=False, in_place=True):
     """Draw a Region Adjacency Graph on an image.
 
     Given a labelled image and its corresponding RAG, draw the nodes and edges
@@ -256,11 +257,11 @@ def draw_rag(labels, rag, img, border_color=None, node_color='yellow',
 
     Parameters
     ----------
-    labels : ndarray, shape(M, N)
+    labels : ndarray, shape (M, N)
         The labelled image.
     rag : RAG
         The Region Adjacency Graph.
-    img : ndarray, shape(M, N, 3)
+    img : ndarray, shape (M, N, 3)
         Input image.
     border_color : colorspec, optional
         Any matplotlib colorspec.
@@ -274,10 +275,16 @@ def draw_rag(labels, rag, img, border_color=None, node_color='yellow',
     thresh : float, optional
         Edges with weight below `thresh` are not drawn, or considered for color
         mapping.
+    desaturate : bool, optional
+        Convert the image to grayscale before displaying. Particularly helps
+        visualiztion when using the `colormap` option.
+    in_place : bool, optional
+        If set, the RAG is modified in place. For each node `n` the function
+        will set a new attribute ``rag.node[n]['centroid']``.
 
     Returns
     -------
-    out : ndarray, shape(M, N, [..., P,] 3)
+    out : ndarray, shape (M, N, 3)
         The image with the RAG drawn.
 
     Examples
@@ -288,7 +295,13 @@ def draw_rag(labels, rag, img, border_color=None, node_color='yellow',
     >>> g =  graph.rag_mean_color(img, labels)
     >>> out = graph.draw_rag(labels, g, img)
     """
-    rag = rag.copy()
+    if not in_place:
+        rag = rag.copy()
+
+    if desaturate:
+        img = color.rgb2gray(img)
+        img = color.gray2rgb(img)
+
     out = util.img_as_float(img)
     cc = colors.ColorConverter()
 
@@ -305,11 +318,10 @@ def draw_rag(labels, rag, img, border_color=None, node_color='yellow',
         offset += 1
 
     rag_labels = map_array[labels]
-
     regions = measure.regionprops(rag_labels)
-    for region in regions:
-        # Because we kept the offset as 1
-        rag.node[region['label'] - 1]['centroid'] = region['centroid']
+
+    for (n, data), region in zip(rag.nodes_iter(data=True), regions):
+        data['centroid'] = region['centroid']
 
     if border_color is not None:
         border_color = cc.to_rgb(border_color)
@@ -328,12 +340,10 @@ def draw_rag(labels, rag, img, border_color=None, node_color='yellow',
             continue
         r1, c1 = map(int, rag.node[n1]['centroid'])
         r2, c2 = map(int, rag.node[n2]['centroid'])
-
         line = draw.line(r1, c1, r2, c2)
 
         if colormap is not None:
-            current_color = smap.to_rgba([data['weight']])[0][:-1]
-            out[line] = current_color
+            out[line] = smap.to_rgba([data['weight']])[0][:-1]
         else:
             out[line] = edge_color