function IU = unwarp(ID, K, Ox, Oy, Sx, Sy, OutColor)
%UNDISTORT Remove a distorção radial de uma imagem.


%  IU = Undistort(ID, K, Ox, Oy, Sx, Sy, OutColor) remove a distorção radial da
%  imagem ID, usando os parâmetros de calibração da câmera:
%
%  - K são os coeficientes de distorção radial. É um vetor de três elementos,
%    usados conforme o seguinte modelo de distorção radial:
%      xd = x * (1 + (K(1) + K(2) * r^2 + K(3) * r^4))
%      yd = y * (1 + (K(1) + K(2) * r^2 + K(3) * r^4)) ,
%    onde r = sqrt(dx^2 + dy^2) é a distância radial ideal em relação ao
%    centro da imagem sem distorção;
%
%  - Ox e Oy são as coordenadas do centro da imagem;
%
%  - Sx e Sy são a relação entre o tamanho do pixel e do elemento sensor da
%    câmera (default = 1 e 1);
%
%  - OutColor define as cores que serão atribuídas aos pontos externos à
%    imagem. Pode ser:
%    - um escalar que define o nível de cinza;
%    - um vetor de três elementos, definindo a cor RGB, para imagens
%      coloridas;
%    - uma estrutura que define um grid, contendo os seguintes campos:
%      - GridColor1 e GridColor2 são as duas cores do grid (cinza ou vetor
%        RGB);
%      - GridSize é o tamanho do grid, em pixels.
%
%  Autor:  Vilar Camara Neto
%  Versão: 1.1 - 2007-03-07


% Valores default para os argumentos opcionais
if nargin < 5, Sx = 1; end
if nargin < 6, Sy = Sx; end
if nargin < 7, OutColor = []; end

% Normalização para imagens fornecidas em inteiros (0 a 255 ao invés de 0.0
% a 1.0)
OrigIsUint8 = isinteger(ID);
if OrigIsUint8, ID = double(ID) / 255; end

% Dimensões e profundidade da imagem original
wd = size(ID, 2);
hd = size(ID, 1);
dd = size(ID, 3);

IDsize = [hd, wd] + 2;

% Ajuste do parâmetro OutColor
if isempty(OutColor)
    if dd == 3
        OutColor.GridColor1 = [1, 1, 1];
    else
        OutColor.GridColor1 = 1;
    end
end

if isnumeric(OutColor)
    OutColor = struct('GridColor1', OutColor);
end
if length(OutColor.GridColor1) < dd
    OutColor.GridColor1 = repmat(OutColor.GridColor1(1), 1, dd);
elseif length(OutColor.GridColor1) > dd
    dd = length(OutColor.GridColor1);
    ID = repmat(ID, [1, 1, dd]);
end
if ~isfield(OutColor, 'GridColor2')
    OutColor.GridColor2 = OutColor.GridColor1;
end
if ~isfield(OutColor, 'GridSize')
    OutColor.GridSize = 8;
end

BandOffset = (0:(dd - 1)) * prod(IDsize);

% Aumento proposital da imagem em 2 linhas e 2 colunas, atribuindo as cores
% do grid externo à imagem aos pixels recém-criados
onv = mod(1:hd, (2 * OutColor.GridSize)) >= OutColor.GridSize;
onh = mod(1:wd, (2 * OutColor.GridSize)) >= OutColor.GridSize;
for band = 1:dd
    ID(:, (wd + 1):(wd + 2), band) = OutColor.GridColor1(band);
    ID((hd + 1):(hd + 2), :, band) = OutColor.GridColor1(band);
    ID(onv, wd + 1, band) = OutColor.GridColor2(band);
    ID(~onv, wd + 2, band) = OutColor.GridColor2(band);
    ID(hd + 1, onh, band) = OutColor.GridColor2(band);
    ID(hd + 2, ~onh, band) = OutColor.GridColor2(band);
    ID(hd + 1, wd + 2, band) = OutColor.GridColor2(band);
    ID(hd + 2, wd + 1, band) = OutColor.GridColor2(band);
end

% Obtenção dos limites da imagem não-distorcida. Usa-se um método iterativo
% disponível no Matlab para este fim, já que a função que mapeia pontos
% não-distorcidos para distorcidos é não-linear.
Options = optimset('Display', 'off');
PT = [0, 0; 0, size(ID, 1) - 1; size(ID, 2) - 1, 0; ...
    size(ID, 2) - 1, size(ID, 1) - 1];
PL = fsolve(@(X) optim_func(X, PT, K, Ox, Oy, Sx, Sy), PT(:), Options);

% Margens ao redor da imagem calculada
ImageMargin = 0;

% Coordenadas do canto superior esquerdo da imagem não-distorcida...
Dx = floor(min(PL(1:4))) - ImageMargin;
Dy = floor(min(PL(5:8))) - ImageMargin;

% ...e dimensões da imagem
w = ceil(max(PL(1:4))) - Dx + 1 + ImageMargin;
h = ceil(max(PL(5:8))) - Dy + 1 + ImageMargin;

% Pré-alocação da imagem não-distorcida
IU = zeros(h, w, dd);

% As operações são feitas em "fatias" verticais da imagem, ou seja: a cada
% iteração deste "for", um conjunto de colunas é processado de uma vez. O
% parâmetro xstep indica o número de colunas a serem processadas em cada
% iteração. Valores altos para xstep aceleram o processamento, porém
% aumentam o consumo de memória.
xstep = 10;
for xmin = 0:xstep:(w - 1)
    xmax = min(w, xmin + xstep) - 1;

    % Inicialização das coordenadas não-distorcidas
    [Px, Py] = meshgrid((xmin:xmax) + Dx, (0:(h - 1)) + Dy);
    PD = [Px(:), Py(:)];

    % Transformação de "warp", para calcular as coordenadas distorcidas
    % correspondentes às não-distorcidas
    PU = warp_points(PD, K, Ox, Oy, Sx, Sy);

    % Bases e fatores-base para interpolação bilinear dos pontos
    PUbase = floor(PU);
    PUfac = PU - PUbase;

    % Linhas e colunas para interpolação bilinear
    left = PUbase(:, 1) + 1;
    top = PUbase(:, 2) + 1;
    right = left + 1;
    bottom = top + 1;

    % Pesos da interpolação para os quatro pontos
    fright = PUfac(:, 1);
    fbottom = PUfac(:, 2);
    fleft = 1 - fright;
    ftop = 1 - fbottom;

    % Ajuste dos pesos: se alguns pontos estiverem fora da imagem, estes
    % serão desconsiderados na interpolação.
    out = (left < 1) | (left > wd);
    left(out) = mod(floor((left(out) - 1) / OutColor.GridSize), 2) + wd + 1;
    out = (right < 1) | (right > wd);
    right(out) = mod(floor((right(out) - 1) / OutColor.GridSize), 2) + wd + 1;
    out = (top < 1) | (top > hd);
    top(out) = mod(floor((top(out) - 1) / OutColor.GridSize), 2) + hd + 1;
    out = (bottom < 1) | (bottom > hd);
    bottom(out) = mod(floor((bottom(out) - 1) / OutColor.GridSize), 2) + hd + 1;

    % Índices para interpolação bilinear
    TLindex = sub2ind(IDsize, top, left);
    BLindex = sub2ind(IDsize, bottom, left);
    TRindex = sub2ind(IDsize, top, right);
    BRindex = sub2ind(IDsize, bottom, right);

    % Interpolação bilinear
    for band = 1:dd
        IU(:, (xmin:xmax) + 1, band) = reshape( ...
            ID(TLindex + BandOffset(band)) .* ftop .* fleft + ...
            ID(BLindex + BandOffset(band)) .* fbottom .* fleft + ...
            ID(TRindex + BandOffset(band)) .* ftop .* fright + ...
            ID(BRindex + BandOffset(band)) .* fbottom .* fright, ...
            h, []);
    end
end

% Reconversão para valores inteiros, se a imagem foi fornecida nesse
% formato
if OrigIsUint8, IU = uint8(IU * 255); end


% ------------------------------------------------------------------------


function F = optim_func(X, T, K, Ox, Oy, Sx, Sy)
% Função auxiliar para o processo de recuperação das margens da imagem
% não-distorcida.

X = reshape(X, [], 2);
F = warp_points(X, K, Ox, Oy, Sx, Sy);
F = F(:) - T(:);


% ------------------------------------------------------------------------


function PU = warp_points(PD, K, Ox, Oy, Sx, Sy)
% Função para converter coordenadas distorcidas em não-distorcidas.

PU(:, 1) = (PD(:, 1) - Ox) * Sx;
PU(:, 2) = (PD(:, 2) - Oy) * Sy;

Indexes = any(PU ~= 0, 2);
rd2 = sum(PU(Indexes, :) .^ 2, 2);
rd = sqrt(rd2);
rd4 = rd2 .* rd2;
ru = (1 + K(1) + K(2) * rd2 + K(3) * rd4) .* rd;
PU(Indexes, :) = PU(Indexes, :) .* repmat(ru ./ rd, 1, 2);

PU(:, 1) = PU(:, 1) / Sx + Ox;
PU(:, 2) = PU(:, 2) / Sy + Oy;
