Thanks, this is what I initially wanted! However, it turns out the classes overlap after all  Is there another (standard?) way of merging the masks back into one image, I guess giving preference to masks with a higher confidence score?