LIVIAETS / boundary-loss

Official code for "Boundary loss for highly unbalanced segmentation", runner-up for best paper award at MIDL 2019. Extended version in MedIA, volume 67, January 2021.
https://doi.org/10.1016/j.media.2020.101851
MIT License
654 stars 98 forks source link

How to apply to 'instance' segment in neuron segmentation? #22

Closed Keep-Passion closed 4 years ago

Keep-Passion commented 4 years ago

Hi, @HKervadec thanks for your sharing I want to apply your code to EM datasets in ISBI 2012 (The website is here) I find that main difference between EM datasets and WMH datasets is the Intensive degree of each segment. The code in your project is to calculate boundary loss for each class

def one_hot2dist(seg: np.ndarray) -> np.ndarray:
    assert one_hot(torch.Tensor(seg), axis=0)
    C: int = len(seg)

    res = np.zeros_like(seg)
    for c in range(C):
        posmask = seg[c].astype(np.bool)

        if posmask.any():
            negmask = ~posmask
            res[c] = distance(negmask) * negmask - (distance(posmask) - 1) * posmask
    return res

I think distance(distance_transform_edt) will get 1 for all membrane (black region) in EM datasets, so is there any chance to rewrite this function to suit for cell segmentation? That is C in your script is number of neurons instead of class.

HKervadec commented 4 years ago

Hey,

I have been thinking about it for a few days ; this would be an interesting line of research.

That is C in your script is number of neurons instead of class.

This could be done, yes ; replacing C with the number of neurons. You would then end up with one distance map per neuron.

One problem I would see, is that neurons (in the instance segmentation case) are supposed to be interchangeable (the number of the instance doesn't really matter, what you really want is that is it different than the neighbor). This is in contrast with class segmentation ; where there is only one true answer.

Perhaps it would be a problem, perhaps it wouldn't be.

Another direction (I'm definitely not knowledgeable about current instance segmentation methods) could be to use the boundary loss to help predict a "boundary mask": bothering only to find the exact boundaries of the neurons in general (binary problem), and having the instance identification coming from somewhere else (another - more blurry - mask, bounding box for instances, ..). In this way, the boundary loss would be some sort of refinement (similar to a CRF loss/post-processing).

Hoel