...

% LFDecodeLenseletImageSimple - decodes a 2D lenselet image into a 4D light field
% 
% Usage:
%
%   [LF, LFWeight, DebayerLenseletImage, CorrectedLenseletImage] = ...
%      LFDecodeLenseletImageSpace( LenseletImage, WhiteImage, GridParams, DecodeOptions )
%   [LF, LFWeight, DebayerLenseletImage, CorrectedLenseletImage] = ...
%      LFDecodeLenseletImageSpace( LenseletImage, WhiteImage, GridParams, DecodeOptions, LevelLimits )
%
%
% This function follows the simple decode process described in:
% D. G. Dansereau, O. Pizarro, and S. B. Williams, "Decoding, calibration and rectification for
% lenselet-based plenoptic cameras," in Computer Vision and Pattern Recognition (CVPR), IEEE
% Conference on. IEEE, Jun 2013.
%
% This involves demosaicing, devignetting, transforming and slicing the input lenselet image to
% yield a 4D structure. More sophisticated approaches exist which combine steps into joint
% solutions, and they generally yield superior results, particularly near the edges of lenselets.
% The approach taken here was chosen for its simplicity and flexibility.
%
% Input LenseletImage is a raw lenselet iamge as found within Lytro's LFP picture files.  Though
% this function is written with Lytro imagery in mind, it should be possible to adapt it for use
% with other lenselet-based cameras.  LenseletImage is ideally of type uint16.
% 
% Input WhiteImage is a white image as found within Lytro's calibration data.  A white image is an
% image taken through a diffuser, and is useful for removing vignetting (darkening near edges of
% images) and for locating lenselet centers.  The white image should be taken under the same zoom
% and focus settings as the light field.  In the case of Lytro imagery, the helper functions
% LFUtilProcessLytroCalData and LFSelectWhiteImage are provided to assist in forming a list of
% available white images, and selecting the appropriate image based on zoom and focus settings.
% 
% Input GridParams is a model of the lenselet grid, as estimated, for example, using
% LFBuildLenseletGridModel -- see that function for more on the required structure.
% 
% Input DecodeOptions is a structure defining several decoding options:
%                  DehexMethod: 'fast' or 'triangulation', the latter is very slow
%                    Precision: 'single' or 'double'
% 
% Optional input LevelLimits is a two-element vector defining the black and white levels in the
% input images.  The appropriate values can be read from the Lytro metadata as demonstrated in
% LFExampleDecodeLenseletImage.
% 
% Output LF is a 5D array of size [Nj,Ni,Nl,Nk,3].  See [1] and the documentation accompanying this
% toolbox for a brief description of the light field structure.
%
% Optional output LFWeight is of size [Nj,Ni,Nl,Nk].  LFWeight contains a confidence measure
% suitable for use in filtering applications that accept a weight parameter.  This parameter is kept
% separate from LF rather than building in as a fourth channel in order to allow an optimization:
% when the parameter is not requested, it is not computed, saving significant processing time and
% memory.
% 
% Optional output DebayerLenseletImage is useful for inspecting intermediary results.  The debayered
% lenselet image is the result of devignetting and debayering, with no further processing.  Omitting
% this output variable saves memory.
% 
% Optional output CorrectedLenseletImage is useful for inspecting intermediary results.  The
% corrected lenselet image has been rotated and scaled such that lenselet centers lie on straight lines, and
% every lenselet center lies on an integer pixel spacing.  See [1] for more detail.  Omitting
% this output variable saves memory.
% 
% See also:  LFExampleDecodeLenseletImage, LFSelectWhiteImage, LFUtilProcessLytroCalData,
% LFBuildLenseletGridModel

% Part of LF Toolbox v0.1, released 26-Apr-2013
% Copyright (c) 2013, Donald G. Dansereau

function [LF, LFWeight, DebayerLenseletImage, CorrectedLenseletImage] = ...
    LFDecodeLenseletImageSimple( LenseletImage, WhiteImage, GridParams, DecodeOptions, LevelLimits )

if( ~exist('LevelLimits', 'var') )
    LevelLimits = [min(WhiteImage(:)), max(WhiteImage(:))];
end

BlackLevel = LevelLimits(1);
WhiteLevel = LevelLimits(2);

%---Rescale image values, remove black level---
WhiteImage = cast(WhiteImage, DecodeOptions.Precision);
WhiteImage = (WhiteImage - BlackLevel) ./ (WhiteLevel - BlackLevel);

LenseletImage = cast(LenseletImage, DecodeOptions.Precision);
LenseletImage = (LenseletImage - BlackLevel) ./ (WhiteLevel - BlackLevel);
LenseletImage = LenseletImage ./ WhiteImage; % Devignette
% Clip -- this is aggressive and throws away bright areas; there is a potential for an HDR approach here
LenseletImage = min(1, max(0, LenseletImage)); 

if( nargout < 2 )
    clear WhiteImage
end

%---Demosaic---
% This uses Matlab's demosaic, which is "gradient compensated".  This likely has implications near
% the edges of lenselet images, where the contrast is due to vignetting / aperture shape, and is not
% a desired part of the image
LenseletImage = cast(LenseletImage.*double(intmax('uint16')), 'uint16');
LenseletImage = demosaic(LenseletImage, 'bggr');
LenseletImage = cast(LenseletImage, DecodeOptions.Precision);
LenseletImage = LenseletImage ./  double(intmax('uint16'));

if( nargout >= 3 )
    DebayerLenseletImage = LenseletImage;
end

%---Tranform to an integer-spaced grid---
fprintf('\nAligning image to lenselet array...');
InputSpacing = [GridParams.HSpacing, GridParams.VSpacing];
NewLenseletSpacing = ceil(InputSpacing);
% Force even so hex shift is a whole pixel multiple
NewLenseletSpacing = ceil(NewLenseletSpacing/2)*2;
XformScale = NewLenseletSpacing ./ InputSpacing;  % Notice the resized image will not be square

NewOffset = [GridParams.HOffset, GridParams.VOffset] .* XformScale;
RoundedOffset = round(NewOffset);
XformTrans =  RoundedOffset-NewOffset;

NewGridParams = struct('HSpacing',NewLenseletSpacing(1), 'VSpacing',NewLenseletSpacing(2), ...
    'HOffset',RoundedOffset(1), 'VOffset',RoundedOffset(2), 'Rot',0);

%---Fix image rotation and scale---
RRot = LFRotz( GridParams.Rot );

RScale = eye(3);
RScale(1,1) = XformScale(1);
RScale(2,2) = XformScale(2);
ScaledAspect = XformScale(1) / XformScale(2);

RTrans = eye(3);
RTrans(end,1:2) = XformTrans;

% The following rotation can rotate parts of the lenselet image out of frame. 
% todo[optimization]: attempt to keep these regions
FixAll = maketform('affine', RRot*RScale*RTrans);
NewSize = size(LenseletImage(:,:,1)) .* XformScale;
LenseletImage = imtransform( LenseletImage, FixAll, 'XData',[1 NewSize(1)], 'YData',[1 NewSize(2)]);
if( nargout >= 2 )
    WhiteImage = imtransform( WhiteImage, FixAll, 'XData',[1 NewSize(1)], 'YData',[1 NewSize(2)]);
end

%---Construct slicing index array---
fprintf('\nConstructing slicing matrix...');
USize = floor(size(LenseletImage,2) / NewGridParams.HSpacing) - 1;
VSize = floor(size(LenseletImage,1) / NewGridParams.VSpacing) - 1;
MaxSpacing = max(NewGridParams.HSpacing, NewGridParams.VSpacing);  % Enforce square output in s,t
SSize = MaxSpacing + 1; % force odd for centered middle pixel -- H,VSpacing are even, so +1 is odd
TSize = MaxSpacing + 1;

TVec = cast(floor((-(TSize-1)/2):((TSize-1)/2)), 'int16');
SVec = cast(floor((-(SSize-1)/2):((SSize-1)/2)), 'int16');
VVec = cast(0:VSize-1, 'int16');
UVec = cast(0:USize-1, 'int16');

[tt,ss,vv,uu] = ndgrid( TVec, SVec, VVec, UVec );

%---Lenselet mask in s,t---
R = sqrt((cast(tt,DecodeOptions.Precision)*ScaledAspect).^2 + cast(ss,DecodeOptions.Precision).^2);
ValidIdx = find(R < NewLenseletSpacing(1)/2);

XStart = mod(NewGridParams.HOffset, NewGridParams.HSpacing);
YStart = mod(NewGridParams.VOffset, NewGridParams.VSpacing);

LFSliceIdxX = XStart + uu.*NewGridParams.HSpacing + ss;
LFSliceIdxY = YStart + vv.*NewGridParams.VSpacing + tt;

%--Use correct hex grid phase--
HexShiftStart = mod( floor(NewGridParams.VOffset/NewGridParams.VSpacing), 2 ) + 1;
LFSliceIdxX(:,:,HexShiftStart:2:end,:) = LFSliceIdxX(:,:,HexShiftStart:2:end,:) + NewGridParams.HSpacing/2;

%--clip--
LFSliceIdxX = max(1, min(size(LenseletImage,2), LFSliceIdxX ));
LFSliceIdxY = max(1, min(size(LenseletImage,1), LFSliceIdxY ));

if( nargout >= 2 )
    NColChans = 4;
else
    NColChans = 3;
end
LF = zeros(TSize, SSize, VSize, USize, NColChans, DecodeOptions.Precision);
LFSliceIdx = sub2ind(size(LenseletImage), cast(LFSliceIdxY,'int32'), ...
    cast(LFSliceIdxX,'int32'), ones(size(LFSliceIdxX),'int32'));

tt = tt - min(tt(:)) + 1;
ss = ss - min(ss(:)) + 1;
vv = vv - min(vv(:)) + 1;
uu = uu - min(uu(:)) + 1;
LFOutSliceIdx = sub2ind(size(LF), cast(tt,'int32'), cast(ss,'int32'), ...
    cast(vv,'int32'),cast(uu,'int32'), ones(size(ss),'int32'));

clear LFSliceIdxX LFSliceIdxY tt ss vv uu R

% todo[optimization]: The variable LFSliceIdx could be precomputed for each white image, saving much
% of the above computation

%---
fprintf('\nSlicing lenselets into LF...');
for( ColChan = 1:3 )
    LF(LFOutSliceIdx(ValidIdx) + numel(LF(:,:,:,:,1)).*(ColChan-1)) = ...
        LenseletImage( LFSliceIdx(ValidIdx) + numel(LenseletImage(:,:,1)).*(ColChan-1) );
end
if( NColChans > 3 )
    LF(LFOutSliceIdx(ValidIdx) + numel(LF(:,:,:,:,1)).*(NColChans-1)) = ...
        WhiteImage( LFSliceIdx(ValidIdx) );
end
if( nargout >= 4 )
    CorrectedLenseletImage = LenseletImage;
end
clear ValidIdx LFOutSliceIdx LFSliceIdx WhiteImage LenseletImage

%---Correct for hex grid and resize to square u,v pixels---
if( strcmp(DecodeOptions.DehexMethod,'fast') )
    fprintf('\nCorrecting hex sampling (1D approximation) and resizing to square u,v pixels...');
    FixHexAspect = 2/sqrt(3);
    LFSize = size(LF);
    NewUVec = 0:1/FixHexAspect:(size(LF,4)+1);  % overshoot then trim for identical U,V sizes
    NewUVec = NewUVec(1:LFSize(3));
    LFSize(4) = length(NewUVec);
    LF2 = zeros(LFSize, DecodeOptions.Precision);
    
    for( ColChan = 1:size(LF,5) )
        ShiftRows = squeeze(LF(:,:,1:2:end,:, ColChan));
        SliceSize = size(ShiftRows);
        SliceSize(4) = length(NewUVec);
        ShiftRows = reshape(ShiftRows, [size(ShiftRows,1)*size(ShiftRows,2)*size(ShiftRows,3), size(ShiftRows,4)]);
        ShiftRows = interp1q( (0:size(ShiftRows,2)-1)', ShiftRows', (0.5 + NewUVec)' )';
        ShiftRows(isnan(ShiftRows)) = 0;
        LF2(:,:,1:2:end,:,ColChan) = reshape(ShiftRows,SliceSize);
        
        ShiftRows = squeeze(LF(:,:,2:2:end,:, ColChan));
        SliceSize = size(ShiftRows);
        SliceSize(4) = length(NewUVec);
        ShiftRows = reshape(ShiftRows, [size(ShiftRows,1)*size(ShiftRows,2)*size(ShiftRows,3), size(ShiftRows,4)]);
        ShiftRows = interp1q( (0:size(ShiftRows,2)-1)', ShiftRows', NewUVec')';
        ShiftRows(isnan(ShiftRows)) = 0;
        LF2(:,:,2:2:end,:,ColChan) = reshape(ShiftRows,SliceSize);
    end
    LF = LF2;
    clear LF2
elseif( strcmp(DecodeOptions.DehexMethod,'triangulation') )
    fprintf('\nCorrecting hex sampling (tri approx.) and resizing to square u,v pixels...');
    HexAspect = 2/sqrt(3);
    OldVVec = (0:size(LF,3)-1);
    OldUVec = (0:size(LF,4)-1) * HexAspect;
    
    LFSize = size(LF);
    NewUVec = 0:1:max(OldUVec);
    NewVVec = (0:1:size(LF,3)-1);
    LFSize(4) = length(NewUVec);
    LFSize(3) = length(NewVVec);
    LF2 = zeros(LFSize, DecodeOptions.Precision);
    
    [Oldvv,Olduu] = ndgrid(OldVVec,OldUVec);
    [Newvv,Newuu] = ndgrid(NewVVec,NewUVec);
    Olduu(1:2:end,:) = Olduu(1:2:end,:) - HexAspect/2;

    for( ColChan = 1:size(LF,5) )
        fprintf('.');
         for( tidx= 1:size(LF,1) )
            for( sidx= 1:size(LF,2) )
                CurUVSlice = squeeze(LF(tidx,sidx,:,:,ColChan));
                Triangulation = TriScatteredInterp(Olduu(:), Oldvv(:), double(CurUVSlice(:)), 'natural');
                CurUVSlice = Triangulation(Newuu, Newvv);
                
                CurUVSlice = cast(CurUVSlice, DecodeOptions.Precision);
                CurUVSlice(isnan(CurUVSlice)) = 0;
                LF2(tidx,sidx, :,:, ColChan) = CurUVSlice;
            end
        end
    end
    LF = LF2;
    clear LF2
end

%---Resize to square s,t pixels---
fprintf('\nResizing to square s,t pixels using 1D linear interp...');
LFSize = size(LF);
NewTVec = cast(TVec, DecodeOptions.Precision) ./ ScaledAspect; % 1:1/ScaledAspect:size(LF,1);
LF2 = zeros(LFSize, DecodeOptions.Precision);

for( ColChan = 1:size(LF,5) )
    CurSlice = squeeze(LF(:,:,:,:, ColChan));
    CurSlice = permute(CurSlice, [4,3,2,1]);
    
    SliceSize = size(CurSlice);
    SliceSize(4) = length(NewTVec);
    
    CurSlice = reshape(CurSlice, [size(CurSlice,1)*size(CurSlice,2)*size(CurSlice,3), size(CurSlice,4)]);
    CurSlice = interp1q( cast(TVec, DecodeOptions.Precision)', CurSlice', NewTVec' )';
    CurSlice(isnan(CurSlice)) = 0;
    CurSlice = reshape(CurSlice,SliceSize);
    
    CurSlice = ipermute(CurSlice, [4,3,2,1]);
    LF2(:,:,:,:,ColChan) = CurSlice;
end
LF = LF2;
clear LF2

%---Trim s,t deadspace---
LF = LF(2:end-1,2:end-1,:,:,:);

%---Slice out LFWeight if it was requested---
if( nargout >= 2 )
    LFWeight = LF(:,:,:,:,4);
    LFWeight = LFWeight./max(LFWeight(:));
    LF = LF(:,:,:,:,1:3);
end

end