I'm working on a semantic segmentation project, and want to add some domain knowledge to the system. I want to ensure that for segmentation, there can only be one group of pixels that are predicted as a certain class. An example would be segmentation of the face. When we want to segment different parts of the face, assuming regular faces, we have the following domain knowledge: one group of pixels can be predicted as the mouth; one group can be predicted as the nose; two groups of pixels can be predicted as the eyes.
A possible solution I currently think of is as follows: Suppose we find two groups of pixels, both groups having pixels with the highest probability for being the nose. Then, we would take the average score for each pixel group to obtain two average scores. Then we compute the probabilities again and select the group with the highest probability to be the nose. The group of pixels having the lower probability will then be predicted to be the class having the second best probability. Doing it in this way, we apply some post-processing (look at scores and determine what is our final prediction, using some logic). However, it does not seem to be learned by the neural network (or maybe it is learned implicitly, by looking at the final predicted class obtained from the postprocessing and updating the weights accordingly).
Are there any solutions, or could you point me to any papers which describe this problem and add domain knowledge to the neural network such that it incorporates this information and becomes better at predicting?