function RegImgSequence2
% RegImgSequence register an image sequence by minimum spanning tree method

input_folder = 'Images\';
output_folder = 'Images_reg\';


files = dir([input_folder '*jpg']);
I = double(rgb2gray(imread([input_folder files(1).name])));
J = I;
N = length(files);

g_width = 50; % graph width

Wt = zeros(N,N); % edge weights of a graph
for i=1:N,
    I = double(rgb2gray(imread([input_folder files(i).name])));
    I(:) = medfilt2(I,[3 3]);
    for j=max(1,i-g_width):i-1,
        J(:) = double(rgb2gray(imread([input_folder files(j).name])));
        J(:) = medfilt2(J,[3 3]);
        Wt(i,j) = mean2(abs(imresize(I-J,0.25)));
        Wt(j,i) = Wt(i,j);
    end
    disp(['Building graph for image ' num2str(i) ' of ' num2str(N)]);
end
disp(' ');
% call minimum spanning tree
[~, pred]= graphminspantree(sparse(Wt));
% MST construction
Tr = zeros(N,N);
for i=1:N,
    if pred(i) ~= 0,
        Tr(i,pred(i)) = 1;
        Tr(pred(i),i) = 1;
    end
end

% decide on an anchor image
path = graphpred2path(pred, N);
anchor = path(floor(length(path)/2));
disp(['Anchor image: ' files(anchor).name]);
disp(' ');
% traverse the tree to register images; 
[disc, Tpred,~] = graphtraverse(sparse(Tr),anchor,'Method','BFS','Directed',false);
% always register a child to its parent
I = imread([input_folder files(anchor).name]);
J = I;
imwrite(J,[output_folder files(anchor).name],'jpg');
xform = eye(3);
trans = zeros(3,3,N);
trans(:,:,anchor) = xform;
for i=2:N,
    child = disc(i);
    I(:) = imread([input_folder files(child).name]);
    parent = Tpred(disc(i));
    J(:) = imread([input_folder files(parent).name]);
%     xform = Reg2Im(I,J);
    xform = FeatReg(rgb2gray(I),rgb2gray(J));
    xform(:) = xform*trans(:,:,parent);
    trans(:,:,child) = xform;
    tform = affine2d(xform);
    K = double(I);
    K(:,:,1) = imwarp(double(I(:,:,1)),tform,'OutputView',imref2d(size(I(:,:,1))));
    K(:,:,2) = imwarp(double(I(:,:,2)),tform,'OutputView',imref2d(size(I(:,:,2))));
    K(:,:,3) = imwarp(double(I(:,:,3)),tform,'OutputView',imref2d(size(I(:,:,3))));
    
    imwrite(K/max(K(:)),[output_folder files(child).name],'jpg');
    disp(['Registering image ' num2str(i) ' of ' num2str(N)]);
    
end

end


function xform = FeatReg(J,I)
    % I: fixed image
    % J: moving image
    % we consider a rigid body transform here
    
    [h,w] = size(I);
    % RANSAC parameters
    numTrials = 10000;
    dThresh = 10;
    I(:) = medfilt2(I,[3 3]);
    J(:) = medfilt2(J,[3 3]);
    [frames1, descrs1] = vl_sift(single(I)) ;
    [frames2, descrs2] = vl_sift(single(J)) ;
    nn = findNeighbours(double(descrs1), double(descrs2)) ;
    matches = [1:size(descrs1,2) ; nn(1,:)] ;

    pfixed  = double(frames1(1:2, matches(1,:))) ; 
    pmoving = double(frames2(1:2, matches(2,:))) ; 
    pfixed = pfixed';
    pmoving = pmoving';

    % Estimate the transformation between J and I.
    xform = eye(3); 
    K = imwarp(double(J),affine2d(xform),'OutputView',imref2d(size(I)));
    maxerr = mean2(abs(imresize(double(I(25:h-25,25:w-25))-K(25:h-25,25:w-25),0.5)));
    for i=1:numTrials,
        p = randperm(size(pfixed,1));
        p = p(:);
        pf = pfixed(p(1:4),:);
        pm = pmoving(p(1:4),:);
        xform2 = computeRigidTransfrom(pf,pm);
        
        X = xform2(1,1)* pmoving(:,1) + xform2(1,2)* pmoving(:,2) + xform2(1,3);
        Y = xform2(2,1)* pmoving(:,1) + xform2(2,2)* pmoving(:,2) + xform2(2,3);      
        dists = sqrt((pfixed(:,1) - X).^2 + (pfixed(:,2) - Y).^2);
        inliers = dists<=dThresh;
        if sum(inliers) >= 5, %length(inliers) * 0.1,
            K(:) = imwarp(double(J),affine2d(xform2'),'OutputView',imref2d(size(I)));
            err = mean2(abs(imresize(double(I)-K,0.25)));
        else
            err = maxerr+1;
        end
        if err < maxerr,
            maxerr = err;
            xform = xform2';
        end 
    end

end

function xform = computeRigidTransfrom(pfixed,pmoving)
M0 = pfixed;
M0 = M0 - repmat(mean(M0),length(pfixed),1);
M = pmoving;
M = M - repmat(mean(M),length(pmoving),1);

% Best fit:
R = transpose(M) * M0 / length(pmoving);
% Consider noise:
[U, S, V] = svd(R);
if det(U * V) > 0
  Rotation = U * transpose(V);
else
  Rotation = U * [1,0; 0,-1] * transpose(V);
end

Translation(1) = (sum(pmoving(:,1)) - sum(pfixed(:,1) * Rotation(1) - pfixed(:,2) * Rotation(2)))/length(pmoving);
Translation(2) = (sum(pmoving(:,2)) - sum(pfixed(:,1) * Rotation(2) + pfixed(:,2) * Rotation(1)))/length(pmoving);

xform = zeros(3,3);
xform(1) = Rotation(1);
xform(2) = Rotation(2);
xform(4) = Rotation(3);
xform(5) = Rotation(4);
xform(7) = Translation(1);
xform(8) = Translation(2);
xform(9) = 1;
end

function xform = Reg2Im(J,I)
    % I: fixed image
    % J: moving image
    
    winsiz = 25;
    A = medfilt2(double(rgb2gray(J)),[5 5]);
    A = A(winsiz+1:end-winsiz,winsiz+1:end-winsiz);
    B = medfilt2(double(rgb2gray(I)),[5 5]);
    c = normxcorr2(A,B);
    [ypeak, xpeak] = find(c(2*winsiz:end-2*winsiz,2*winsiz:end-2*winsiz)==max(max(c(2*winsiz:end-2*winsiz,2*winsiz:end-2*winsiz))));
    yoffSet = ypeak(1) - size(A,1) + winsiz - 1;
    xoffSet = xpeak(1) - size(A,2) + winsiz - 1;
    
    xform=eye(3);
    xform(3,1) = xoffSet;
    xform(3,2) = yoffSet;    
end

function [nn, dist2] = findNeighbours(descrs1, descrs2, numNeighbours)
% FINDNEIGHBOURS  Find nearest neighbours
%   NN = FINDNEIGHBOURS(DESCRS1, DESCRS2) finds for each column of
%   DESCRS1 the closest column of DESRS2 (in Eulclidean distance)
%   storing the index to NN. The function uses a KDTree for
%   apporximate but fast matching.
%
%   NN = FINDNEIGHBOURS(DESCRS1, DESCRS2, NUMNEIGHBOURS) returns
%   NUMNEIGHBOUSRS neighbours by increasing distance, storing them
%   as successive rows of NN.
%
%   [NN, DIST2] = FINDNEIGHBOURS(...) returns the corresponding
%   matrix of distances DIST2 as well.

% Authors: Andrea Vedaldi

if nargin <= 2, numNeighbours = 1 ; end
vl_twister('state',0) ;
tree = vl_kdtreebuild(descrs2,'numTrees',2) ;
[nn, dist2] = vl_kdtreequery(tree, descrs2, descrs1, ...
                             'maxNumComparisons', 100, ...
                             'numNeighbors', numNeighbours) ;
end