
%% Lens Distortion Correction  (Script)
%
% 1) Choose coefficient files containing model parameters trained using
% specific LDC method (ex. blind) and specific training set. 
% 
% ex. methodName_(fullset)_cameraModel_subject.mat 'coefs' 'scale' 'offset'
%
% 2) Choose image files (jpg/png/ppm) with identical camera model
%
% 3) Create "undistorted" image using methodName-LDC


%% Script Setup 

addpath('./common/');
addpath('../standard_lens_correction_c/');
addpath('../blind_image_lens_correction_c/');
addpath('../alvarez_image_lens_correction_c/');

close all;
clc;

rng(0);


%% Find Supported Coefficient and Image Files 

testPath = '../test_environment/';
coefsPath = './coefs/';
corrPath = './corr/';
imgPath = './images/';
linePath = './line_points/';
suppCoefFiles = {'mat'};
suppImgFiles = {'jpg', 'png', 'ppm'};

% Find all supported coefficient files within directory
fprintf('\tSearching directory %s...\n', coefsPath);
[coefNames, nMats] = find_supported_files(coefsPath, suppCoefFiles);
fprintf('\t\tFound %d coefs .mat files\n', nMats);

% Find all supported image files within directory
fprintf('\tSearching directory %s...\n', imgPath);
[imgNames, nImgs] = find_supported_files(imgPath, suppImgFiles);
fprintf('\t\tFound %d image files\n', nImgs);


%% Select Files 

% Decide whether to manually specify line point .mat files used, or to
%have program automatically segment files into 'camera-model' groups and
%have user choose specific camera model ('single-only-selection')

[selectImgNames, selectCamModelName] = select_files(imgNames, '^(?<group>[a-zA-Z0-9-]+)_', 'Image');
nImgs = length(selectImgNames);    

[selectCoefNames, selectMethodName] = select_files(coefNames, '^(?<group>[a-zA-Z0-9-]+)_', 'Coefficient');
nMats = length(selectCoefNames);

%% Construct output file names

% Check if Alvarez method selected, if so then corresponding line-points
%file should be chosen instead of coefficient file.

lineFiles = cell(nMats, nImgs);
outImgFiles = cell(nMats, nImgs);
methodNames = cell(nMats, 1);
for i = 1:nImgs,   
    for c = 1:nMats,

        fileComps = strsplit(selectCoefNames{c}, '[_.]', 'DelimiterType', 'RegularExpression');
        methodNames{c} = fileComps{1};

        if strcmp('fullset', fileComps{2})                        
            camModel = fileComps{3};            
            if length(fileComps) > 3 && ~strcmp(fileComps{4}, '.mat')
                subject = fileComps{4};
            else
                subject = '';
            end
            outImgFiles{c,i} = strcat(methodNames{c}, '_corr', '_fullset', subject, '_', selectImgNames{i}); 
        else
            camModel = fileComps{2};
            subject = fileComps{3};
            outImgFiles{c,i} = strcat(methodNames{c}, '_corr', '_', subject, '_', selectImgNames{i});
        end

        if strcmp(methodNames{c}, 'alvarez')
            if strcmp('fullset', fileComps{2})
                lineFiles{c} = strcat('fullset', '_', camModel, '.mat');
            else
                lineFiles{c} = strcat('manip', '_', camModel, '_', subject, '.mat');
            end
        end

        [~, outImgFiles{c,i}, ~] = fileparts(outImgFiles{c,i});

    end
end


%% Compute Lens Distortion Model Parameters/Coefficients 

methodPaths = {'../standard_lens_correction_c/', '../blind_image_lens_correction_c/', '../alvarez_image_lens_correction_c/'};
models = {'std', 'std', 'rad'};

for i = 1:nImgs,    

    for c = 1:nMats,
        
        coefFile = fullfile(coefsPath, selectCoefNames{c});
        load(coefFile, 'coefs', 'scale', 'offset');   
                                
        switch(methodNames{c})                        

            case 'std'

                standard_image_lens_correction(selectImgNames{i}, outImgFiles{c,i}, coefs, 'dirPath', imgPath, 'dirOutPath', corrPath, ...
                                               'scale', scale, 'offset', offset);                                                                          

            case 'blind'
               
                blind_image_lens_correction(selectImgNames{i}, outImgFiles{c,i}, coefs, 'dirPath', imgPath, 'dirOutPath', corrPath, ...
                                            'scale', scale, 'offset', offset);                        

            case 'alvarez'
                                
                lineFile = fullfile(linePath, lineFiles{c});
                load(lineFile, 'pts', 'lineSizeMat');
                
                cd(methodPaths{3});
                [coefs, err] = alvarez_image_lens_correction('seqName', selectImgNames{i}, 'seqOutName', outImgFiles{c,i}, ...
                                                             'dirPath', fullfile(testPath, imgPath), 'dirOutPath', fullfile(testPath, corrPath), ...
                                                             'pts', pts, 'lineSizeMat', lineSizeMat,'optimizeCenter', 1);                                                                                          
                cd(testPath);                
                
            otherwise
                error('LensDistortionCorrection:InvalidMethod', 'Unknown method specified');

        end
                
    end
    
end
  





