Repository: kornia/kornia

Issue:
No automatic `nearest` interpolation for masks in `RandomElasticTransform`
### Describe the bug

Usually, the data type `DataKey.MASK` uses the `nearest` interpolation method to preserve label index values (e.g. when using `RandomAffine`). However, this does not work for the `RandomElasticTransform` augmentation since the corresponding parameter is named `mode` and not `resample`. Hence the default argument to `AugmentationSequential`
```python
extra_args={DataKey.MASK: dict(resample=Resample.NEAREST, align_corners=True)}
```
does not trigger the `nearest` interpolation mode.

### Reproduction steps

Here is an example reproducing the issue of the current default interpolation for mask input types:
```python
import torch
import kornia.augmentation as K
features = torch.rand(1, 1, 4, 4)
labels = torch.ones(1, 1, 4, 4, dtype=torch.float32) * 10
labels[:, :, :, :2] = 0
labels[:, :, :2, :] = 0
# tensor([[[[ 0.,  0.,  0.,  0.],
#           [ 0.,  0.,  0.,  0.],
#           [ 0.,  0., 10., 10.],
#           [ 0.,  0., 10., 10.]]]])

compose = K.AugmentationSequential(K.RandomElasticTransform(alpha=(10, 10)))
torch.manual_seed(0)
compose(features, labels, data_keys=["input", "mask"])[1]
# tensor([[[[0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
#           [0.0000e+00, 0.0000e+00, 4.4417e-03, 4.3297e-03],
#           [0.0000e+00, 0.0000e+00, 9.8832e+00, 1.0000e+01],
#           [0.0000e+00, 0.0000e+00, 9.8833e+00, 1.0000e+01]]]])
```


### Expected behavior

The default behavior for the elastic transformation should be similar to the other transformations, i.e. `nearest` interpolation for mask types.

### Environment
- PyTorch Version (e.g., 1.0): 1.13
- OS (e.g., Linux): Linux
- How you installed PyTorch (`conda`, `pip`, source): pip
- Build command you used (if compiling from source):
- Python version: 3.10.4
- CUDA/cuDNN version: 11.7
- GPU models and configuration: 3090
- Any other relevant information:


### Additional context

Current workaround for this issue is to manually set the interpolation mode:
```python
compose = K.AugmentationSequential(*transforms, extra_args={DataKey.MASK: dict(resample=Resample.NEAREST, align_corners=True, mode="nearest")})
```

One option would be to rename the `mode` parameter to `resample`, than the default mask handling triggers as well. This would be a breaking change but I see that elastic transformation might get some exciting new functionality in the future anyway (issue #994) so this could be a combined change.


Discussion:
@shijianjian 
@JanSellner I believe this is a mistake for `RandomElasticTransform`. Can you send a PR to deprecate the mode argument and to use resample as the default?

Resolving patch:
diff --git a/kornia/augmentation/_2d/geometric/elastic_transform.py b/kornia/augmentation/_2d/geometric/elastic_transform.py
--- a/kornia/augmentation/_2d/geometric/elastic_transform.py
+++ b/kornia/augmentation/_2d/geometric/elastic_transform.py
@@ -1,9 +1,11 @@
-from typing import Any, Dict, Optional, Tuple
+import warnings
+from typing import Any, Dict, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
 
 from kornia.augmentation._2d.geometric.base import GeometricAugmentationBase2D
+from kornia.constants import Resample
 from kornia.geometry.transform import elastic_transform2d
 
 
@@ -19,7 +21,9 @@ class RandomElasticTransform(GeometricAugmentationBase2D):
         alpha: The scaling factor that controls the intensity of the deformation
           in the y and x directions, respectively.
         align_corners: Interpolation flag used by `grid_sample`.
-        mode: Interpolation mode used by `grid_sample`. Either 'bilinear' or 'nearest'.
+        resample: Interpolation mode used by `grid_sample`. Either 'nearest' (0) or 'bilinear' (1).
+        mode: Deprecated: Interpolation mode used by `grid_sample`. Either 'bilinear' or 'nearest'.
+          Please use the `resample` argument instead.
         padding_mode: The padding used by ```grid_sample```. Either 'zeros', 'border' or 'refection'.
         same_on_batch: apply the same transformation across the batch.
         p: probability of applying the transformation.
@@ -49,7 +53,8 @@ def __init__(
         sigma: Tuple[float, float] = (32.0, 32.0),
         alpha: Tuple[float, float] = (1.0, 1.0),
         align_corners: bool = False,
-        mode: str = "bilinear",
+        resample: Union[str, int, Resample] = Resample.BILINEAR.name,
+        mode: Optional[str] = None,
         padding_mode: str = "zeros",
         same_on_batch: bool = False,
         p: float = 0.5,
@@ -59,12 +64,21 @@ def __init__(
         super().__init__(
             p=p, return_transform=return_transform, same_on_batch=same_on_batch, p_batch=1.0, keepdim=keepdim
         )
+        if mode is not None:
+            resample = mode
+            if not torch.jit.is_scripting():
+                warnings.warn(
+                    "`mode` argument has been deprecated. Please use the `resample` argument instead",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+
         self.flags = dict(
             kernel_size=kernel_size,
             sigma=sigma,
             alpha=alpha,
             align_corners=align_corners,
-            mode=mode,
+            resample=Resample.get(resample),
             padding_mode=padding_mode,
         )
 
@@ -90,6 +104,6 @@ def apply_transform(
             flags["sigma"],
             flags["alpha"],
             flags["align_corners"],
-            flags["mode"],
+            flags["resample"].name.lower(),
             flags["padding_mode"],
         )


Test patch:
diff --git a/test/augmentation/test_augmentation.py b/test/augmentation/test_augmentation.py
--- a/test/augmentation/test_augmentation.py
+++ b/test/augmentation/test_augmentation.py
@@ -3580,6 +3580,24 @@ def test_same_on_batch(self, device, dtype):
         res = f(input)
         assert (res[0] == res[1]).all()
 
+    def test_mask_transform(self, device, dtype):
+        torch.manual_seed(0)
+        features = torch.rand(1, 1, 4, 4, dtype=dtype, device=device)
+        labels = torch.ones(1, 1, 4, 4, dtype=dtype, device=device) * 10
+        labels[:, :, :, :2] = 0
+        labels[:, :, :2, :] = 0
+
+        compose = AugmentationSequential(RandomElasticTransform(alpha=(10, 10)))
+
+        # Use an example which would produce invalid label values
+        torch.manual_seed(0)
+        labels_transformed = compose(features, labels, data_keys=["input", "input"])[1]
+        assert len(labels_transformed.unique()) > 2
+
+        # The transformed values are fine if we use mask input type
+        labels_transformed = compose(features, labels, data_keys=["input", "mask"])[1]
+        assert torch.all(labels_transformed.unique() == torch.tensor([0, 10], dtype=dtype, device=device))
+
 
 class TestRandomThinPlateSpline:
     def test_smoke(self, device, dtype):