In [1]:
from PIL import Image
from numpy import *
from pylab import *
In [2]:
# The data is still available from web archive
# https://web.archive.org/web/20161203110733/http://research.microsoft.com/en-us/um/cambridge/projects/visionimagevideoediting/segmentation/grabcut.htm
In [3]:
from scipy.misc import imresize
import graphcut
graphcut = reload(graphcut)
In [4]:
im = array(Image.open('music.JPG'))
In [5]:
scale = 0.1
im = imresize(im, scale, interp='bilinear')
In [6]:
figure()
imshow(im)
axis('off')
show()
print im.shape
In [7]:
labels = zeros(im.shape[:2])
labels[10:35, 20:40] = 1
labels[:, :10] = -1
labels[:, -8:] = -1
figure()
graphcut.show_labeling(im, labels)
show()
In [8]:
figure()
gray()
subplot(1, 3, 1)
imshow(im)
axis('off')
subplot(1, 3, 2)
imshow(labels)
axis('off')
show()
In [9]:
from pygraph.classes.digraph import digraph
from pygraph.algorithms.minmax import maximum_flow
import bayes
def build_bayes_graph(im, labels, sigma=1e-2, kappa=2, weight=1):
""" Build a graph from 4-neighborhood of pixels.
Foregraound and background is determined from
labels (1 for foreground, -1 for background, 0 othewise)
and is modeled with naive Bayes classifiers. """
m, n = im.shape[:2]
# RGB vector version (one pixel per row)
vim = im.astype('float')
vim = vim.reshape((-1, 3))
# RGB for foreground and background
foreground = im[labels == 1].reshape((-1, 3))
background = im[labels == -1].reshape((-1, 3))
train_data = [foreground, background]
# train naive Bayes classifier
bc = bayes.BayesClassifier()
bc.train(train_data, labels)
# get probabilities for all pixels
bc_lables, prob = bc.classify(vim)
prob_fg = prob[0]
prob_bg = prob[1]
# create graph with m*n+2 nodes
gr = digraph()
gr.add_nodes(range(m*n+2))
source = m*n # second to last is source
sink = m*n+1 # last node is sink
# normalize
pos = m*n/2-100
for i in range(vim.shape[0]):
vim[i] = vim[i] / linalg.norm(vim[i])
# go through all nodes and add edges
lb = labels.copy()
lb = lb.flatten()
for i in range(m*n):
if weight>0:
if lb[i]==1:
gr.add_edge((source, i), wt=weight)
elif lb[i]==-1:
gr.add_edge((i, sink), wt=weight)
else:
# add edge from source
gr.add_edge((source, i), wt=(prob_fg[i]/(prob_fg[i] + prob_bg[i])))
# add edge to sink
gr.add_edge((i, sink), wt=(prob_bg[i]/(prob_fg[i] + prob_bg[i])))
else:
if (prob_fg[i]>prob_bg[i]):
gr.add_edge((source, i), wt=(prob_fg[i]/(prob_fg[i] + prob_bg[i])))
else:
# add edge to sink
gr.add_edge((i, sink), wt=(prob_bg[i]/(prob_fg[i] + prob_bg[i])))
# add edges to neighbors
if i % n != 0: # left exists
edge_wt = kappa*exp(-1.0*sum((vim[i] - vim[i-1])**2)/sigma)
gr.add_edge((i, i-1), wt=edge_wt)
if (i+1) % n != 0: # right exists
edge_wt = kappa*exp(-1.0*sum((vim[i] - vim[i+1])**2)/sigma)
gr.add_edge((i, i+1), wt=edge_wt)
if i//n != 0: # up exists
edge_wt = kappa*exp(-1.0*sum((vim[i] - vim[i-n])**2)/sigma)
gr.add_edge((i, i-n), wt=edge_wt)
if i//n != m-1: # down exists
edge_wt = kappa*exp(-1.0*sum((vim[i] - vim[i+n])**2)/sigma)
gr.add_edge((i, i+n), wt=edge_wt)
return gr
In [10]:
g = build_bayes_graph(im, labels, sigma=1e-2, kappa=2, weight=100)
In [11]:
res = graphcut.cut_graph(g, im.shape[:2])
In [12]:
figure()
graphcut.show_labeling(im, res)
show()
In [13]:
seg = array(Image.open('music.bmp'))
seg = imresize(seg, 0.1, interp='nearest')
def score_match(rin, sin):
r = rin.flatten().copy()
s = sin.flatten().copy()
s[s==0] = 1
s[s==255] = 0
score = sum(1.0*(r==s))/r.shape[0]
return score
print score_match(res, seg)
In [14]:
for i in range(4):
lb = res.copy()
lb[lb==1] = -1
lb[lb==0] = 1
lb[lb==-1] = 0
lb[:8, :] = -1
lb[-8:, :] = -1
lb[:, :2] = -1
lb[:, -4:] = -1
g = build_bayes_graph(im, lb, sigma=1e-2, kappa=2, weight=0)
res = graphcut.cut_graph(g, im.shape[:2])
figure()
graphcut.show_labeling(im, res)
show()
print score_match(res, seg)
In [15]:
figure()
graphcut.show_labeling(im, lb)
show()
In [68]:
# This doesn't work well
In [ ]: