Created
January 12, 2016 23:02
-
-
Save vighneshbirodkar/ca5908d513ec7993893a to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment