HorizontalFlip and VerticalFlip inconsistent with multilabel masks

🐛 Bug

Given a pair of image and its corresponding mask, the generated output for the augmented mask through augmentation is not the same as the image.(it is inconsistent) when HorizontalFlip() and VerticalFlip() are included in the augmentations.

To Reproduce

The following snippet is a small dataloader that i usually use. Can’t share code.

    self.aug = Compose([
                        RandomBrightnessContrast(),
                        HueSaturationValue(),
                        RandomGamma(),
                        GaussNoise(),
                        GaussianBlur(),
                        # HorizontalFlip(),
                        # VerticalFlip(),
                        ])
def __getitem__(self, patient_id):
    image_path = os.path.join(self.df.iloc[patient_id, 0])
    z = np.load(image_path)
    image = z['patch']
    gt_data = z['patch_gt']
    # print("Pre : ", image.shape, gt_data.shape)
    gt_data = gt_data.swapaxes(0, 2)
    # print("Pre swapped: ", image.shape, gt_data.shape)
    # gt_data = gt_data[:4, :, :]
    if not self.valid:
        augmented = self.aug(image=image, mask=gt_data)
        image = augmented['image']
        gt_data = augmented['mask']
    image = (image/255).astype(np.float32)
    # print("Post Augment :", image.shape, gt_data.shape)
    image = image.swapaxes(0, 2)
    gt_data = gt_data.swapaxes(0, 2)
    # print("Post Augment Swapped: ", image.shape, gt_data.shape)
    image = torch.FloatTensor(image)
    gt_data = torch.FloatTensor(gt_data)
    #mask.shape = (5, 256, 256)
    #image.shape = (256, 256, 3)
    return image, gt_data

Expected behavior

The augmentation over the images for horizontal and vertical flip should be working fine for both the mask and the image, but for some reason, there are errors in mask augmentations during horizontal and vertical flips.

Image shape : 256, 256, 3 Mask Shape : 5, 256, 256

Environment

• Albumentations version : 0.43.
• Python version : 3.6
• OS (e.g., Linux): Linux
• How you installed albumentations (conda, pip, source): pip
• Any other relevant information:

Additional context

[Dipet]

As @BloodAxe said, please, give a minimal example to reproduce the problem.

[Dipet]

Problem with swapaxes use np.transpose(gt_data, [1, 2, 0])