% SHACORR - Performs illumination correction that preserves medium size objects.
%           The medium sized object are typically darker or brighter than
%           the average in the image. Each type can be handled separately.
%
% Usage:
%
% [nim, pbq, qq] = shaCorr(im,f,bright,dark, d, st, cor,contrast,pfact)
%
%
% Arguments:
%     im        - Input image. If it is a RGB image, it is converted to HSV
%                 and the process is performed on the V component only.
%                 Correction in RGB can be done by calling the function for
%                 each component.
%     f         - scale the filter size to get a high pass filter.
%                 Typically less than 1. If less than zero no medium sized
%                 objects will be preserevd.
%     bright    - preserve bright parts. Typically a value between 0 and 1.
%                 But can be higher in order to enhance it further.
%     dark      - preserve dark parts. Typically a value between 0 and 1.
%                 But can be higher in order to enhance it further.
%     d         - used to change the mask size. Default 6. Will be halfed
%                 in each iteration.
%     st        - Number of correction iterations. Default 3.
%     cor       - correct contrast, especially in dark or bright areas.
%                 Can be set to:
%                 cor = 0, no contrast correction.
%                 cor = 1, contrast cocrrection of both high and low
%                          contrast areas.
%                 cor = 2, contrast compensation with no supression of high
%                          contrast areas. Pump up low contrast areas.
%                 cor = 3, works as cor = 2, but do not pump up too hard.
%                          Use parameter pfact to scale the increase.
%
%     contrast  - How much should contrast be enhanced? 
%                 Typically a value between 0 and 1, but can be higher.
%     pfact     - pumping up factor.
%     msize     - Preferred size of down scaled image. Working on a smaller
%                 image will be many times faster than working on the whole
%                 image.
%
%
% Returns:
%     cim       - corrected image.
%
%
% For an example of the use of this function see RANSACFITHOMOGRAPHY or
% RANSACFITPLANE 
%
%
% Copyright (c) 2011-2019 Anders Hast
% Uppsala University
% http://www.cb.uu.se/~aht
% 
% 
% Permission is hereby granted, free of charge, to any person obtaining a copy
% of this software and associated documentation files (the "Software"), to deal
% in the Software without restriction, subject to the following conditions:
% 
% The above copyright notice and this permission notice shall be included in 
% all copies or substantial portions of the Software.
%
% The Software is provided "as is", without warranty of any kind.
%
% History
% 
%
%

function [cim] = shaCorr(im,f,bright,dark,d,st,cor,contrast,pfact,msize)
    % Test number of parameters
    narginchk( 1,10);
    nargoutchk(1,1);
    % Set defualt values
    if nargin<10
       msize=100; 
    end
    if nargin<7
        cor=0;
        contrast=0;
        pfact=0;
    end
    if nargin<5
       d=6;
       st=3;
    end
    if nargin<2
       f=-1;
       bright=0;
       dark=0;
    end
    % Image size
    [n,m,c]=size(im);
    So=[n m];
    
    % Do we have an RGB image?
    if c>1
        imhsv=rgb2hsv(im);
        im=imhsv(:,:,3);
    end

    % Setup for the lowpass filter
    [N,sigma, S]=mask(So,msize); 
    g2 = fspecial('gauss',[N 1],sigma);
    div2 = conv2(conv2(ones(S),g2,'same'),g2','same'); 

    if f>0
        % Setup for the highpass filter
        [N,sigma]=mask(ceil(N*f),0);
        g1 = fspecial('gauss',[N 1],sigma);
        div1 = conv2(conv2(ones(S),g1,'same'),g1','same');
    end
    
    % Perform iterative correction
    k=0;
    while k<st
        % Removing the bright/dark areas before every correction and then
        % adding them after in each step makes a better result.

        if f>0
            % Downscale the image in order to compute the ligthing faster
            imq=double(imresize(im,[S(1), S(2)],'bicubic'));

            % Compute the bandpass filter
            p1=conv2(conv2(imq,g1,'same'),g1','same')./div1;
            p2=conv2(conv2(imq,g2,'same'),g2','same')./div2;
            im2=p2-p1;
        
            % Scale up the bandpass image containing medium size objects
            q=imresize(im2,[n, m],'bicubic');
       
            % Separate the negative (bright) and postive (dark) parts
            s1=find(q<0);
            s2=find(q>0);
            q(s1)=q(s1)*bright;
            q(s2)=q(s2)*dark;

            % Remove the brightest and darkest parts so the average illumination
            % can be computed more accurately for the whole image.
            % They will be added after illumination correction
            im=im+q;
        end
        
        % Compute a new mask in each iteration
        [N,sigma]=mask(ceil(max(S)/d),0);
        g = fspecial('gauss',[N 1],sigma);
        div = conv2(conv2(ones(S),g,'same'),g','same');

        % Scale down the image in order to compute the illumination faster
        ims=double(imresize(im,[S(1), S(2)],'bicubic'));

        % Compute the illumination
        p=conv2(conv2(ims,g,'same'),g','same')./div;
        pm=mean(mean(p));
  
        % Compute the contrast
        s=abs(ims-p);
        w=conv2(conv2(s,g,'same'),g','same')./div;
        
        % Scale up comtrast
        if cor>0
            wb=imresize(w,[n, m],'bicubic');
        end
        % Scale up illumination
        pb=imresize(p,[n, m],'bicubic');
    
        % Correct the image by removing the illumination
        % This removes the mean locally
        im=im-pb;
         
        % Mean illumination
        mm=(mean(mean(w)));

        % Correct the contrast
        if cor > 1 
            wb(find(wb>mm))=mm; % Do not supress, just increase
        end
        % Do not increase too hard
        if cor > 2
            wb(find(wb<mm*pfact))=mm*pfact;
        end
        
        % Correct the contrast
        if cor >0
            % How much contrast do we want?
            wb=(wb-mm)*contrast+mm;
            im=mm*(im./(wb));
        end
        
        % Add the mean globally to obtain the corrected image
        im=im+pm;

        if f>0
            % Add the medium sized object that were temporarily removed
            im=im-q;
        end
        % Compute illumination and contrast using bigger and bigger masks.
        % This will give a better result!
        if d>1
            d=d/2;
        end
       
        % then repeat the process untill an even ligthing is obtained        
        k=k+1;
    end

    % Do we have an RGB image?    
    if c>1
        imhsv(:,:,3)=im;
        cim=hsv2rgb(imhsv);
    else
        cim=im;        
    end
      
    
end

% Function to compute mask size and sigma depending on the window size
function [N,sigma,S]=mask(So,sz)
    if sz==0
        N=So;
        S=0;
    else
        % What scale shall we use
        scale = min(So)/sz;   
        S=ceil(So/scale);
        % Compute mask size depending on the size of the image
        N=ceil(max(S));
    end
    % Make sure N is odd!
    if mod(N,2) == 0
        N=N+1;
    end
    % Compute sigma depending on mask size
    sigma=N/6;
end