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