trackFile.m

%TRACKFILE track a mouse in the openfield experiment
%
%TRACKFILE(FILENAME, ARENASIZE, THRESH, BORDER)
%
% FILENAME the input video file
% ARENASIZE the size of the (square) experimental arena in mm
% THRESH the threshold to apply for segmentation (0-1)
% BORDER an additional border that is substracted from the image on all
%    sides to avoid tracking problems due to minimal misalignements
%
% the function will generate a .mat file containing the x,y position
% of the mouse (mousePos), the affine transformation matrix (tform)
% that was used to warp the image based on the four corners of the arena
% and the given size of the arena (arenaSize).
%
% Friedrich Kretschmer
% Max-Planck Institute for Brain Research
% sciclist@brain.mpg.de

function trackFile(fileName, arenaSize, thresh, border)
    %Save a video containing the tracking result
    doSaveDemo = true;

    [pathstr, name, ~] = fileparts(fileName);

    vid = VideoReader(fileName);
    
    %First frame is used to estimate the affine transformation matrix
    im = vid.read(1);

    %Threshold image
    bim = im2bw(im, thresh);
    
    %Estimate affine transformation
    tform = calculateTForm(bim, arenaSize);

    nFrames = round(vid.Duration * vid.FrameRate);
    frameRate = vid.FrameRate;
    mousePos = NaN(nFrames, 2);
    
    if doSaveDemo
        vidWriter = VideoWriter([pathstr, filesep, name, '_tracked.avi']);
        open(vidWriter);
        fig = figure('NumberTitle', 'off', 'Position', [100,100,482,500]);
        ax = axes('Parent', fig, 'Units', 'normalized', 'Position', [0,0,1,1]);
        pim = imshow(zeros(arenaSize-(2*border), arenaSize-(2*border), 3), 'Parent', ax);
        hold on
        p = plot(0,0, 'r+', 'markerSize', 14, 'LineWidth', 2);
        axis equal
    end
    
    for frameIdx = 1:nFrames
        try
            im = vid.read(frameIdx);
            
            %Threshold image
            bim = im2bw(im, thresh);

            %And warp the image according to the affine transformation
            wim = imwarp(bim, tform, 'OutputView', imref2d([arenaSize, arenaSize]));

            %Remove border from image (might not be perfectly aligned and
            %black affects tracking
            bwim = wim((border+1):end-border,(border+1):end-border);
            
            % Detect the mouse
            CC = bwconncomp(~bwim);
            %Deteremine the number of pixels belonging to each component
            numPixels = cellfun(@numel, CC.PixelIdxList);
            [~, idx] = max(numPixels);

            %Determine the largest blob (should be the mouse)
            if ~isempty(idx)
                CC.NumObjects = 1;
                CC.PixelIdxList = {CC.PixelIdxList{idx}};
                stats = regionprops(CC, 'Centroid');
                mousePos(frameIdx, :) = stats.Centroid + [border+1, border+1];
            else
                mousePos(frameIdx,:) = [NaN,NaN];
            end
            
            %Generate an image overlayed with the tracking result
            if doSaveDemo
                wim = imwarp(im, tform, 'OutputView', imref2d([arenaSize, arenaSize]));
                set(pim, 'CData', wim((border+1):end-border,(border+1):end-border,:));
                set(fig, 'Name', [name, ': ', num2str(frameIdx),'/',num2str(nFrames)]);
                hold on
                if ~isempty(idx)
                    set(p, 'XData', stats.Centroid(1), 'YData', stats.Centroid(2));
                else
                    set(p, 'XData', 0, 'YData', 0);
                end
                drawnow();
                writeVideo(vidWriter, getframe(fig));
            end
        catch exception
            %The number of frames is often incorrectly read. 
            if(strcmp(exception.identifier, 'MATLAB:audiovideo:VideoReader:invalidFrameVarFrameRate'))
                %Just interrupt if file ends early
                break
            else
                rethrow(exception);
            end
        end
    end

    %Save results
    save([pathstr, filesep, name], 'mousePos', 'tform', 'arenaSize', 'frameRate');
    
    if doSaveDemo
        close(vidWriter);
        close(fig);
    end
end