Thank you for your response as always!
I believe this post has an implementation of your predicted class indices suggestion: How to save multi-class segmentation prediction as image? - #2 by eqy.
I believe the idea is correct but I am having trouble getting it to run. If you could take a look at that then that would be really amazing! Thank you very much.