Can we limit the predicted segmentation masks to specific pixels intensities (eg 0, 128, 255)?

[GXcells]

I am running training on a dataset of immunofluorescence images of cells in which the segmentation masks have 3 values for pixel intensities in the masks (0, 128, 255). 0 being background, 128 nuclei, 255 cytoplasm

https://github.com/edwardcao3026/SegBenchmark/blob/main/datasets/datasets.tar

Is it possible to limit the mask prediction to these 3 intensities?
For now I ran up to 50 epochs on 411 image/masks pairs and I kind of gradients with many different gry values instead of only 3

[xiongxyowo]

Hi, you can try to change the output channels of segmentation heads from 1 to 3 to support multi-class segmentation. In this case, 0, 128, 255 correspond to classes 0, 1, 2 respectively. We recommend referring to some U-Net examples (such as Pytorch-UNet) that support both binary and multi-class segmentation to confirm what modifications need to be made to SAM2-UNet.

[GXcells]

Thanks, I'll try to work on it :)

[SijieY]

请问您是怎么解决的呢？

[wangxiang-wx]

请问您是怎么解决的呢？

Hello, we have roughly made the following changes:

1. Modify the number of output channels in SAM2-UNet.py
2. Modify the data loading method in dataset.py (ensure that the values ​​of the training mask are 0, 1, 2, 3)
3. Modify the loss function, use celoss, and add iou loss
   However, we found that the detection effect after multi-class training is relatively poor compared to the results of binary classification training. Do you think the training parameters need to be adjusted? Or the weight of the loss function needs to be improved, etc.? Looking forward to your suggestions