Documenting RandomTransplantation for cut and paste augmentation

[robmarkcole]

Cut and paste augmentation is particularly useful for ship datasets, where ships can be placed in other areas of the ocean.
This requires using from kornia.augmentation.container import AugmentationSequential as the torchgeo implementation adds a channel dim to the mask which is not allowed.

Note kornia.__version__ == 0.7.2

import matplotlib.pyplot as plt
import kornia.augmentation as ka
from kornia.augmentation.container import AugmentationSequential
import torch


datamodule.setup(stage="fit") # not shown, segmentation dataset with image, mask
train_dataloader = datamodule.train_dataloader()
batch = next(iter(train_dataloader)) # a dict with keys image & mask, nothing else!

# Note the mask musn't have a channel dim for this aug
# batch["image"].shape, batch["mask"].shape == (torch.Size([10, 3, 64, 64]), torch.Size([10, 64, 64]))

aug = ka.RandomTransplantation(p=1., excluded_labels=[0]) # 0 is background
image_out, mask_out = aug(batch['image'], batch['mask'])

for i in range(10):
    fig, ax = plt.subplots(1, 4, figsize=(10, 5))
    mask_before = batch["mask"][i].squeeze().numpy()
    image_before = batch["image"][i].squeeze().permute(1, 2, 0).numpy()
    image_before = image_before / image_before.max()

    mask_after = mask_out[i].squeeze().numpy()
    image_after = image_out[i].squeeze().permute(1, 2, 0).numpy()
    image_after = image_after / image_after.max()

    ax[0].imshow(image_before)
    ax[1].imshow(mask_before, cmap="gray")
    ax[2].imshow(image_after)
    ax[3].imshow(mask_after, cmap="gray")
    # ticks of
    for a in ax:
        a.axis("off")
    ax[0].set_title("Image Before")
    ax[1].set_title("Mask Before")
    ax[2].set_title("Image After")
    ax[3].set_title("Mask After")
    plt.show()

And some results - this will work well enough for open water, but can place ships in odd locations.

And to use in training:

        self.train_aug = AugmentationSequential(
            Ka.Normalize(mean=self.mean, std=self.std, keepdim=True),
            Ka.RandomTransplantation(p=1.0, excluded_labels=[0]),
            data_keys=None,
        )