% Author: Cem Emre Akbas
% e-mail: akbas@fi.muni.cz
% Center for Biomedical Image Analysis
% Faculty of Informatics, Masaryk University
% Written & Tested on Matlab 2017a

%% computes SIMPLE segmentation estimation based on global weighted voting
% see paper
% "Label fusion in atlas-based segmentation using a selective and
% iterative method for performance level estimation (SIMPLE)."
% Langerak et. al. (2010)
%
clc; clear all; close all;
tStart=tic;
%% Parameters to be changed
% Do not forget to change the dataset number BEFORE EACH RUN 
dataset_number = 6;     % should be between 1-26 (see lines 22-47 for dataset names)

% parent_directory should be the parent directory of all available results
% It should contain alphabetically ordered folders named as isbiXXX which respresent CTC participants
% Each isbiXXX folder contains submitted results by that participant in
% different subfolders that are named as <dataset_name> (see line 58)
% Change the parent_directory ONLY BEFORE THE FIRST RUN
parent_directory = 'C:/Users/ceakb/Documents/CTC_Training_Results/';

%% Constants
% SEG scores of CTC participants - in the order: isbi001, isbi005, isbi011, isbi014, isbi020, isbi038, isbi043, isbi094, isbi115_01, isbi115_02, isbi170_01, isbi170_02, isbi170_03, isbi170_04, isbi263, isbi333, isbi442, isbi496, isbi496, isbi765, isbi770, isbi926
weights = cell(26,2);
weights{1,1} = '/DIC-C2DH-HeLa/01_RES/'; weights{1,2} = [0.3451, 0.8918, 0.3441, 0.4880];
weights{2,1} = '/DIC-C2DH-HeLa/02_RES/'; weights{2,2} = [0.1252, 0.8756, 0.1998, 0.5452];
weights{3,1} = '/Fluo-C2DL-MSC/01_RES/'; weights{3,2} = [0.4212, 0.1930, 0.3815, 0.6332, 0.0101, 0.2161, 0.5214, 0.4190, 0.3211, 0.3591, 0.4496, 0.0098];
weights{4,1} = '/Fluo-C2DL-MSC/02_RES/'; weights{4,2} = [0.5296, 0.1881, 0.2635, 0.6357, 0.1806, 0.1834, 0.7081, 0.4757, 0.4306, 0.5979, 0.5379];
weights{5,1} = '/Fluo-C3DH-H157/01_RES/'; weights{5,2} = [0.3684, 0.4584, 0.7449, 0.8911, 0.3864, 0.6010, 0.7525];
weights{6,1} = '/Fluo-C3DH-H157/02_RES/'; weights{6,2} = [0.1739, 0.5567, 0.4478, 0.7793, 0.5486, 0.4983, 0.5729, 0.5317];
weights{7,1} = '/Fluo-C3DL-MDA231/01_RES/'; weights{7,2} = [0.2668, 0.3484, 0.4580, 0.6836, 0.5268, 0.1998, 0.1957, 0.6521];
weights{8,1} = '/Fluo-C3DL-MDA231/02_RES/'; weights{8,2} = [0.6539, 0.4286, 0.5404, 0.7067, 0.5748, 0.2775, 0.2898, 0.6484];
weights{9,1} = '/Fluo-N2DH-GOWT1/01_RES/'; weights{9,2} = [0.6996, 0.5158, 0.3686, 0.7032, 0.3692, 0.3253, 0.7724, 0.5280, 0.5232, 0.4483, 0.7221, 0.5448, 0.7461];
weights{10,1} = '/Fluo-N2DH-GOWT1/02_RES/'; weights{10,2} = [0.8126, 0.8227, 0.5559, 0.7980, 0.6866, 0.4782, 0.9413, 0.7693, 0.8537, 0.5335, 0.9140, 0.8980, 0.8650];
weights{11,1} = '/Fluo-N2DL-HeLa/01_RES/'; weights{11,2} = [0.6525, 0.6927, 0.6689, 0.5183, 0.6267, 0.8270, 0.4557, 0.2026, 0.8251, 0.4594, 0.4203, 0.6498, 0.7435, 0.7772];
weights{12,1} = '/Fluo-N2DL-HeLa/02_RES/'; weights{12,2} = [0.7248, 0.7302, 0.7987, 0.5570, 0.7085, 0.8457, 0.4703, 0.2805, 0.8469, 0.5099, 0.3195, 0.7272, 0.8140, 0.6205];
weights{13,1} = '/Fluo-N3DH-CE/01_RES/'; weights{13,2} = [0.3887, 0.0955, 0.3847, 0.6346, 0.5662];
weights{14,1} = '/Fluo-N3DH-CE/02_RES/'; weights{14,2} = [0.3818, 0.0592, 0.3546, 0.5924, 0.4858];
weights{15,1} = '/Fluo-N3DH-CHO/01_RES/'; weights{15,2} = [0.5319, 0.7824, 0.6250, 0.7043, 0.8518, 0.6600, 0.8447, 0.8188, 0.5990, 0.8138, 0.8167];
weights{16,1} = '/Fluo-N3DH-CHO/02_RES/'; weights{16,2} = [0.6123, 0.8622, 0.6817, 0.7480, 0.9119, 0.8264, 0.9204, 0.8202, 0.7579, 0.9034, 0.8787];
weights{17,1} = '/Fluo-N3DL-DRO/01_RES/'; weights{17,2} = [0.0010, 0.0863, 0.4935];
weights{18,1} = '/Fluo-N3DL-DRO/02_RES/'; weights{18,2} = [0.0035, 0.2446, 0.5638];
weights{19,1} = '/PhC-C2DH-U373/01_RES/'; weights{19,2} = [0.8748, 0.3562, 0.9387, 0.2536, 0.8138];
weights{20,1} = '/PhC-C2DH-U373/02_RES/'; weights{20,2} = [0.7567, 0.3587, 0.8315, 0.2488, 0.7285];
weights{21,1} = '/PhC-C2DL-PSC/01_RES/'; weights{21,2} = [0.5453, 0.5417, 0.5209, 0.5136, 0.5911, 0.0479, 0.6206, 0.4635];
weights{22,1} = '/PhC-C2DL-PSC/02_RES/'; weights{22,2} = [0.5657, 0.6217, 0.0898, 0.4772, 0.5926, 0.0511, 0.6690, 0.4647];
weights{23,1} = '/Fluo-N2DH-SIM+/01_RES/'; weights{23,2} = [0.7140, 0.7988, 0.8071, 0.7778, 0.6711, 0.8489, 0.3804, 0.8652, 0.8260, 0.7864, 0.8508, 0.8440, 0.8640];
weights{24,1} = '/Fluo-N2DH-SIM+/02_RES/'; weights{24,2} = [0.3910, 0.5790, 0.5740, 0.3110, 0.4753, 0.6406, 0.3438, 0.6577, 0.1432, 0.4047, 0.3520, 0.6298, 0.6055];
weights{25,1} = '/Fluo-N3DH-SIM+/01_RES/'; weights{25,2} = [0.6533, 0.7348, 0.5867, 0.7754, 0.8359, 0.7895, 0.7251];
weights{26,1} = '/Fluo-N3DH-SIM+/02_RES/'; weights{26,2} = [0.3186, 0.3166, 0.3185, 0.3784, 0.6261, 0.3214, 0.5306];

% threshold parameter for simple segmentation
alpha = 9;

% number of reconsideration iterations
x = 1;

% iteration counter
k = 1;

%% Define the dataset and find all submitted results on it
% Change the dataset number after each run
dataset_name = weights{dataset_number, 1};
initialWeights = weights{dataset_number, 2};

ctc_result_folders = [];
subdirs = dir(parent_directory);
num_dirs = length(subdirs);
for ii = 1:num_dirs
    data_dir = [parent_directory, subdirs(ii).name, '/', dataset_name];
    if exist(data_dir,'dir') == 7 % check if this subdirectory exists (i.e. has results for this dataset)
        ctc_result_folders = [ctc_result_folders; string([parent_directory, subdirs(ii).name])]; % collect the paths of available result folders for this dataset
    end
end                 

imlist_path = (strcat(ctc_result_folders(1), dataset_name));
imlist = dir(fullfile(char(imlist_path), '*.tif'));
num_of_images = length(imlist);
num_of_results = length(ctc_result_folders);

for ii = 1:num_of_images

    for j = 1:num_of_results
      Segmentations(j).Volume = imread(char(strcat(ctc_result_folders(j), dataset_name, imlist(ii).name))); %Read the jth image
    end

    % outer loop: image numbers
    % inner loop: result folders


    %% real code

    % struct for selected segmentations in each iteration
    SelectedSegmentations                          = struct;
    idxSelected                                    = 1 : length(Segmentations);
    SelectedSegmentations(k).SelectedSegmentations = idxSelected;

    % propagate labels to subject (initial ground truth estimation)
    CurrentSegmentations = Segmentations;

    GroundTruthEstimate = globalWeightedVoting(CurrentSegmentations, initialWeights);       

    boolNoConsensusFound=0;


    %% iterate until the number of selected segmentations converges
    boolSelectionChanged = 1;

    while ( boolSelectionChanged )


        %% (1) estimate performance of each segmentation
        performances = NaN(length(CurrentSegmentations),1);
        for iSegmentation = 1 : length(CurrentSegmentations)
            performances(iSegmentation) = segmentationDice(GroundTruthEstimate, CurrentSegmentations(iSegmentation).Volume);
        end

        SelectedSegmentations(k).performances = performances;

        %% (2) keep only well performing segmentations

        % compute threshold
        theta = nanmean(performances) - alpha * nanstd(performances);

        % select well performing segmentations
        idxSelected = SelectedSegmentations(k).SelectedSegmentations(performances > theta);

        if ~isempty(idxSelected)
            % increase iteration count
            k = k + 1;


            %% (3) compute weighted majority voting for ground truth update
            GroundTruthEstimate = globalWeightedVoting(Segmentations(idxSelected), performances(performances > theta));


            %% re-consider early discarded segmentations in the first x iterations
            if ( x > 0 )

                x = x - 1;

                CurrentSegmentations = Segmentations;
                idxSelected          = 1 : length(Segmentations);
                SelectedSegmentations(k).SelectedSegmentations = idxSelected;

                % compute convergence criterion
                boolSelectionChanged = 1;
            else

                SelectedSegmentations(k).SelectedSegmentations = idxSelected;
                CurrentSegmentations                           = Segmentations(idxSelected);

                % compute convergence criterion
                boolSelectionChanged = ( numel(SelectedSegmentations(k-1).SelectedSegmentations) ~= numel(SelectedSegmentations(k).SelectedSegmentations) );
            end
        else
            boolSelectionChanged=0;
            boolNoConsensusFound = 1;
        end
    end

    if boolNoConsensusFound
        SelectedSegmentations = [];

    else
        SelectedSegmentations = SelectedSegmentations(1:end-1);
    end

    ctc_folder = "C:\Users\ceakb\Documents\CTC_Training_Results\";
    output_folder = strcat(ctc_folder, "SIMPLE_Results\SEG", dataset_name);
    if exist(output_folder,'dir') == 0
        mkdir(char(output_folder));
    end
    LabeledGroundTruthEstimate = bwlabel(GroundTruthEstimate,4);
    imwrite(uint16(LabeledGroundTruthEstimate), char(strcat(output_folder, imlist(ii).name)));
end
tElapsed=toc(tStart);
splitted_name = strsplit(dataset_name, '/');
fid = fopen([char(output_folder), 'Report_', splitted_name{2}, '-', splitted_name{3}, '.txt'],'w');
fprintf(fid,'Dataset: %s \nComputation time: %d seconds \nCreated %d fused segmentation masks with SIMPLE algorithm. \nUsed %d CTC participant results stored in the paths below:\n', dataset_name, tElapsed, num_of_images, num_of_results);       %# Print the string
for j = 1:num_of_results
    fprintf(fid,'%s \n', ctc_result_folders(j));
end
fclose(fid);

function [ SegmentationEstimation ] = globalWeightedVoting(Segmentations, weights)
%% computes per voxel global weighted majority voting (one weight for the entire image region)
%
% Inputs:
%  Segmentations, struct-array with size 1xN
%   Segmentations(i).Volume, matrix: the segmentations estimation (label image)
%  weights, 1xN double vector: the weiths for the corresponding
%  segmentations
%
% Outputs:
%  SegmentationEstimation, matrix: the resulting label image
%
%
% Author: Matthias Dorfer @ CIR Lab (Medical University of Vienna)
% email:  matthias.dorfer@meduniwien.ac.at
%

    % image or volume that counts how often a voxel was suggested as
    % foreground
    MajorityCounter = zeros(size(Segmentations(1).Volume));

    % normalize weiths to range [0 1]
    weights = weights / max(weights);
    
    % iterate all segmentations
    for iSeg = 1 : length ( Segmentations )
    
        WeightedSegmentation = (Segmentations(iSeg).Volume > 0) * weights(iSeg);
        MajorityCounter      = MajorityCounter + WeightedSegmentation;
    end
    
    % estimate segmentation based on majority voting
    SegmentationEstimation = MajorityCounter >= (sum(weights) / 2);
    SegmentationEstimation = double(SegmentationEstimation);
end
