ExtractTensorPatches for Masks in AugmentationSequential

🚀 Feature

Modify kornia.contrib.ExtractTensorPatches to apply the same operation to masks in AugmentationSequential.

Motivation

In remote sensing it’s common to work with imagery which can be relatively large in shape (e.g. 3000x3000+). Therefore during training images are broken into patches. For segmentation tasks it’s needed to apply the same patch extract operation to both images and masks.

Pitch

Modify kornia.contrib.ExtractTensorPatches to apply the same operation to masks in AugmentationSequential.

Alternatives

If there is a better way to chain augmentations and ExtractTensorPatches please provide suggestions.

Additional context

import torch
import kornia.augmentation as K
from kornia.contrib import ExtractTensorPatches

image = torch.randn(2, 3, 128, 128)
mask = torch.randn(2, 1, 128, 128)

augs = K.AugmentationSequential(
    K.ColorJitter(0.1, 0.1, 0.1, 0.1, p=1.0),
    K.RandomAffine(360, p=1.0),
    ExtractTensorPatches(window_size=32, stride=32, padding=0),
    data_keys=["input", "mask"]
)
image, mask = augs(image, mask)

# NotImplementedError: data_key DataKey.MASK is not implemented for ExtractTensorPatches().

@adamjstewart @calebrob6

[shijianjian]

Hi @isaaccorley, please checkout #1347.

Then, the required operation can be implemented as

PatchExtractor = LambdaAugmentation.from_modules(
    kornia.contrib.ExtractTensorPatches(),
    inverse_module=kornia.contrib.CombineTensorPatches(),
    for_data_keys=["input", "mask"],
)

Please comment on the PR if any 😃

[isaaccorley]

I think lambda makes sense here and helps support chaining other general ops in a AugmentationSequential container.