Jump Flooding Algorithm for generating voronoi diagram

jfa.m

function B = jfa(n, w, h, nsamples)

sites = round([rand(n, 1)*(w-1)+1, rand(n,1)*(h-1)+1]);
rawsites = sites;

if nargin < 4
    nsamples = 16;
end
offsets=zeros(nsamples, 2);
sqrtNSamples = round(sqrt(nsamples));
jstep = 1.0 / sqrtNSamples;
for i=1:sqrtNSamples
    for j=1:sqrtNSamples
        offsets((i-1)*sqrtNSamples+j, :) = [(i-1+rand())*jstep, (j-1+rand())*jstep];
    end
end
B = zeros(w, h, nsamples);

parfor oidx=1:nsamples
    A = zeros(w, h);
    for i=1:n
        x = rawsites(i,1);
        y = rawsites(i,2);
        A(x, y) = i;
    end

    oidx
    offset = offsets(oidx,:);
    sites = rawsites + repmat(offset, n, 1);
    
    % repeat k steps
    k = ceil(log2(max(w, h)));

    for i=1:k
        j = pow2(i);

        for y=1:h
            for x=1:w
                if A(x, y) ~= 0 
                    % pass it to other grid points
                    for xx=-1:1
                        for yy=-1:1
                            x1 = x + (w/j)*xx;
                            y1 = y + (h/j)*yy;
                            if inRange(x1, y1, w, h)
                                if A(x1, y1) == 0 
                                    A(x1, y1) = A(x, y);
                                else
                                    % distance test
                                    idx0 = A(x, y);
                                    idx1 = A(x1, y1);

                                    dx0 = x1 - sites(idx0, 1);
                                    dy0 = y1 - sites(idx0, 2);
                                    dist0 = dx0*dx0+dy0*dy0;

                                    dx1 = x1 - sites(idx1, 1);
                                    dy1 = y1 - sites(idx1, 2);
                                    dist1 = dx1*dx1+dy1*dy1;

                                    if dist0 < dist1
                                        A(x1, y1) = A(x, y);
                                    end
                                end
                            end
                        end
                    end
                end
            end
        end        
    end

    B(:, :, oidx) = A;
end

% color the image
disp('colorizing the image...');
colors = hsv(n);
V = zeros(w, h, 3);
parfor y=1:h
    for x=1:w
        V(x, y, :) = color(B(x, y, :), colors);
    end
end

% mark the sites
for i=1:n
    x = rawsites(i,1);
    y = rawsites(i,2);
    B(x, y) = 0;
end

figure;imshow(V);
end

function c = color(B, colors)
    c = zeros(1, 3);
    n = length(B);
    for i=1:n
        c = c + colors(B(i), :);
    end
    c = c / n;
end

function flag = inRange(x, y, w, h)
    flag = (x>0 && x<=w && y>0 && y<=h);
end