function LABELING = image_labeling_gui(image,nr_labels)
%IMAGE_LABELING_GUI   Manual labeling of the image
%
% LABELING = IMAGE_LABELING_GUI(IMAGE,NR_LABELS)
%
% IMAGE, input image, may be uint8 or 0-to-1 double, grayscale or rgb
% NR_LABELS, nubmer of not-background labels, default 4
% LABELING, label image with elements from 0 (background) to NR_LABELS
%
% Author: vand@dtu.dk, 2015

if nargin<1
    image = imread('tape.png');
end

if nargin<2
    nr_labels = 4;
end

% DRAWING SETTINGS
thickness_options = [2 5 10 20 50 0]; % last option (value 0) is 'fill'
thickness_initial = 3; % initial pencil thickness 10 pixels
RADIUS = floor(thickness_options(thickness_initial)/2)+0.5; % pencil radius

labels_options = 0:nr_labels; % 0 is background
labels_initial = 2; % initial label is 1
LABEL = labels_options(labels_initial);

nr_circle_pts = 16; % nr points defining a circular pencil
colors = [0.5,0.5,0.5;cool(nr_labels)]; % visualization colors for labels

[image,gray] = normalize_image(image); % rgb og grayscale image
OVERLAY = image; % image overlaid labeling
LABELING = zeros(size(gray)); % labeling

%%%%%%%%%% INITIALIZATION %%%%%%%%%%
fmar = [0.2 0.2]; % discance from screen edge to figure (x and y)
amar = [0.05 0.1]; % distance from figure edge to axes, relative to figure
adim = [0.5*(1-4*amar(1)),1-2*amar(2)]; % axes dimensions
menuw = 0.06; % width of the menu column, centered, smaller than 2*amar(1)

fig = figure('Units','Normalized','Position',[fmar,1-2*fmar],...
    'WindowButtonDownFcn',@start_draw,...
    'WindowButtonMotionFcn',@pointer_motion);
overlay_axes = axes('Units','Normalized','Position',[amar,adim]);
imagesc(OVERLAY), axis image off, hold on
labeling_axes = axes('Units','Normalized','Position',...
    [1-(amar(1)+adim(1)),amar(2),adim]);
imagesc(LABELING,[0,nr_labels]), colormap(colors), axis image off, hold on

uicontrol('Style','text','String','Thickness',...
    'Units','Normalized','Position',[0.5*(1-menuw),0.65,menuw,0.03]);
thickness_options_str = num2cell(thickness_options);
thickness_options_str{end} = 'fill';
uicontrol('Style','popupmenu',...
    'String',thickness_options_str,'Value',thickness_initial,...
    'Units','Normalized','Position',[0.5*(1-menuw),0.62,menuw,0.03],...
    'Callback',{@change_thickness,thickness_options});
uicontrol('Style','text','String','Label',...
    'Units','Normalized','Position',[0.5*(1-menuw),0.55,menuw,0.03]);
uicontrol('Style','popupmenu',...
    'String',num2cell(labels_options),'Value',labels_initial,...
    'Units','Normalized','Position',[0.5*(1-menuw),0.52,menuw,0.03],...
    'Callback',@change_label);
uicontrol('Style','pushbutton','String','Save',...
    'Units','Normalized','Position',[0.5*(1-menuw),0.42,menuw,0.04],...
    'Callback',@save_labeling);

LIMITS = [1 size(gray,2)-0.5 1 size(gray,1)+0.5];
zoom_handle = zoom(fig);
pan_handle = pan(fig);
set(zoom_handle,'ActionPostCallback',@adjust_limits)
set(pan_handle,'ActionPostCallback',@adjust_limits)

XO = [];
uiwait % waits with assigning output until a figure is closed

%%%%%%%%%% CALLBACKS AND HELPING FUNCTIONS %%%%%%%%%%

    function adjust_limits(~,~)
        LIMITS([1,2]) = get(overlay_axes,'XLim');
        LIMITS([3,4]) = get(overlay_axes,'YLim');
    end

    function change_thickness(source,~,thickness_opts)
        RADIUS = floor(thickness_opts(get(source,'Value'))/2)+0.5;
    end

    function change_label(source,~)
        str = get(source,'String');
        LABEL = str2double(str{get(source,'Value')});
    end

    function save_labeling(~,~)
        [file,path] = uiputfile('labeling.mat','Save labeling as');
        if ~isequal(file,0) && ~isequal(path,0)
            save(fullfile(path,file),'LABELING')
        end
    end

    function start_draw(~,~)
        x = get_pixel;
        if is_inside(x)
            if RADIUS>0.5 % thickness>0
                XO = x;
                M = disc(XO,RADIUS,nr_circle_pts,size(gray));
                set(fig,'WindowButtonMotionFcn',@drag_and_draw,...
                    'WindowButtonUpFcn',@end_draw)
            else % fill
                M = bwselect(LABELING==LABELING(x(2),x(1)),x(1),x(2),4);
            end
            update_overlay_and_labeling(M);
        end
    end

    function drag_and_draw(~,~)
        x = get_pixel;
        M = stadium(XO,x,RADIUS,nr_circle_pts,size(gray));
        update_overlay_and_labeling(M);
        XO = x;
    end

    function end_draw(~,~)
        M = stadium(XO,get_pixel,RADIUS,nr_circle_pts,size(gray));
        update_overlay_and_labeling(M);
        P = disc(get_pixel,RADIUS,nr_circle_pts,size(gray));
        overlay_pointer(P);
        XO = [];
        set(fig,'WindowButtonMotionFcn',@pointer_motion,...
            'WindowButtonUpFcn','')
    end

    function pointer_motion(~,~)
        if strcmp(zoom_handle.Enable,'off') && strcmp(pan_handle.Enable,'off')
            x = get_pixel;
            if is_inside(x)
                set(fig,'Pointer','crosshair')
            else
                set(fig,'Pointer','arrow')
            end
            if RADIUS>0.5 % thickness>0
                P = disc(x,RADIUS,nr_circle_pts,size(gray));
                overlay_pointer(P);
            end
        end
    end

    function a = is_inside(x)
        a = inpolygon(x(1),x(2),LIMITS([1,2,2,1]),LIMITS([3,3,4,4]));
    end

    function p = get_pixel
        p = get(overlay_axes,'CurrentPoint');
        p = round(p(1,[1,2]));
    end

    function update_overlay_and_labeling(M)
        if LABEL~=0 % not a background colors
            col = colors(LABEL+1,:);
            w = 0.4; % color weight
            OVERLAY(repmat(M(:),[3,1])) = [col(1)*(w+(1-w)*gray(M(:)));...
                col(2)*(w+(1-w)*gray(M(:)));col(3)*(w+(1-w)*gray(M(:)))];
        else % background restores original image
            OVERLAY(repmat(M(:),[3,1])) = image(repmat(M(:),[3,1]));
        end
        LABELING(M) = LABEL;
        axes(overlay_axes), cla, imagesc(OVERLAY)
        axes(labeling_axes), cla, imagesc(LABELING)
    end

    function overlay_pointer(P)
        shown = OVERLAY;
        shown(repmat(P(:),[3,1])) = 0.5+0.5*OVERLAY(repmat(P(:),[3,1]));
        axes(overlay_axes), cla, imagesc(shown)
    end

    function M = disc(x,r,N,dim)
        % disc shaped mask in the image
        angles = (0:2*pi/N:2*pi*(1-1/N));
        X = x(1)+r*cos(angles);
        Y = x(2)+r*sin(angles);
        M = poly2mask(X,Y,dim(1),dim(2));
    end

    function M = stadium(x1,x2,r,N,dim)
        % stadium shaped mask in the image
        angles = (0:2*pi/N:pi)-atan2(x1(1)-x2(1),x1(2)-x2(2));
        X = [x1(1)+r*cos(angles), x2(1)+r*cos(angles+pi)];
        Y = [x1(2)+r*sin(angles), x2(2)+r*sin(angles+pi)];
        M = poly2mask(X,Y,dim(1),dim(2));
    end

    function [I,G] = normalize_image(I)
        if isa(I,'uint8')
            I = double(I)/255;
        end
        if size(I,3)==3 % rgb image
            G = rgb2gray(I);
        else % assuming grayscale image
            G = I;
            I = repmat(G,[1,1,3]);
        end
    end

end