% function labels = mygraphcut()

%  revised from the folloing copyright by Kecheng Yang

% TEST2 Demonstrates gradient-based image
%   min-cut partitioning.
%
%   (c) 2008 Michael Rubinstein, WDI R&D and IDC
%   $Revision: 140 $
%   $Date: 2008-09-15 15:35:01 -0700 (Mon, 15 Sep 2008) $
%

clear all;

im = imread('source.jpg');

imshow(im);
mask = imfreehand();
front = mask.createMask();

% imshow(im);
% mask = imfreehand();
% back = mask.createMask();


m = double(rgb2gray(im));
[height,width] = size(m);

% r = double(im(:,:,1));
% g = double(im(:,:,2));
% b = double(im(:,:,3));

disp('building graph');
N = height*width;

unary = 9e9;

% construct graph
E = edges4connected(height,width);
figdis = gmdistribution.fit(m(:),1);
V = exp(-((m(E(:,1))-m(E(:,2))).^2)/(2*figdis.Sigma*figdis.Sigma));
A = sparse(E(:,1),E(:,2),V,N,N,4*N);

% terminal weights
% connect source to leftmost column.
% connect rightmost column to target.
% T = sparse([1:height;N-height+1:N]',[ones(height,1);ones(height,1)*2],ones(2*height,1)*9e9);
T = sparse(N,2);

for y = 1:height
    for x = 1:width
%         if back(y,x)
        if y==1 || y==height || x==1 || x==height
            T((x-1)*height+y,1)=unary;
        end
        
        if front(y,x)
            T((x-1)*height+y,2)=unary;
        end
    end
end


% T(1,1)=9e9;
% T(N/2,2)=9e9;

% T = sparse([1;height*400+270],[1;2],ones(2,1)*9e9,N,2);

disp('calculating maximum flow');

[flow,labels] = maxflow(A,T);
labels = reshape(labels,[height width]);

old_labels = int32(front);


% while sum(sum(labels)) ~= sum(sum(old_labels))
% for i =1:3
%     for x = 1:width
%         for y = 1:height       
%             i1 = zeros(N,1);
%             i2 = zeros(N,1);
%             v = zeros(N,1);
%             i1((x-1)*height+y) = (x-1)*height+y;
%             v((x-1)*height+y) = 9e9;
%             if labels(y,x) > 0 
%                 i2((x-1)*height+y) = 2;
%             else
%                 i2((x-1)*height+y) = 1;
%             end
%         end
%     end
%     T = sparse(i1,i2,v,N,2);
%     old_labels = labels;
%     [flow,labels] = maxflow(A,T);
%     labels = reshape(labels,[height width]);
% end

count = 0;

while sum(labels(:)) ~= sum(old_labels(:))
    count = count + 1;
    fi = find(labels(:)==1);
    bi = find(labels(:)==0);
    fg = m(fi);
    bg = m(bi);
    fgdis = gmdistribution.fit(fg,1);
    bgdis = gmdistribution.fit(bg,1);
    
    tmp = m(:);
    unary_v = -log(fgdis.pdf(tmp)./bgdis.pdf(tmp));
    unary_b = unary_v;
    unary_v = abs(unary_v);
    pos = find(unary_b >= 0);
    neg = find(unary_b < 0);
    unary_b(pos) = 1;
    unary_b(neg) = 2;
    
    i(1:N,1)=[1:N];
    v(1:N,1)=unary_v;
    j(1:N,1)=unary_b;
    T = sparse(i,j,v,N,2);
    old_labels = labels;
    [flow,labels] = maxflow(A,T);
    labels = reshape(labels,[height width]);
end

imagesc(labels); title('labels');

bMask = logical(labels);
