colormap.py

from skimage import data, color, util, measure, segmentation
from matplotlib import pyplot as plt
from skimage.future import graph
from matplotlib import colors
from matplotlib import cm
from skimage.util.colormap import viridis

import numpy as np


def draw_rag(rag, img, labels, node_color='#00ff00', edge_color='#ff0000',
             border_color="#000000", colormap='afmhot'):
    plt.figure()

    img = color.rgb2gray(img)

    out = util.img_as_float(img)
    cc = colors.ColorConverter()

    edge_color = cc.to_rgb(edge_color)
    node_color = cc.to_rgb(node_color)

    # Handling the case where one node has multiple labels
    # offset is 1 so that regionprops does not ignore 0
    offset = 1
    map_array = np.arange(labels.max() + 1)
    for n, d in rag.nodes_iter(data=True):
        for label in d['labels']:
            map_array[label] = offset
        offset += 1

    rag_labels = map_array[labels]
    regions = measure.regionprops(rag_labels)

    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)
        out = segmentation.mark_boundaries(out, rag_labels, color=border_color)

    if colormap is not None:
        edge_weight_list = [d['weight'] for x, y, d in rag.edges_iter(data=True)]
        norm = colors.Normalize()
        norm.autoscale(edge_weight_list)
        smap = cm.ScalarMappable(norm, colormap)
    
    plt.imshow(out)
    for n1, n2, data in rag.edges_iter(data=True):

        r1, c1 = map(int, rag.node[n1]['centroid'])
        r2, c2 = map(int, rag.node[n2]['centroid'])
        col = smap.to_rgba([data['weight']])[0][:-1]
        plt.gca().add_artist(plt.Line2D([c1, c2], [r1, r2],  color=col, lw='2.5'))
        #line = draw.line(r1, c1, r2, c2)

        #if colormap is not None:
        
        #else:
        #    out[line] = edge_color

        #circle = draw.circle(r1, c1, 2)
        #out[circle] = node_color

img = data.coffee()
labels = segmentation.slic(img, compactness=30, n_segments=400)
rag = graph.rag_mean_color(img, labels)
draw_rag(rag, img, labels)
plt.show()