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 :)