I am trying to segment 5 classes in Fundus images.I use a Unet which has an output of 5 channels and have created a custom rgb_to_onehot function.I want to augment the images wi