function [wseg, edges, nseg] = pb2watershed(pbim, min_area, do_fill)
% [wseg, edges, nseg] = pb2watershed(pbim, min_area)
%
% Computes a watershed oversegmentation from the soft boundary map given by pbim
% and removes isolated regions smaller than min_area. Information about the 
% boundaries between regions is also returned.  Fills boundaries if do_fill is set.
%
% INPUT
% pbim(imh, imw): pixel indices provide boundary confidence
% min_area: pixel area of smallest allowed isolated region
% do_fill: set if watershed boundary pixels should be filled 
% 
% OUTPUT
% wseg(imh, imw): pixel indices of regions (uint16)
% edges.(ne, spLR(ne, 1:2), index{ne,1}, ...): number of edges, 
%          regions on each side of boundary, image indices of boundaries, etc.
% nseg: number of image regions
%
% License: This software comes with no warrantees.  Anyone (commercial or not) may 
% use pb2watershed.m free of charge and may make any changes to it at your own risk.
% This software is provided by Derek Hoiem.  No restrictions are attached to its use.
%

if ~exist('min_area', 'var') || isempty(min_area)
  min_area = 25;
end

wseg = watershed(pbim);

%figure(2), hold off, imagesc(wseg==0), axis image, colormap gray

% Compute edge information and remove small regions
[edges, nseg] = getEdgeInfo(wseg);  
if min_area > 0 
  [wseg, edges, nseg] = removeIsolatedRegions(wseg, edges, min_area);
end
edges = addGradientInfo(edges, pbim, etheta);


do_display = false;  
if do_display
  tmpim1 = zeros(size(wseg));
  tmpim2 = zeros(size(wseg));
  for e = 1:edges.ne
    tmpim1(edges.index{e}) = edges.mag_mean(e);
    tmpim2(edges.index{e}) = edges.magdir_mean(e);
  end
  figure(1), hold off, imagesc(tmpim1, [0 1]), axis image, colormap gray
  figure(2), hold off, imagesc(tmpim2, [0 1]), axis image, colormap gray
  figure(4), hold off, imagesc(wseg), axis image, colormap jet
  disp(num2str([nseg edges.ne]))
end

% fill in edge pixels with neighboring region of most similar color
if do_fill
  while (any(wseg(:)==0))
    wseg = fill_in_segmentation(im, wseg, 0, 8);
  end
end
wseg = uint16(wseg);
  



%%%%%%%%% removeIsolatedRegions %%%%%%%%%%%
function [wseg, edges, nseg] = removeIsolatedRegions(wseg, edges, min_area)

stats = regionprops(wseg, 'Area', 'PixelIdxList');
area = [stats.Area]';
indices = {stats.PixelIdxList};

nsurround = zeros(numel(area), 1);
for e = 1:edges.ne
  nsurround(edges.spLR(e, :)) = nsurround(edges.spLR(e, :))+1;
end

rkeep = (area >= min_area) | (nsurround>1);
ekeep = true(edges.ne, 1);
for r = find(~rkeep(:))'  
  ind = edges.spLR(:, 1)==r;
  if any(ind), 
    r2 = edges.spLR(ind, 2);      
  else
    ind = edges.spLR(:, 2)==r;
    r2 = edges.spLR(ind, 1);
  end
  wseg(indices{r}) = r2;
  ekeep(ind) = false;
end
edges.npix = edges.npix(ekeep);
edges.ne = sum(ekeep);
edges.index = edges.index(ekeep);
edges.theta = edges.theta(ekeep);
rnums = cumsum(rkeep);
wseg(wseg>0) = rnums(wseg(wseg>0));
edges.spLR = rnums(edges.spLR);
nseg = rnums(end);


%%%%%%%%% addGradientInfo %%%%%%%%%%%
function edges = addGradientInfo(edges, esoft, etheta);
edges.mag_max = zeros(edges.ne, 1);
edges.mag_mean = zeros(edges.ne, 1);
edges.magdir_max = zeros(edges.ne, 1);
edges.magdir_mean = zeros(edges.ne, 1);
for e = 1:edges.ne
  emag = esoft(edges.index{e});
  edges.mag_max(e) = max(emag); % / edges.npix(e);
  edges.magdir_max(e) = max(emag.*cos(etheta(edges.index{e})-edges.theta(e))); %/edges.npix(e);
  edges.mag_mean(e) = sum(emag) / edges.npix(e);
  edges.magdir_mean(e) = sum(emag.*cos(etheta(edges.index{e})-edges.theta(e)))/edges.npix(e);
end
edges.magdir_max = abs(edges.magdir_max);
edges.magdir_mean = abs(edges.magdir_mean);


%%%%%%%%% getEdgeInfo %%%%%%%%%%%
function [edges, nseg] = getEdgeInfo(seg)

nseg = max(seg(:));
seg = negpad(seg);
[nrows,ncols] = size(seg);
eind = find(seg==0);
ne = numel(eind);

%neighbor_offset = [-1 1 -nrows nrows];
%neighbor_index = eind*ones(1,4) + ones(ne,1)*neighbor_offset;
neighbor_offset = [-1 1 -nrows nrows -1-nrows -1+nrows 1-nrows 1+nrows];
neighbor_index = eind*ones(1,8) + ones(ne,1)*neighbor_offset;

neighbors = seg(neighbor_index);

% set neighbors to the neighboring region with the smallest and largest
% indices (in case there are multiple or multiple neighboring pixels from
% the same region)
tmpn = neighbors;
tmpn(tmpn<=0) = Inf;
n1 = min(tmpn, [], 2);
n2 = max(neighbors, [], 2);

[eval, sind] = sort(n1+(n2-1)*nseg);

% tmpim = zeros(size(seg));
% tmpim(eind(sind)) = n2;
% figure(2), hold off, imagesc(tmpim); axis image, colormap jet;
% pause;

%eind = eind(sind)-nrows-1;
[y, x] = ind2sub([nrows ncols], eind);
eind = (y-1) + (x-2)*(nrows-2);


%%% Get the indices for each edge (this code is complicated but fast)
deval = [double(eval(2:end)~=eval(1:end-1)) ; 1]; % marks the last index of each edge
ne=sum(deval);  %sum(deval)+1;
nepix = zeros(ne, 1);
e = 1; c = 0;
% count number of pixels in each edge
for k = 1:numel(eval)
  c = c+1;  % increment count of number of pixels per edge
  nepix(e) = c; % set npix to current count  
  c = c*(1-deval(k)); % reset c if dsval(k) = 1; (new edge)    
  e = e + deval(k); % increment e if dval(k)=1;  
end
edges.npix = nepix;
edges.ne = ne;

% get left and right side superpixels for each boundary
deval = logical(deval);
edges.spLR = [n1(sind(deval)) n2(sind(deval))];

% assign edge indices to cell array and get angles
eind = uint32(eind(sind)); %(sval));
edges.index = cell(edges.ne, 1);
edges.theta = zeros(edges.ne, 1);
k = 0;

[imx, imy] = meshgrid(1:ncols, 1:nrows);
for e = 1:edges.ne  
  % assign indices
  edges.index{e} = eind(k+(1:nepix(e)));    
  k = k + nepix(e);

  % get angles           
  y = imy(edges.index{e});
  x = imx(edges.index{e});    
  zmx = (x-sum(x)/nepix(e));
  zmy = -(y-sum(y)/nepix(e));
  D = [sum(zmx.*zmx) sum(zmx.*zmy); sum(zmx.*zmy) sum(zmy.*zmy)];
  [v, lambda] = eig(D);
  edges.theta(e, 1) = atan2(v(2, 2) , v(1, 2));  
end

%%%% This is a much simpler, but 400x slower version of getting edge indices
% uval = unique(eval);
% ne = numel(uval);
% edges = cell(ne, 1);
% for k = 1:ne    
%     edges{k} = uint32(eind(eval==uval(k)));
% end



% tmpim = -seg(2:end-1, 2:end-1);
% rp = randperm(ne);
% for k = 1:ne
%   tmpim(edges{k}) = k;
% end
% figure(4), hold off, imagesc(tmpim), axis image, colormap jet





%%%%%%%%% negpad %%%%%%%%%%%
function B = negpad(A)
[nrows,ncols,nchannels] = size(A);
B = -ones(nrows+2,ncols+2,nchannels, class(A));
B(2:end-1,2:end-1,:) = A;