%_______________________________________________________________________________
%
%   MixtureModel
%
%   This function implements the Gaussian mixture model explained in:
%   Stauffer, Grimson.  Learning Patterns of Activity Using Real-Time
%   Tracking. PAMI. Volume 22. Number 8. August 2000.
%
%   NOTE: This code is a bit convoluted due to heavly performance oriented
%   code being used.  The intuitive scheme for this algorithm has been
%   heavily modified to lower cache misses in memory to improve speed!
%   
%   TODO: Fill the inputs and outputs in.
%   [bSuccess] = MixtureModel(strFilename, fAlpha, iK, iT)
%   Input:
%`  Output: 1,0 depending on function success.
%
%   Author:         Isaac Gerg
%   Version:        0.02
%   Revision Date:  3.15.2005
%
%   0.02 3.15.2005  idg - Modified code for better performance.
%   0.01 2.8.2005   Isaac Gerg(idg) - Initial release.
%_______________________________________________________________________________
%old main0 -> 0.001 2, 0.5, 4.0
%patee -> same

clear; clc;

% Parameters
strInputFilename = 'video/hub2.avi';
strOutputFilename = 'hub2-test.avi';
fAlpha = 0.01;
iK = 3;
fT = 0.8;

% Read in video sequence
movInput = aviread(strInputFilename);
MovieInfo = aviinfo(strInputFilename);

% Specify the size of the movie we want to process.
iVideoWidth = 160;
iVideoHeight = 120;
iNumberOfFrames = length(movInput);

% Intialize pixel gaussian pamameters.
% [weight, mu, sigma^2]
% Start with a small sigma value so we do not get division by zero errors
% later on.  This value could probably be set to 5 intially, but is not for
% debugging puposes.
WEIGHT = 1; MU = 2; SIGMA = 3;
pixelParameters = zeros(iK*3, iVideoHeight * iVideoWidth);
pixelParameters(SIGMA:3:iK*3, :) = 0.001;

% Initialize list used to sort a pixel's gaussians.
distributionList = ones(iK, 5);

% Setup output movie stream
movOutput = avifile(strOutputFilename);
    movOutput.Quality = 100;
    movOutput.FPS = 15;

% Record time to determine frames/minute.
start = clock;

for i=1:iNumberOfFrames
    
    % Fetch frame of video and begin processing.
    % Convert RGB to grayscale.
    imInput = movInput(i).cdata;
    imInput = imresize(imInput, [iVideoHeight iVideoWidth]);
    imInput = double(rgb2gray(imInput));

    i
    
    % Intialialize output frame.
    imOutput = ones(iVideoHeight, iVideoWidth, 3);
    
    % Uncomment the following lines if you want an overlay of the
    % segmentation on the actual video.
    %imOutput(:,:,1) = imInput(:,:,1);
    %imOutput(:,:,2) = imInput(:,:,1);
    %imOutput(:,:,3) = imInput(:,:,1);
    
    % Convert 2D input frame to 1D array. This is done for performance
    % considerations.
    imInput = reshape(imInput, iVideoHeight*iVideoWidth, 1);
        
    % Update mu and sigma^2
    for x=1:iVideoHeight*iVideoWidth 
        % Initialize pixel distribution parameters.
        Xt = imInput(x);       
        iSelectedDistribution = -1;
        iMinWeightDistribution = -1;

        fMinWeight = 9999;
        % See what distrubution the current pixel matches up with best.
        % Find the distribution with the min weight.  While we are at it,
        % begin to update the weights of all distributions.
        for z=0:iK-1
            fWeight = pixelParameters(3*z+WEIGHT, x);
            fMean = pixelParameters(3*z+MU, x);
            fSigma = pixelParameters(3*z+SIGMA, x);
            if (fWeight < fMinWeight)
                fMinWeight = fWeight;
                iMinWeightDistribution = z;
            end
            pixelParameters(3*z+WEIGHT, x) = (1 - fAlpha)*fWeight;
            if (abs(Xt-fMean)  <= 5.0*fSigma)
                iSelectedDistribution = z;
            end;                   
        end    
        % See if matched distribution exists. If not, create one.
        if (iSelectedDistribution == -1)
            % Remove the lowest weight distribution and replace with a new
            % one.
            % 'Low weight' is defined here as 0.05.  This can be tuned.
            pixelParameters(3*iMinWeightDistribution+WEIGHT, x) = 0.05;
            pixelParameters(3*iMinWeightDistribution+SIGMA, x) = 5;
            pixelParameters(3*iMinWeightDistribution+MU, x) = Xt;   
            iSelectedDistribution = iMinWeightDistribution;
        else
            % A suitable distribution exists.  Update its parameters given
            % the new point, Xt.
            % Finish updating the weight of the selected distribution.
            z=iSelectedDistribution;
            pixelParameters(3*z+WEIGHT, x) = pixelParameters(3*z+ WEIGHT, x) + fAlpha;       
           
            % Calculated the new mean and sigma.
            fMean = pixelParameters(3*z+MU,x);
            fSigma = pixelParameters(3*z+SIGMA, x);
            covarianceMatrix = fSigma; %* eye(iK);
            nu = (1/((2*pi)^0.5 * sqrt(det(covarianceMatrix)))) * exp(-0.5 * (Xt - fMean)' * inv(covarianceMatrix) * (Xt - fMean));
            rho = fAlpha * nu;
            fNewMean = (1 - rho) * fMean + (rho * Xt);
            fNewSigma = (1 - rho) * fSigma + rho*(Xt - fNewMean)'*(Xt - fNewMean);
            
            % Update the measurments.
            pixelParameters(3*z+MU, x) = fNewMean;
            pixelParameters(3*z+SIGMA, x) = fNewSigma;                
        end

        % Normalize weights.
        a = 0;
        for z=0:iK-1
            a = a + pixelParameters(3*z+WEIGHT, x);
        end
        f = 1 / a;
        for z=0:iK-1
            pixelParameters(3*z+WEIGHT, x) = pixelParameters(3*z+WEIGHT, x) * f;
        end
        
        % Determine if forground pixel or not. Sort the distributions and
        % then take the first B of them until the summation of their
        % weights is just greater than T.
        for z=0:iK-1
            distributionList(z+1,1:3) = pixelParameters(3*z+1:3*z+3, x)';
            distributionList(z+1,5) = distributionList(z+1, WEIGHT) / distributionList(z+1,SIGMA);
            distributionList(z+1,4) = z+1;
        end
        distributionList = sortrows(distributionList, 5);
        iTemp = iSelectedDistribution+1;
        for z=0:iK-1
            pixelParameters(3*z+1:3*z+3, x) = distributionList(z+1,1:3)';
            if (distributionList(z+1,4) == iTemp)
                iSelectedDistribution = z+1;
            end
        end    
        B = iK;
        fSum  = 0.0;
        while (fSum < fT)
            B = B - 1;
            fSum = fSum + pixelParameters(3*B+WEIGHT,x);
        end
            B = B+1;
        if (iSelectedDistribution < B)
            % Current pixel is foreground.  Label as such.
            imOutput(mod(x-1,iVideoHeight)+1, ceil(x/iVideoHeight),1) = 255;
            imOutput(mod(x-1,iVideoHeight)+1, ceil(x/iVideoHeight),2) = 255;
            imOutput(mod(x-1,iVideoHeight)+1, ceil(x/iVideoHeight),3) = 255;
        end
    end
    
    % Normalize the current output frame and then add it to the output
    % movie.
    imOutput = imOutput ./ 255;
    movOutput = addframe(movOutput, imOutput);
end

stop = clock;
disp('Frames per minute:');
iNumberOfFrames/(etime(stop, start)/60)   

% Close movie when done and report success.
movOutput = close(movOutput)

disp('Done!');
bSuccess = 0;
