% LEARNING ACTIVE BASIS MODEL FROM NONALIGNED IMAGES -- keep-track version
%% Load in exponential model, mex C codes, and set parameters
% ExponentialModel; % this line is not needed if 'storedExponentialModel' exists
clear; close all; 
load 'storedExponentialModel'; % load in exponential model 
mex CsharedSketch.c;  % learning from aligned images 
mex CgetSUM2.c;  % detect object from multi-resolution images  
mex Caverage.c; % compute within window average for normalization 
mex CgetMAX1track.c; % local maximum pooling for MAX1 maps
mex Ccopy.c; % copy around detect location with sigmoid transformation to prepare for learning
mex Ccopy0.c % simple copy without sigmoid transformation
epsilon = .1; % allowed correlation between selected Gabors 
subsample = 1; % subsample in computing MAX1 and SUM2 maps   
locationShiftLimit = 3; % shift in normal direction = locationShiftLimit*subsample pixels
orientShiftLimit = 1; % shift in orientation
Correlation = CorrFilter(allFilter, epsilon); % correlation between filters 
numElement = 45; % number of Gabors in active basis   
numIteration = 5;  % number of iterations
useWholeImageOrNot = 1; % whether use the whole image of I{start} or use a portion of it to initialize learning
windowNormalizeOrNot = 1; % whether to normalize within scanning window of the template in detection
%% Load in training images and initialize SUM1 maps for learning
resizeFactor = .4; % resize the input images 
imageFolder = 'positiveImage'; % folder of training images  
imageName = dir([imageFolder '/*.jpg']);
numImage = size(imageName, 1); % number of training images 
I = cell(1, numImage); 
for (img = 1 : numImage)
    tmpIm = imread([imageFolder '/' imageName(img).name]); 
    if size(tmpIm,3) == 3
        tmpIm = rgb2gray(tmpIm);
    end
    I{img} = imresize(single(tmpIm), resizeFactor, 'nearest');
end
starting = 1; % initilize by single image learning from I{starting}
startx = 6; endx = 132; % initialize by learning from a given patch
starty = 35; endy = 175; 
tmp = I{starting}; 
if (useWholeImageOrNot>0)  % initialize from the whole image 
    startx = 1; starty = 1; 
    [endx endy] = size(tmp); 
end
I0{1} = tmp(startx:endx, starty:endy); % initialize the template from the starting image 
[sizeTemplatex sizeTemplatey] = size(I0{1}); % size of template
halfTemplatex = floor(sizeTemplatex/2.); halfTemplatey = floor(sizeTemplatey/2.); % half size
sizeTemplatexSubsample = floor(sizeTemplatex/subsample); % subsample size
sizeTemplateySubsample = floor(sizeTemplatey/subsample); 
SUM1mapLearn = cell(numImage, numOrient); % data structures for learning
MAX1mapLearn = cell(numImage, numOrient); 
trackMapLearn = cell(numImage, numOrient);
for (img = 1:numImage) 
    for (orient = 1:numOrient)
       SUM1mapLearn{img, orient} = single(zeros(sizeTemplatex, sizeTemplatey)); 
       MAX1mapLearn{img, orient} = single(zeros(sizeTemplatexSubsample, sizeTemplateySubsample)); 
       trackMapLearn{img, orient} = single(zeros(sizeTemplatexSubsample, sizeTemplateySubsample)); 
    end
end
%% Generate multi-resolution images and initialize SUM1, MAX1 and SUM2 maps
numResolution = 9;  % number of resolutions to search for in detection stage
originalResolution = 5; % original resolution is the one at which the imresize factor = 1, see 11th line beneath this line 
allSizeMatrixx = zeros(numImage, numResolution); allSizeVectorx = zeros(1, numResolution*numImage); 
allSizeMatrixy = zeros(numImage, numResolution); allSizeVectory = zeros(1, numResolution*numImage); 
ImageMultiResolution = cell(1, numResolution*numImage); 
MAX1mapFind = cell(numResolution*numImage, numOrient);  % data structures for detection
trackMapFind = cell(numResolution*numImage, numOrient);
SUM2map = cell(1, numResolution*numImage);
translatedTemplate = cell(1, numResolution*numImage); 
averageMap = cell(1, numResolution*numImage); % average within scanning window
for (img = 1:numImage)
  for(j=1:numResolution)
    resolution = .6+(j-1)*.1; % so that .6+(originalResolution-1)*.1 = 1
    ind = (img-1)*numResolution+j; 
    ImageMultiResolution{ind} = imresize(I{img}, resolution, 'nearest');  % images at multiple resolutions
    [sizex, sizey] = size(ImageMultiResolution{ind}); 
    allSizeMatrixx(img, j) = sizex; allSizeVectorx(ind) = sizex; 
    allSizeMatrixy(img, j) = sizey; allSizeVectory(ind) = sizey; 
    for (orient = 1:numOrient) 
        MAX1mapFind{ind, orient} = single(zeros(floor(sizex/subsample), floor(sizey/subsample)));
        trackMapFind{ind, orient} = single(zeros(floor(sizex/subsample), floor(sizey/subsample)));
    end
    SUM2map{ind} = single(zeros(floor((sizex+sizeTemplatex)/subsample), floor((sizey+sizeTemplatey)/subsample)));  
    translatedTemplate{ind} = single(zeros(sizex, sizey));  
    averageMap{ind} = single(zeros(sizex, sizey));  % map of within window average for normalization
  end
end
disp(['start filtering training images at all resolutions']); tic
SUM1mapFind = ApplyFilterfft(ImageMultiResolution, allFilter);  % filtering images at multiple resolutions
disp(['filtering time: ' num2str(toc) ' seconds']);
Caverage(numResolution*numImage, allSizeVectorx, allSizeVectory, numOrient, halfFilterSize, SUM1mapFind, averageMap, halfTemplatex, halfTemplatey, windowNormalizeOrNot);
CgetMAX1track(numResolution*numImage, allSizeVectorx, allSizeVectory, numOrient, locationShiftLimit, orientShiftLimit, subsample, ...
                     SUM1mapFind, MAX1mapFind, trackMapFind);
%% Prepare output variables for learning
selectedOrient = zeros(1, numElement);  % orientation and location of selected Gabors
selectedx = zeros(1, numElement); 
selectedy = zeros(1, numElement); 
selectedlambda = zeros(1, numElement); % weighting parameter for scoring template matching
selectedLogZ = zeros(1, numElement); % normalizing constant
commonTemplate = single(zeros(sizeTemplatex, sizeTemplatey)); % template of active basis 
deformedTemplate = cell(1, numImage); % templates for training images 
for (img = 1 : numImage)
    deformedTemplate{img} = single(zeros(sizeTemplatex, sizeTemplatey));  
end
%% Initialize learning from the first image 
numImage0 = 1; locationShiftLimit0 = 0; orientShiftLimit0 = 0; 
allSizex0 = sizeTemplatex + zeros(1, numImage); allSizey0 = sizeTemplatey + zeros(1, numImage); 
SUM1mapLearn0 = cell(1, numOrient); 
MAX1mapLearn0 = cell(1, numOrient); 
trackMapLearn0 = cell(1, numOrient); 
ind = (starting-1)*numResolution+originalResolution; 
sizex = allSizeVectorx(ind); sizey = allSizeVectory(ind); 
ave = averageMap{ind}; 
averageHere = ave(startx+halfTemplatex, starty+halfTemplatey); % average within local window
for (orient = 1 : numOrient)      
       SUM1mapLearn0{1, orient} = single(zeros(sizeTemplatex, sizeTemplatey)); 
       Ccopy(SUM1mapLearn0{1, orient}, SUM1mapFind{ind, orient}, startx-1, starty-1, sizeTemplatex, sizeTemplatey, sizex, sizey, averageHere, saturation); 
       MAX1mapLearn0{1, orient} = single(zeros(sizeTemplatexSubsample, sizeTemplateySubsample)); 
       trackMapLearn0{1, orient} = single(zeros(sizeTemplatexSubsample, sizeTemplateySubsample)); 
end
CgetMAX1track(numImage0, allSizex0, allSizey0, numOrient, locationShiftLimit0, orientShiftLimit0, subsample, ...
                     SUM1mapLearn0, MAX1mapLearn0, trackMapLearn0);
deformedTemplate0{1} = single(zeros(sizeTemplatex, sizeTemplatey));
disp(['start from single image learning']); tic
CsharedSketch(numOrient, locationShiftLimit0, orientShiftLimit0, subsample, ... % about active basis  
       numElement, numImage0, sizeTemplatex, sizeTemplatey, ...
       SUM1mapLearn0, MAX1mapLearn0, trackMapLearn0, ... % about training images 
       halfFilterSize, Correlation, allSymbol(1, :), ... % about filters
       numStoredPoint, storedlambda, storedExpectation, storedLogZ, ... % about exponential model 
       selectedOrient, selectedx, selectedy, selectedlambda, selectedLogZ, ... % learned parameters
       commonTemplate, deformedTemplate0); % learned templates 
disp(['mex-C learning time: ' num2str(toc) ' seconds']);
Outepsgif(-double(commonTemplate), 'ex5c_monkey_sym0'); 
%% Iteration part 1: detect the object in each image 
for(it = 1:numIteration)
    disp(['detection for step ' num2str(it)]);
    for (img=1:numImage)
      if (img~=starting)
        disp(['    start detecting in image ' num2str(img)]); tic
        MAX2score = single(zeros(1, numResolution)); Fx = zeros(1, numResolution); Fy = zeros(1, numResolution); 
        allSizex = allSizeMatrixx(img, :); allSizey = allSizeMatrixy(img, :); 
        CgetSUM2(img, numImage, numResolution, allSizex, allSizey, numOrient, ...
              locationShiftLimit, orientShiftLimit, subsample, ...
              sizeTemplatex, sizeTemplatey, halfFilterSize, allSymbol(1, :), ...
              numElement, selectedOrient, selectedx, selectedy, selectedlambda, selectedLogZ, ...
              MAX1mapFind, trackMapFind, SUM2map, translatedTemplate, MAX2score, Fx, Fy, ...
              averageMap, saturation);
        disp(['mex-C finding time: ' num2str(toc) ' seconds']);
        [maxOverResolution j] = max(MAX2score);   % most likely resolution
        maxFx = Fx(j); maxFy = Fy(j); 
      else
        j = originalResolution; 
        maxFx = startx-1; maxFy = starty-1; 
      end
        ind = (img-1)*numResolution+j; 
        sizex = allSizeVectorx(ind); sizey = allSizeVectory(ind); 
        ave = averageMap{ind}; 
        xHere = maxFx+halfTemplatex+1; yHere = maxFy+halfTemplatey+1; 
        if ((xHere>=1)&&(xHere<=sizex)&&(yHere>=1)&&(yHere<=sizey))
            averageHere = ave(xHere, yHere);  % average within local window
        else
            averageHere = single(1.); 
        end
        Jshow{img} = ImageMultiResolution{ind}+translatedTemplate{ind}*100.;  % overlaid sketch 
        for (orient = 1 : numOrient)
               Ccopy(SUM1mapLearn{img, orient}, SUM1mapFind{ind, orient}, maxFx, maxFy, sizeTemplatex, sizeTemplatey, sizex, sizey,  averageHere, saturation); 
              % Ccopy(MAX1mapLearn{img, orient}, MAX1mapFind{ind, orient}, round(maxFx/subsample), round(maxFy/subsample), sizeTemplatexSubsample, sizeTemplateySubsample, floor(sizex/subsample), floor(sizey/subsample),  averageHere, saturation); 
              % Ccopy0(trackMapLearn{img, orient}, trackMapFind{ind, orient}, round(maxFx/subsample), round(maxFy/subsample), sizeTemplatexSubsample, sizeTemplateySubsample, floor(sizex/subsample), floor(sizey/subsample)); 
        end
    end
%% Iteration part 2: re-learn the template
    disp(['multiple image learning for step ' num2str(it)]);
    tic
    CgetMAX1track(numImage, allSizex0, allSizey0, numOrient, locationShiftLimit, orientShiftLimit, subsample, ...
                     SUM1mapLearn, MAX1mapLearn, trackMapLearn);
    CsharedSketch(numOrient, locationShiftLimit, orientShiftLimit, subsample, ... % about active basis  
           numElement, numImage, sizeTemplatex, sizeTemplatey, ...
           SUM1mapLearn, MAX1mapLearn, trackMapLearn, ... % about training images 
           halfFilterSize, Correlation, allSymbol(1, :), ... % about filters
           numStoredPoint, storedlambda, storedExpectation, storedLogZ, ... % about exponential model 
           selectedOrient, selectedx, selectedy, selectedlambda, selectedLogZ, ... % learned parameters
           commonTemplate, deformedTemplate); % learned templates 
    disp(['mex-C learning time: ' num2str(toc) ' seconds']);
    Outepsgif(-double(commonTemplate), ['ex5c_monkey_sym' num2str(it)]); 
end
%% display and save results
tmp(startx:endx, starty:endy) = tmp(startx:endx, starty:endy) + deformedTemplate{starting}*100.; 
Jshow{starting} = tmp; 
Outepsgif(-double(commonTemplate), 'ex5c_monkey_symfinal'); 
for(img=1:numImage)
       Outepsgif(double(Jshow{img}), ['ex5c_monkey' num2str(img)]);  % overlaid sketch 
end
save 'learnedTemplate' numElement selectedOrient selectedx selectedy selectedlambda selectedLogZ ...
      numOrient halfFilterSize allFilter allSymbol locationShiftLimit orientShiftLimit saturation subsample ...
      resizeFactor numResolution;

% CleanUp; 
  
  