Code not following research paper equations in staple.py file

[ayush1298]

Description

In the staple.py file in the staple folder, there is a function called def _run_expectation(self). In this function, we are computing values of a and b using the equations given in the research paper ' Simultaneous truth and performance level estimation (STAPLE): an algorithm for the validation of image segmentation. ' But the code of equations does not follow the equations in a research paper. The following is the code given in the staple.py file.

Compute a

    p_matrix = p.repeat(N, axis=0)
    p1_matrix = 1 - p_matrix
    m_1 = np.ma.masked_array(p_matrix, D == 0).prod(axis=1)  
    m_2 = np.ma.masked_array(p1_matrix, D == 1).prod(axis=1)  
    a = f_1 * m_1.data * m_2.data

What I Did :

The problem with the above equation is that in m_1, there should be D==1; in the m_2 equation, there should be D==0, as per the equation given in the research paper. So correct equation should be

Compute a

    p_matrix = p.repeat(N, axis=0)
    p1_matrix = 1 - p_matrix
    m_1 = np.ma.masked_array(p_matrix, D == 1).prod(axis=1)  
    m_2 = np.ma.masked_array(p1_matrix, D == 0).prod(axis=1) 
    a = f_1 * m_1.data * m_2.data

I have also attached an image of the original equation, equation No.14, in the above research paper. Similarly, there should be changes done when computing b in the same function(for reference, see equation no. 15 from the research paper) and also in computing p and q (for reference, see equations no. 18 and 19 from the research paper) in _run_maximization(self) function.

[fepegar]

Hi, @ayush1298. I wrote this code a long time ago and am currently lacking a lot of context. Have yo tried this? I took a look at the code and it seems that the flipped numbers in both the expectation and maximization steps make the current implementation work. If this is important to you, feel free to open a pull request and add tests.

[ayush1298]

Hi @fepegar . I was working on this project 7-8 months before, and I have made these changes, but even after this, the algorithm is not converging. I don't know what the possible reasons for this are. Do you have any suggestions for this?

[fepegar]

This algorithm seems to work for me. If you can share some data and a minimal reproducible example, maybe I'll be able to help.