Applying the same random transformation on the images and on the masks for segmentation task

🐛 Bug

I’m working on a segmentation task and I’d like to apply the same random transformation on the images and on the masks. When I inspect the datasets, the images and masks don’t match. In the following code I just used HorizontalFlip for simplicity. Thank you very much in advance!

To Reproduce

from PIL import Image
from torchvision.transforms import Resize
import matplotlib.pyplot as plt
from albumentations.augmentations.transforms import HorizontalFlip

class MyDataset(Dataset):

    def __init__(self, list_images_input, list_images_target, transform=None):
        super(MyDataset, self).__init__()
        self.images_input = list_images_input
        self.images_target = list_images_target
        self.resize = Resize(224)
        self.transform = transform

    def __len__(self):
        return len(self.images_input)

    def __getitem__(self, index):
        img_path = str(self.images_input[index])
        img = Image.open(img_path).convert('RGB')
        seg_path = self.images_target[index]
        seg = Image.open(seg_path)
        seg = seg.convert('L')
        img = self.resize(img)
        seg = self.resize(seg)
        img = np.array(img)
        seg = np.array(seg)

        if self.transform:
            transformed = self.transform(image=img, mask=seg)
            img = transformed['image']
            seg = transformed['mask']
        seg = (seg > 0)
        return img, seg

def visualize_label_onto_image(img, label, alpha):
    plt.imshow(img)
    plt.imshow(label.squeeze(), cmap='Greys', alpha=alpha)
    plt.axis('off')
    plt.show()

def read_list_images_path_from_txt(txt_path):
  with open(txt_path) as f:
    content = f.readlines()
    content = [x.strip() for x in content]
    return content
  f.close()

input_path='input.txt'
target_path='target.txt'

imgs_training_input_split=read_list_images_path_from_txt(input_path)
imgs_training_target_split=read_list_images_path_from_txt(target_path)

affine_transform = HorizontalFlip()
training_dataset_plain = MyDataset(imgs_training_input_split, imgs_training_target_split)
training_dataset_affine = MyDataset(imgs_training_input_split, imgs_training_target_split, transform=affine_transform)

for i in range(10):
    visualize_label_onto_image(training_dataset_plain[i][0], training_dataset_plain[i][1], 0.3)
    visualize_label_onto_image(training_dataset_affine[i][0], training_dataset_affine[i][1], 0.3)

Environment

• Albumentations version 0.5.2
• Python version 3.7.5

[Dipet]

visualize_label_onto_image(training_dataset_plain[i][0], training_dataset_plain[i][1], 0.3)

There you get different image and mask, because it is 2 different calls of transforms. Change to:

res = training_dataset_plain[i]
visualize_label_onto_image(res[0], res][1], 0.3)

[ematao]

Thank you so much!