limo_STAPLE.m

function [W, p, q] = limo_STAPLE(varargin)

% This code is an adaptation of the original code to work directly with
% outputs from LIMO tools - Vecorized MATLAB implementation of the STAPLE 
% algorithm by Warfield et al. 

% Function:  [W, p, q] = STAPLE(stat_files, threshold)
%
% IMPUT : stat_files is a cell array of statitical maps generated by LIMO
%         threshold to apply to stat_files to get an estimated ground truth
%         (0.5 by default)
%
% OUTPUT: W weight matrix for each voxel
%         p vector of sensitivities for each expert
%         q vector of specificities for each expert
%         staple_mask, a binary image writen on the drive
%
% Reference: Warfield, Simon K., Kelly H. Zou, and William M. Wells. 
%            "Simultaneous truth and performance level estimation (STAPLE): 
%            an algorithm for the validation of image segmentation." 
%            Medical Imaging, IEEE Transactions on 23.7 (2004): 903-921.
%
% Original code edited by Cyril Pernet to work with EEG/LIMO results
% ------------------------------
%  Copyright (C) LIMO Team 2020

%% check inputs
if nargin == 0
    limo_STAPLE
    return
else 
    stat_files = varargin{1};
    if ~iscell(stat_files)
        error('argument in must be a cell array of names')
    end
    
    if nargin == 1
        threshold = 0.5;
    else
        threshold = varargin{2};
    end
    clear varargin
end

if nargin > 2
    warn('more arguments in than needed??? only using arguments 1 and 2')
end

%% how many maps
if size(stat_files,1) == 1
    stat_files=stat_files'; 
end
N = size(stat_files,1);

%% make the matrix DD (keep all clusters separate)
for map = N:-1:1
    data = load(stat_files{map});
    data = data.(cell2mat(fieldnames(data)));
    nclusters(map) = max(data(:));
    if map == 1
       map_dim = size(data); 
    end
    try
        DD(:,N) = data(:);
    catch load_err
       error('couldn''t load map %g\n %s\n',map,load_err.message) 
    end
end


%% STAPLE-Algorithm by Warfield et al. for binary segementations
% --> Andreas Husch's code
     
%% Parameters
MAX_ITERATIONS = 30;
EPSILON = 0.00001; % convergence criterion

for cluster = max(nclusters):-1:1
    D = double(DD==cluster);
    
    
    % Initial sensitivity and specificity parameter p(j),q(j) for all
    % raters j
    p(1:N) = 0.99999;
    q(1:N) = 0.99999;
    Tprior = (sum(D(:))/length(D(:))); % note dependence on (sub)volume size, final result depends on this prior (which is not an implementation issue but a basic limitation of the EM approach)
    
    avgW = 1;
    W = zeros(1,length(D));
    
    %% EM
    
    for step=1:MAX_ITERATIONS
        
        % E-Step
        % The following vectorized code is equivalent to this loop by MUCH
        % faster on current CPUs
        %     for i = 1:length(D)
        %         W(i) = ((prod(p(D(i,:))) * prod(1 - p(~D(i,:)))) * Tprior) / ...
        %                ((prod(p(D(i,:))) * prod(1 - p(~D(i,:)))) * Tprior + (prod(q(~D(i,:))) * prod(1 - q(D(i,:))))) * (1- Tprior) ;
        %         %NOTE that prod([]) = 1
        %     end
        
        
        P = repmat(p,length(D), 1);
        Q = repmat(q,length(D), 1);
        P_given_D = P .* D; %TODO: use bsxfun instead of repmat?
        P_given_D(P_given_D(:)== 0) = 1; %
        Q_given_D = 1 - Q .* D;
        Q_given_D(Q_given_D(:)== 0) = 1; % alternative: initialize with 1 and set Q_given_D(D) = 1- P(D)
        compP_given_not_D  = 1 - P .* ~D;
        compP_given_not_D(compP_given_not_D(:)== 0) = 1;
        compQ_given_not_D  = Q .* ~D;
        compQ_given_not_D(compQ_given_not_D(:)== 0) = 1;
        
        % W(i) can be interpretated as the prob. of cell i beeing true (i.e. is part of the groundtruth y) for given p(1:N), q(1:N)
        W = (prod(P_given_D') .* prod(compP_given_not_D') * Tprior) ./ ...
            ((prod(P_given_D') .* prod(compP_given_not_D') * Tprior) + (prod(Q_given_D') .* prod(compQ_given_not_D') * (1 - Tprior)));
        
        % Convergence?
        if(abs(avgW - sum(W) / length(W)) < EPSILON)
            break;
        end
        avgW = sum(W) / length(W);
        
        % M-Step
        p = (W * D) / sum(W(:)); % W * D = sum(W(D))
        q = ((1 -  W) * ~D) / sum(1 - W(:));
    end
    
    %% write the result image
    staple_wmap = reshape(W,map_dim);
    save(fullfile(pwd,['staple_wmap_cluster' num2str(cluster) '.mat']),'staple_wmap')
    staple_thmap = staple_wmap.*(staple_wmap>threshold/N);
    save(fullfile(pwd,['staple_thmap_cluster' num2str(cluster) '.mat']),'staple_thmap')
    staple_thmap(find(staple_thmap)) = cluster; % label cluster
    if ndims(staple_thmap) == 2 %#ok<ISMAT>
        all_maps(:,:,cluster) = staple_thmap;
    else
        all_maps(:,:,:,cluster) = staple_thmap;
    end
end
staple_thmap = sum(all_maps,ndims(all_maps));
save(fullfile(pwd,'staple_thmap_labelled.mat'),'staple_thmap')
figure; imagesc(staple_thmap);
title('Binary STAPLE map')