Repository: kornia/kornia
Pull request: Fix params computation for `LongestMaxSize` and `SmallestMaxSize`
Fixes #2123


#### Type of change
- [x] 🧪 Tests Cases
- [x] 🐞 Bug fix (non-breaking change which fixes an issue)

Problem statement:
Weird behavior of LongestMaxSize
### Describe the bug

Hello me again,

I might be doing something wrong with the way I use kornia augmentations, please let me know if it is the case.

I was expecting `LongestMaxSize` in kornia to perform similarily as the albumentation implementation. Meaning that I can throw any images with different shapes to the the transformation function and get an image with different shapes but similar ratios. The largest size being equal to the value given to `LongestMaxSize`.

See bellow a small code sample that disturbs me.

### Reproduction steps

```bash
import kornia.augmentation as K
a = torch.ones((512, 256))
b = torch.ones((512, 756))

print("first try")
transfo = K.LongestMaxSize(max_size=256, p=1.)

print(transfo(a).shape)
print(transfo(b).shape)

print("second try")

a = torch.ones((512, 256))
b = torch.ones((512, 756))

transfo = K.LongestMaxSize(max_size=256, p=1.)
print(transfo(b).shape)
print(transfo(a).shape)

Outputs:
first try
torch.Size([1, 1, 256, 128])
torch.Size([1, 1, 256, 128])
second try
torch.Size([1, 1, 173, 256])
torch.Size([1, 1, 173, 256])
```


### Expected behavior

I would expect to have the same values for the transformations no matter the order of the elements.

ie `transfo(a).shape == torch.Size([1, 1, 256, 128])` and `transfo(b).shape ==torch.Size([1, 1, 173, 256])`

Am I missing something here ?

### Environment

```shell
kornia='0.6.9'
torch='1.12.1+cu113'
```


### Additional context

_No response_

Hints from issue discussion:
Seems like Resize behave similarly when size is an integer.
@johnnv1 @shijianjian Is this the expected behavior ?
@twsl Can you take a look?
Unfortunately I don't have much time right now, but a quick peek at the source code makes me wonder if the resize generator in the resize augmentation does include any checks to recalculate the parameters
@johnnv1 Can you help fix it if you got some time?
as @twsl commented, seems for this case we need to recompute the params -- i'm not sure of the best way to try avoid generating the params in each call

Commit message: Fix the side resize for `SmallestMaxSize`

Patch to review:
diff --git a/kornia/augmentation/_2d/geometric/resize.py b/kornia/augmentation/_2d/geometric/resize.py
index eb6bb059cb..efef0526a6 100644
--- a/kornia/augmentation/_2d/geometric/resize.py
+++ b/kornia/augmentation/_2d/geometric/resize.py
@@ -114,6 +114,12 @@ def __init__(
             p=p,
         )
 
+    def apply_transform(
+        self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
+    ) -> Tensor:
+        self._param_generator = rg.ResizeGenerator(resize_to=self.flags['size'], side=self.flags['side'])
+        return super().apply_transform(input, params, flags, transform)
+
 
 class SmallestMaxSize(Resize):
     """Rescale an image so that minimum side is equal to max_size, keeping the aspect ratio of the initial image.
@@ -139,3 +145,9 @@ def __init__(
             align_corners=align_corners,
             p=p,
         )
+
+    def apply_transform(
+        self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
+    ) -> Tensor:
+        self._param_generator = rg.ResizeGenerator(resize_to=self.flags['size'], side=self.flags['side'])
+        return super().apply_transform(input, params, flags, transform)
diff --git a/test/augmentation/test_augmentation.py b/test/augmentation/test_augmentation.py
index ab8119896b..136fbf1f7a 100644
--- a/test/augmentation/test_augmentation.py
+++ b/test/augmentation/test_augmentation.py
@@ -3727,18 +3727,30 @@ def test_smoke(self, device, dtype):
 
 class TestSmallestMaxSize:
     def test_smoke(self, device, dtype):
-        img = torch.rand(1, 1, 4, 6, device=device, dtype=dtype)
+        img_A = torch.rand(1, 1, 4, 6, device=device, dtype=dtype)
+        img_B = torch.rand(1, 1, 9, 6, device=device, dtype=dtype)
         aug = SmallestMaxSize(max_size=2)
-        out = aug(img)
-        assert out.shape == (1, 1, 2, 3)
+
+        assert aug(img_A).shape == (1, 1, 2, 3)
+        assert aug(img_B).shape == (1, 1, 3, 2)
+
+        aug = SmallestMaxSize(max_size=2)
+        assert aug(img_B).shape == (1, 1, 3, 2)
+        assert aug(img_A).shape == (1, 1, 2, 3)
 
 
 class TestLongestMaxSize:
     def test_smoke(self, device, dtype):
-        img = torch.rand(1, 1, 4, 6, device=device, dtype=dtype)
+        img_A = torch.rand(1, 1, 4, 6, device=device, dtype=dtype)
+        img_B = torch.rand(1, 1, 8, 6, device=device, dtype=dtype)
         aug = LongestMaxSize(max_size=3)
-        out = aug(img)
-        assert out.shape == (1, 1, 2, 3)
+
+        assert aug(img_A).shape == (1, 1, 2, 3)
+        assert aug(img_B).shape == (1, 1, 3, 2)
+
+        aug = LongestMaxSize(max_size=3)
+        assert aug(img_B).shape == (1, 1, 3, 2)
+        assert aug(img_A).shape == (1, 1, 2, 3)
 
 
 class TestRandomPosterize:


Code review comments:
diff_hunk: @@ -114,6 +114,12 @@ def __init__(
             p=p,
         )
 
+    def apply_transform(
+        self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
+    ) -> Tensor:
+        self._param_generator = rg.ResizeGenerator(resize_to=self.flags['size'], side=self.flags['side'])
path: kornia/augmentation/_2d/geometric/resize.py
text: @author:
```suggestion
        self._param_generator.output_size = self.flags['size']
```

I don't know what makes more sense here, creating another ParamGenerator for these cases, or update of the resize size as this suggestion, or some other thing

for default, the `ResizeGenerator` set the `output_size` on the init (https://github.com/kornia/kornia/blob/master/kornia/augmentation/random_generator/_2d/resize.py#L34), and then when expands it from int to a shape (https://github.com/kornia/kornia/blob/master/kornia/augmentation/random_generator/_2d/resize.py#L69)

as I understand it, for these MaxSize cases we need to update this size on each call, because the `output_size` depends on the input size (which side is bigger)

@user1

diff_hunk: @@ -114,6 +114,12 @@ def __init__(
             p=p,
         )
 
+    def apply_transform(
+        self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
+    ) -> Tensor:
+        self._param_generator = rg.ResizeGenerator(resize_to=self.flags['size'], side=self.flags['side'])
+        return super().apply_transform(input, params, flags, transform)
+
path: kornia/augmentation/_2d/geometric/resize.py
text: @author:
Add some check here before recompute the params? 

@user1:
out of curiosity, what kind of checks are thinking of ?

@author:
Try cache at least the last one, like check if the last aspectio ratio and long/small side is equals the actual

@user1:
Isn't it just as complex to check that you don't need to update the parameters than simply updating every time ?

@author:
Depends on the case, that's the reason I don't know the best case to assume here

If we always need to check and recompute, just recompute will be the best case -- but if we never needs to recompute we can save some computation -- here https://github.com/kornia/kornia/blob/598444893212899f6f3678f45a10b5315ef0f5f9/kornia/augmentation/random_generator/_2d/resize.py#L46



@user2:
Ideally. We shall not call the random generator again. There must be somewhere in the code that weirdly remembers the parameters. 

@author:
oh, theoretically this makes a lot more sense  😅 

@user3:
then, should we wait to find the issue to merge; or we proceed and fix later ?

@author:
I will investigate this further in the next few days