% SimpleIllumEstimation
%
% Very simple gray world illumination estimation program. Basically follows
% method introduced by Buchsbaum, 1980, J. Franklin Institute.
%
% This is not meant to be a particularly great algorithm, but the program
% estimates how to read the data and the format we'd like to get the 
% estimates in.  It also provides a baseline score that someone should
% be able to beat.
%
% 1/17/11  dhb  Wrote it.
% 1/18/11  dhb  Write out and read tab delimited text, to illustrate
%               desired contest format.

%% Clear and close
clear; close all

%% Where the data are
dataDir = 'ContestImageData/';

%% Initialize answer index
actualIndex = 1;

%% Load basis and cone files
load([dataDir 'T_cones_osa']);
load([dataDir 'B_illum']);
load([dataDir 'B_sur']);

nImages = 10;
estimatedIllumSpds = zeros(31,nImages);

%% Flags
LOOKATIMAGE = 0;
if (LOOKATIMAGE)
    imgFig = figure; clf;
end

%% Load calibration image and use it to, well, calibrate
load([dataDir 'ImageCal']);
load([dataDir 'spd_ImageCal']);
if (LOOKATIMAGE)
    figure(imgFig); clf;
    imshow(theImage/max(theImage(:)));
    pause;
end

% We know the calibration image illuminant, so use this
% to find mean surface reflectance for the calibration
% image.  We can do this, since the mean LMS responses
% are equal to the cone coordinates of the mean surface
% under the known calibration illuminant, and we have
% three dimensional linear model constratin on surfaces.
%
% For processing the rest of the images, we'll assume that
% mean surface reflectance is equal to that derived here.
% This isn't necessarily the case, but seems more likely to
% be right than pulling a mean out of thin air.
meanLMS = zeros(3,1);
for k = 1:3
    meanLMS(k) = mean(mean(theImage(:,:,k)));
end
meanSurWgts = inv(T_cones_osa*diag(spd_illumSpdCal)*B_sur)*meanLMS;

%% Estimate the illuminant for each image, on the assumption that
% the spatial mean of the surface reflectances in the images matches
% the mean derived above.
for i = 1:nImages
    % Load the image
    load([dataDir 'Image' num2str(i)]);
    if (LOOKATIMAGE)
        figure(imgFig); clf;
        imshow(theImage/max(theImage(:)));
        pause;
    end
    
    % All we need from the image for this simple algorithm is the
    % mean LMS values
    meanLMS = zeros(3,1);
    for k = 1:3
        meanLMS(k) = mean(mean(theImage(:,:,k)));
    end
    
    % The mean LMS values let us get the illuminant, within the illuminant
    % linear model
    estimatedIllumSpds(:,i) = B_illum*inv(T_cones_osa*diag(B_sur*meanSurWgts)*B_illum)*meanLMS;
end

%% Save out the estimates.  This is simple tab delimited text, with no
% headers, and is the format we use for the contest.
save estimatedIllumSpds.txt estimatedIllumSpds -ascii -tabs

%% Load the estimates we just saved, simply to prove we know how to 
% do it.
clear estimatedIllumSpds
estimatedIllumSpds = load('estimatedIllumSpds.txt','-ascii');

%% Compute the error, if the answers are available.
PLOTESTIMATES = 0;
if (PLOTESTIMATES)
    illumFig = figure;
end
if (exist([dataDir 'actualIllumSpds.mat'],'file'))
    SumMSE = 0;
    load([dataDir 'actualIllumSpds']);
    if (size(actualIllumSpds,2) ~= nImages)
        error('Oops.  Inconsistent number of contest images');
    end
    for i = 1:nImages
        % Scale the estimate to get the best MSE fit to the actual
        scaledEstimatedSpd = estimatedIllumSpds(:,i)*(estimatedIllumSpds(:,i)\actualIllumSpds(:,i));
        
        % Take difference
        estimationDiff = actualIllumSpds(:,i)-scaledEstimatedSpd;
        
        % Compute MSE
        MSE = sum(estimationDiff.^2)/length(estimationDiff);
        SumMSE = SumMSE + MSE;
        
        % Plot, just as a reality check
        if (PLOTESTIMATES)
            figure(illumFig); clf; hold on
            plot(400:10:700,actualIllumSpds(:,i),'k','LineWidth',2);
            plot(400:10:700,estimatedIllumSpds(:,i),'r:','LineWidth',1);
            plot(400:10:700,scaledEstimatedSpd,'r','LineWidth',2);
            xlabel('Wavelength');
            ylabel('Illuminant');
            xlim([400 700]);
            axis('square');
            pause;
        end
    end
    fprintf('Mean of MSE over illuminant estimates is %g\n',SumMSE/nImages);
else
    fprintf('Sorry, you don''t get to know the actual illuminants until the contest is over!\n');
end
