Setting particular pixel to specified color in the output of segmentation map

unclear, how you obtain this, and what you get there, exactly, please explain !

shouldn’t there be binary, per-class maps ?
imo, trying to parse an already ‘rendered’ result mapping (like your 2nd image) is a bad idea.