% Directory setting
frame_directory = 'data/frames/';
mask_directory = 'data/predicts/';
video_names = dir( [frame_directory 'sample*'] );

% Enhancing figure visibility
set(gcf, 'Position', [0 0 800 800], 'Color', [1 1 1] );

% Save result
csv_file_directory = 'data/result_csv/';
if ~isdir(csv_file_directory), mkdir(csv_file_directory); end

% Figure to video 
output_video_option = 0;
if output_video_option
    output_video_directory = 'data/output/';
    if ~isdir(output_video_directory), mkdir(output_video_directory); end
end

% Main loop
for videoIdx = 1:length( video_names )
    frame_names = dir( [frame_directory video_names(videoIdx).name '/*.jpg']);
    mask_names = dir( [mask_directory video_names(videoIdx).name '/*.jpg']);
    
    %% For one video
    xy = nan([2,length( frame_names )]);
    bb = nan([4,length( frame_names )]);
    for frameIdx = 1:length( frame_names )
        
        %% Single frame loading
        frame_name = [ frame_names(frameIdx).folder ...
            '/' frame_names(frameIdx).name];
        mask_name = [ mask_names(frameIdx).folder ...
            '/' mask_names(frameIdx).name];
        
        frame = double(imread( frame_name ))/255; % Scaling to 0-1
        mask = double(imread( mask_name ))/255; % Scaling to 0-1
        
        % Manual artifact rejection
        mask(:,1:30) = 0;

        % Visualization
        
        % Original frame
        sp1 = subplot(2,2,1); hold off
        imagesc( frame );
        set(gca, 'Box', 'off', 'LineWidth', 2, 'XTick', [], 'YTick', [], 'FontSize', 15);
        title( sprintf( '(1) Frame #%04d', frameIdx ));
        colormap(sp1, 'gray');
        
        % Mask 
        sp2 = subplot(2,2,2); hold off
        imagesc( mask' );
        set(gca, 'Box', 'off', 'LineWidth', 2, 'XTick', [], 'YTick', [], 'FontSize', 15);
        title( sprintf( '(2) Mask #%04d', frameIdx ));
        colormap(sp2, 'gray');
        
        % Overlay
        overlay = repmat( frame, [1 1 3] );
        overlay(:,:,1) = 0.5*mask';
        
        sp3 = subplot(2,2,3); hold off
        imagesc( overlay );
        set(gca, 'Box', 'off', 'LineWidth', 2, 'XTick', [], 'YTick', [], 'FontSize', 15);
        title( sprintf( '(3) Overlay #%04d', frameIdx ));
        
        %% Extracting image statistics using regionprops
        probability_threshold = 0.5; % Probability, anonymous unit (0-1)
        mask_binary = mask > probability_threshold;
        stats = regionprops( mask_binary', 'BoundingBox', 'Area', 'Centroid', 'Orientation' );
        
        size_threshold = 100; % Pixel
        drawing_regionprops = false;
        
        big_enough = find( [stats.Area]>size_threshold );
        if ~isempty( big_enough )
            stat = stats(big_enough);
            xy(:,frameIdx) = stat.Centroid;
            bb(:,frameIdx) = stat.BoundingBox;
        else % if failed
            xy(:,frameIdx) = xy(:,frameIdx-1);
            bb(:,frameIdx) = bb(:,frameIdx-1);
        end
            
        sp4 = subplot(2,2,4); hold off
        imagesc( frame );
        set(gca, 'Box', 'off', 'LineWidth', 2, 'XTick', [], 'YTick', [], 'FontSize', 15);
        title( sprintf( '(4) Tracking result #%04d', frameIdx ));
        colormap(sp4,'gray');
        hold on;
        % Centroid
        plot( xy(1,frameIdx), xy(2,frameIdx), 'g*', 'MarkerSize', 20);
        % BoundingBox
        rectangle('Position', bb(:,frameIdx),'EdgeColor','g','LineWidth',2 )
        % Text annotation
        text_contents = sprintf( 'X: [%.1f], Y: [%.1f]',  xy(:,frameIdx));
        text( xy(1,frameIdx),xy(2,frameIdx) - 50, text_contents, 'Color', 'g',...
            'HorizontalAlignment', 'center' );
        
        %% Etc.
        drawnow;
        
        % Writing output video
        if output_video_option
            if frameIdx==1
                try vid_obj.close(); clear vid_obj; end
                vid_obj = VideoWriter( [output_video_directory ...
                    video_names(videoIdx).name '.mp4'], 'MPEG-4');
                vid_obj.open();
            end
            capture_frame = getframe(gcf);
            vid_obj.writeVideo(capture_frame.cdata);
            if frameIdx==length( frame_names )
                vid_obj.close();
            end
        end
        
    end
    
    % Save coordinate data into *.csv format
    csv_filename = [csv_file_directory 'result_' video_names(videoIdx).name '.csv'];
    csvwrite( csv_filename, xy );
    
end


%% Sample visualization
csv_files = dir( [ csv_file_directory '*.csv'] );
for fileIdx = 1:length(csv_files)
    
    %% 1. Trace
    % Get coordinate
    xy = csvread( [csv_file_directory csv_files( fileIdx ).name] )/448;
    
    % Calc speed
    speed = [0; smooth( sum(abs(diff(xy,2)),1), 5 )];
    speed = speed/max(speed);
    outliers = find( speed > (nanmean(speed)+nanstd(speed)*3) );
    xy(outliers) = nan;
    
    % Plot
    subplot( 2,4, fileIdx ); hold off
    plt = plot( xy(1,:), xy(2,:), 'k.');

    % Enhance visibility
    set(gca, 'XDir', 'normal', 'YDir', 'reverse');
    set(gca, 'XTick', [], 'YTick', [], 'LineWidth', 2, 'Box', 'on', 'FontSize', 15 );
    axis([0 1 0 1]);
    title( sprintf( 'Trace (File %1d)', fileIdx ) );
    
    %% 2. Histogram
    % Get coordinate
    edges = linspace( 0, 1, 100 );
    [n,xedge,yedge] = hb_histcount2( xy(1,:), xy(2,:), edges, edges );
    try,n=imgaussfilt(n,1);end
    
    % Plot
    subplot( 2,4, 4+fileIdx ); hold off
    imagesc( xedge, yedge, log(n)' );

    % Enhance visibility
    set(gca, 'XDir', 'normal', 'YDir', 'reverse');
    set(gca, 'XTick', [], 'YTick', [], 'LineWidth', 2, 'Box', 'on', 'FontSize', 15 );
    title( sprintf( 'Occupancy (File %1d)', fileIdx ) );
    axis([0 1 0 1]);
    colormap hot;
end

% Subfunction
function [n,xedge,yedge,idx] = hb_histcount2(X, Y, xedge, yedge)
nbin = [length(xedge)-1, length(yedge)-1];
for binIdx_x = 1:nbin(1)
    xIdx = (X>xedge(binIdx_x)) .* (X<=xedge(binIdx_x+1));
    for binIdx_y = 1:nbin(2)
        yIdx = (Y>xedge(binIdx_y)) .* (Y<=xedge(binIdx_y+1));
        jointIdx = find( xIdx .* yIdx );
        n( binIdx_x, binIdx_y ) = length(jointIdx);
        idx{binIdx_x, binIdx_y} = jointIdx;
    end
end
end