search

fepegar / staple

Python implementation of the STAPLE segmentation algorithm
MIT License
23 stars 5 forks source link

Code not following research paper equations in staple.py file #50

Open ayush1298 opened 1 year ago

ayush1298 commented 1 year ago

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. image

fepegar commented 7 months ago

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 commented 7 months ago

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 commented 7 months ago

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.