Image augmentation pipeline, applies random flips incorrectly

[Sajtospoga01]

Happens due to the mask being flattened, such the same random transforms don't apply to the masks what apply to an image.