%% RadialDistortion - Perform radial distortion correction
% Usage: imgOut = RadialDistortion(imgIn, a1, a2, showOutput)
% Inputs:
%    imgIn      = input image (grayscale or color)
%    a1         = amplitude of the quadratic term
%                 a1 < 0: pincushion distortion
%                 a1 > 0: barrel distortion
%    a2         = amplitude of the quartic term
% Outputs:
%    imgOut     = output image
%
% Created:     08/19/09, Hoang Trinh
%--------------------------------------------------------------------------

function imgOut = RadialDistortion(imgIn, a1, a2, showOutput)

if nargin < 4
    showOutput = 0;
end

[nrows,ncols,nz] = size(imgIn);
[xi,yi] = meshgrid(1:ncols,1:nrows);
imid = round(size(imgIn)/2); % Find index of middle element

% radial barrel distortion
xt = xi(:) - imid(2);
yt = yi(:) - imid(1);
[theta,r] = cart2pol(xt,yt);

rMax = sqrt(imid(1)*imid(1) + imid(2)*imid(2));
% normalize so that 0<=r<=1
r = r./rMax;

% a = .001; % Try varying the amplitude of the cubic term.
% a = -.0005; % Try varying the amplitude of the cubic term.

resamp = makeresampler('linear','fill');
%s = a1*r + a2*r.^3;
s = 1 + a1*(r.^2) + a2*(r.^4);

s = s.*r;
s = s.*rMax;

[ut,vt] = pol2cart(theta,s);
u = reshape(ut,size(xi)) + imid(2);
v = reshape(vt,size(yi)) + imid(1);
tmap_B = cat(3,u,v);
imgOut = tformarray(imgIn,[],resamp,[2 1],[1 2],[],tmap_B,.3);

if showOutput
    figure, imshow(imgOut);
    title('barrel');
end

end


%a   =  amplitude of the cubic term
%       a < 0: pincushion distortion
%       a > 0: barrel distortion

% function imgOut = RadialDistortion(imgIn, a1, a2)
% 
% [nrows,ncols, nz] = size(imgIn);
% [xi,yi] = meshgrid(1:ncols,1:nrows);
% imid = round(size(imgIn)/2); % Find index of center pixel
% 
% % radial distortion correction
% xt = xi(:) - imid(2);
% yt = yi(:) - imid(1);
% [theta,r] = cart2pol(xt,yt);
% 
% s = 1 + a1*(r.^2) + a2*(r.^4);
% 
% xNew = round(xt .* s);
% yNew = round(yt .* s);
% 
% ncolsNew = max(xNew) - min(xNew) + 1;
% nrowsNew = max(yNew) - min(yNew) + 1;
% 
% xNew = max(1,round(xNew + ncolsNew/2));
% yNew = max(1,round(yNew + nrowsNew/2));
% 
% imgOut = zeros(nrowsNew, ncolsNew,nz);
% 
% for i = 1:length(xNew)
%     imgOut(yNew(i),xNew(i),:) = imgIn(yi(i),xi(i),:);
% end
% 
% % [ut,vt] = pol2cart(theta,s);
% % u = reshape(ut,size(xi)) + imid(2);
% % v = reshape(vt,size(yi)) + imid(1);
% % tmap_B = cat(3,u,v);
% % resamp = makeresampler('linear','fill');
% % imgOut = tformarray(imgIn,[],resamp,[2 1],[1 2],[],tmap_B,.3);
% 
% figure, imshow(imgOut);
% title('barrel');
% 
% end