# 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.
# 2011-03-18 pi Converted it to python, with automatic file fetching
#
# latest version of this file can be found at:
# http://pirsquared.org/research/osacontest2011/SimpleIllumEstimation.py


import os, sys, logging
import numpy as np

logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
log = logging.getLogger('SimpleIllumEstimation.py')

try:
    import scipy.io as sio
except ImportError:
    log.critical("""Please install scipy in order to read .mat files, or
          contact Paul Ivanov requesting that he provide .npz equivalents, 
          removing the need to have scipy installed.""")
    sys.exit(1)

try:
    import matplotlib.pyplot as plt
    matplotlib_installed = True
except ImportError:
    matplotlib_installed = False


## Where the data are (but we'll fetch it if it's missing -pi)
dataDir = 'ContestImageData';

## Initialize answer index
actualIndex = 1; # this was in the matlab code but not used anywhere -pi

contest_url = "http://color.psych.upenn.edu/osacontest2011/"
def fetch_file(fname):
    "download files from the OSA contest website"
    log.info("%s not found, attempting to fetch it...", fname)
    import urllib
    try:
        u = urllib.urlopen(contest_url+fname)
        if u.code != 200: #if status code is not 'OK'
            import httplib
            raise IOError("Error %d "%u.code+httplib.responses[u.code])
        f = file(fname,'w')
        f.write(u.read())
        f.close(); u.close();
        log.info("download succeeded.")
    except IOError, e:
        log.critical(e)
        log.critical("Unable to fetch %s - please download it from:",fname)
        log.critical(contest_url+fname)
        sys.exit(1)


# ensure that we have the data
if not os.path.isdir(dataDir):
    log.info("Couldn't find '%s' directory,", dataDir)
    zfname= 'ContestImageData.zip'
    dataDir = 'ContestImageData'
    if not os.path.isfile(zfname):
        log.info("or the ~50 MB %s", zfname)
        log.warn("beware that it may take a while to download %s", zfname)
        fetch_file(zfname)
    log.info("will extract %s", zfname)
    import zipfile
    z = zipfile.ZipFile(zfname)
    for member in z.namelist():
        z.extract(member)
    z.close()

## Load basis and cone files
tmp = sio.loadmat(os.path.join(dataDir,'T_cones_osa'));
tmp.update(sio.loadmat(os.path.join(dataDir,'B_illum')))
tmp.update(sio.loadmat(os.path.join(dataDir,'B_sur')))

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

## Flags
LOOKATIMAGE = False;

if LOOKATIMAGE and not matplotlib_installed:
    log.warn("Please install matplotlib if you want to see the images")
    LOOKATIMAGE = False

if (LOOKATIMAGE):
    imgFig = plt.figure();

PAUSE_MSG = "Paused: press Enter to continue"

## Load calibration image and use it to, well, calibrate
tmp.update(sio.loadmat(os.path.join(dataDir,'ImageCal')))
tmp.update(sio.loadmat(os.path.join(dataDir,'spd_ImageCal')))
# drop the meta information we get when loading .mat files
[tmp.pop(k) for k in tmp.keys() if k.startswith('_')];
# now put all of the items we've stored in tmp into the local namespace
locals().update(tmp)

if (LOOKATIMAGE):
    plt.figure(imgFig.number); plt.clf();
    plt.imshow(theImage/theImage.max());
    plt.show()
    #raw_input(PAUSE_MSG)

# 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 constraint 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 = np.matrix(np.zeros((3,1)));
theImage.mean(0).mean(0,out=meanLMS)

T_cones_osa = np.matrix(T_cones_osa, copy=False)
spd_illumSpdCal = np.matrix(spd_illumSpdCal, copy=False)
B_sur = np.matrix(B_sur, copy=False)
B_illum= np.matrix(B_illum, copy=False)
meanSurWgts = np.linalg.inv(T_cones_osa*np.diagflat(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 in range(nImages):
    # Load the image
    fpath =os.path.join(dataDir,'Image%d'%(i+1))
    theImage = sio.loadmat(fpath)['theImage']
    if LOOKATIMAGE:
        plt.figure(imgFig.number); plt.clf();
        plt.imshow(theImage/theImage.max());
        plt.show()
        #raw_input(PAUSE_MSG)
    
    # All we need from the image for this simple algorithm is the
    # mean LMS values
    theImage.mean(0).mean(0,out=meanLMS)
    
    # The mean LMS values let us get the illuminant, within the illuminant
    # linear model
    estimatedIllumSpds[:,i] = (B_illum*np.linalg.inv(T_cones_osa*np.diagflat(B_sur*meanSurWgts)*B_illum)*meanLMS).T

## Save out the estimates.  This is simple tab delimited text, with no
# headers, and is the format we use for the contest.
np.savetxt('myEstimatedIllumSpds.txt', estimatedIllumSpds, fmt="  % .7e", delimiter='\t')
#
### Load the estimates we just saved, simply to prove we know how to 
## do it.
del(estimatedIllumSpds)
estimatedIllumSpds = np.loadtxt('myEstimatedIllumSpds.txt')

if not os.path.isfile('estimatedIllumSpds.txt'):
    fetch_file('estimatedIllumSpds.txt')
estimatedIllumSpdsProvided = np.loadtxt('estimatedIllumSpds.txt')

ediff = estimatedIllumSpdsProvided - estimatedIllumSpds
# SSD: sum of the squares of differences
print "The SSD between file provided and the one calculated is %.9f"%(ediff**2).sum()

# TODO: the code below has not yet been converted to python -pi

#raw_input("done")
### 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