ReplayCompose gives Keyerror when augmenting mask.

[toshi2k2]

The following error pops up if trying to use ReplayCompose with a mask.

Error -

KeyError: Caught KeyError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/home/ubuntu/my_git/lama/inpenv/lib/python3.6/site-packages/torch/utils/data/_utils/worker.py", line 287, in _worker_loop
    data = fetcher.fetch(index)
  File "/home/ubuntu/my_git/lama/inpenv/lib/python3.6/site-packages/torch/utils/data/_utils/fetch.py", line 49, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/ubuntu/my_git/lama/inpenv/lib/python3.6/site-packages/torch/utils/data/_utils/fetch.py", line 49, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/ubuntu/my_git/lama/saicinpainting/training/data/datasets.py", line 164, in __getitem__
    houghed_wireframe_data = A.ReplayCompose.replay(data['replay'], houghed_wireframe=houghed_wireframe)
  File "/home/ubuntu/my_git/lama/inpenv/lib/python3.6/site-packages/albumentations/core/composition.py", line 331, in replay
    return augs(force_apply=True, **kwargs)
  File "/home/ubuntu/my_git/lama/inpenv/lib/python3.6/site-packages/albumentations/core/composition.py", line 321, in __call__
    result = super(ReplayCompose, self).__call__(force_apply=force_apply, **kwargs)
  File "/home/ubuntu/my_git/lama/inpenv/lib/python3.6/site-packages/albumentations/core/composition.py", line 182, in __call__
    data = t(force_apply=force_apply, **data)
  File "/home/ubuntu/my_git/lama/inpenv/lib/python3.6/site-packages/albumentations/core/transforms_interface.py", line 68, in __call__
    return self.apply_with_params(self.params, **kwargs)
  File "/home/ubuntu/my_git/lama/inpenv/lib/python3.6/site-packages/albumentations/core/transforms_interface.py", line 96, in apply_with_params
    params = self.update_params(params, **kwargs)
  File "/home/ubuntu/my_git/lama/inpenv/lib/python3.6/site-packages/albumentations/imgaug/transforms.py", line 38, in update_params
    params = super(BasicIAATransform, self).update_params(params, **kwargs)
  File "/home/ubuntu/my_git/lama/inpenv/lib/python3.6/site-packages/albumentations/core/transforms_interface.py", line 146, in update_params
    params.update({"cols": kwargs["image"].shape[1], "rows": kwargs["image"].shape[0]})
KeyError: 'image'

Code used -

transform = A.ReplayCompose([
            IAAPerspective2(scale=(0.0, 0.06)),
            IAAAffine2(scale=(0.7, 1.3),
                       rotate=(-40, 40),
                       shear=(-0.1, 0.1)),
            A.PadIfNeeded(min_height=out_size, min_width=out_size),
            A.OpticalDistortion(),
            A.RandomCrop(height=out_size, width=out_size),
            A.HorizontalFlip(),
            A.CLAHE(),
            A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2),
            A.HueSaturationValue(hue_shift_limit=5, sat_shift_limit=30, val_shift_limit=5),
            A.ToFloat()
        ],
        additional_targets={'wireframe':'mask', 'hough_wireframe':'mask'}
        )
....
data = self.transform(image=img, wireframe=wireframe) 
img, wireframe = data['image'], data['wireframe']
...
houghed_wireframe_data = A.ReplayCompose.replay(data['replay'], houghed_wireframe=houghed_wireframe)

[Dipet]

Now, when you call A.ReplayCompose.replay you must provide image to

[toshi2k2]

Do I need to provide image again? Will that change the augmentations from the first function call?

[Dipet]

Yes, now you also have to provide an image too. It looks like we need to change this logic for ReplayCompose.

Try this:

houghed_wireframe_data = A.ReplayCompose.replay(data['replay'], image=np.emty_like(img), houghed_wireframe=houghed_wireframe)