Repository: kornia/kornia

Issue:
Incorrect shape while using CombineTensorPatches
### Describe the bug

Currently `CombineTensorPatches` does not accurately invert the results of `ExtractTensorPatches`. This can be observed by running the following script.

### Reproduction steps

```bash
import torch

from kornia.contrib import CombineTensorPatches, ExtractTensorPatches

# Original size
h, w = 16, 16
# Window size
win = 2

image = torch.randn(2, 3, h, w)
print(image.shape)
tiler = ExtractTensorPatches(window_size=win, stride=win, padding=0)
merger = CombineTensorPatches(window_size=win, unpadding=0)
image_tiles = tiler(image)
new_image = merger(image_tiles)
print(new_image.shape)
assert (image == new_image).all()
```

Instead of returning the original shape (2,3,16,16), `CombineTensorPatches` will return (32, 3, 4, 4)


### Expected behavior

image and new_image should have the same shape - (2,  3, 16, 16)

### Environment

```shell
I used the devcontainer


PyTorch version: 1.10.2+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 11 (bullseye) (x86_64)
GCC version: (Debian 10.2.1-6) 10.2.1 20210110
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.8.12 (default, Dec 21 2021, 10:45:09)  [GCC 10.2.1 20210110] (64-bit runtime)
Python platform: Linux-5.4.0-96-generic-x86_64-with-glibc2.2.5
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy==0.931
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.22.2
[pip3] pytest-mypy==0.8.1
[pip3] torch==1.10.2
[pip3] torchvision==0.11.3
[conda] Could not collect


```
### Additional context

This error is caused by `CombineTensorPatches` not taking into account the original size of the image. It can be solved by setting `window_size = (original_size[0] // window_size[0], original_size[1] // window_size[1])`. I've opened a PR to implement this fix.


Discussion:

Resolving patch:
diff --git a/kornia/contrib/extract_patches.py b/kornia/contrib/extract_patches.py
--- a/kornia/contrib/extract_patches.py
+++ b/kornia/contrib/extract_patches.py
@@ -90,14 +90,14 @@ class CombineTensorPatches(nn.Module):
         defined in the function signature.
         left-right and top-bottom order.
 
-    * :attr:`window_size` is the size of the sliding window and controls the
-      shape of the output tensor and defines the shape of the output patch.
-    * :attr:`stride` controls the stride to apply to the sliding window and
-      regulates the overlapping between the extracted patches.
-    * :attr:`padding` controls the amount of implicit zeros-paddings on both
-      sizes at each dimension.
-
-    The parameters :attr:`window_size`, :attr:`stride` and :attr:`padding` can
+    * :attr:`original_size` is the size of the original image prior to
+      extracting tensor patches and defines the shape of the output patch.
+    * :attr:`window_size` is the size of the sliding window used while
+      extracting tensor patches.
+    * :attr:`unpadding` is the amount of padding to be removed. This value
+      must be the same as padding used while extracting tensor patches.
+
+    The parameters :attr:`original_size`, :attr:`window_size`, and :attr:`unpadding` can
     be either:
 
         - a single ``int`` -- in which case the same value is used for the
@@ -107,7 +107,8 @@ class CombineTensorPatches(nn.Module):
 
     Args:
         patches: patched tensor.
-        window_size: the size of the sliding window and the output patch size.
+        original_size: the size of the original tensor and the output patch size.
+        window_size: the size of the sliding window used.
         unpadding: remove the padding added to both side of the input.
 
     Shape:
@@ -116,29 +117,38 @@ class CombineTensorPatches(nn.Module):
 
     Example:
         >>> out = extract_tensor_patches(torch.arange(16).view(1, 1, 4, 4), window_size=(2, 2), stride=(2, 2))
-        >>> combine_tensor_patches(out, window_size=(2, 2), stride=(2, 2))
+        >>> combine_tensor_patches(out, original_size=(4, 4), window_size=(2, 2), stride=(2, 2))
         tensor([[[[ 0,  1,  2,  3],
                   [ 4,  5,  6,  7],
                   [ 8,  9, 10, 11],
                   [12, 13, 14, 15]]]])
+
+    .. note::
+        This function is supposed to be used in conjunction with :class:`ExtractTensorPatches`.
+
     """
 
     def __init__(
         self,
+        original_size: Union[int, Tuple[int, int]],
         window_size: Union[int, Tuple[int, int]],
         unpadding: Union[int, Tuple[int, int]] = 0,
     ) -> None:
         super().__init__()
+        self.original_size: Tuple[int, int] = _pair(original_size)
         self.window_size: Tuple[int, int] = _pair(window_size)
         pad: Tuple[int, int] = _pair(unpadding)
         self.unpadding: Tuple[int, int, int, int] = (pad[0], pad[0], pad[1], pad[1])
 
-    def forward(self, input: torch.Tensor) -> torch.Tensor:  # type: ignore
-        return combine_tensor_patches(input, self.window_size, stride=self.window_size, unpadding=self.unpadding)
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return combine_tensor_patches(
+            input, self.original_size, self.window_size, stride=self.window_size, unpadding=self.unpadding
+        )
 
 
 def combine_tensor_patches(
     patches: torch.Tensor,
+    original_size: Tuple[int, int] = (16, 16),
     window_size: Tuple[int, int] = (4, 4),
     stride: Tuple[int, int] = (4, 4),
     unpadding: Optional[Tuple[int, int, int, int]] = None,
@@ -147,7 +157,8 @@ def combine_tensor_patches(
 
     Args:
         patches: patched tensor with shape :math:`(B, N, C, H_{out}, W_{out})`.
-        window_size: the size of the sliding window and the output patch size.
+        original_size: the size of the original tensor and the output patch size.
+        window_size: the size of the sliding window used while extracting patches.
         stride: stride of the sliding window.
         unpadding: remove the padding added to both side of the input.
 
@@ -156,28 +167,42 @@ def combine_tensor_patches(
 
     Example:
         >>> out = extract_tensor_patches(torch.arange(16).view(1, 1, 4, 4), window_size=(2, 2), stride=(2, 2))
-        >>> combine_tensor_patches(out, window_size=(2, 2), stride=(2, 2))
+        >>> combine_tensor_patches(out, original_size=(4, 4), window_size=(2, 2), stride=(2, 2))
         tensor([[[[ 0,  1,  2,  3],
                   [ 4,  5,  6,  7],
                   [ 8,  9, 10, 11],
                   [12, 13, 14, 15]]]])
+
+    .. note::
+        This function is supposed to be used in conjunction with :func:`extract_tensor_patches`.
     """
     if stride[0] != window_size[0] or stride[1] != window_size[1]:
         raise NotImplementedError(
             f"Only stride == window_size is supported. Got {stride} and {window_size}."
             "Please feel free to drop a PR to Kornia Github."
         )
+
+    if original_size[0] % 2 != 0 or original_size[1] % 2 != 0:
+        raise NotImplementedError(f"Original image size must be divisible by 2. Got {original_size}")
+
     if unpadding is not None:
+        hpad_check = (original_size[0] + unpadding[0] + unpadding[1]) % window_size[0] == 0
+        wpad_check = (original_size[1] + unpadding[2] + unpadding[3]) % window_size[1] == 0
+
+        if not hpad_check or not wpad_check:
+            raise NotImplementedError("Insufficient padding")
+
         window_size = (
-            window_size[0] + (unpadding[0] + unpadding[1]) // window_size[0],
-            window_size[1] + (unpadding[2] + unpadding[3]) // window_size[1]
+            (original_size[0] + (unpadding[0] + unpadding[1])) // window_size[0],
+            (original_size[1] + (unpadding[2] + unpadding[3])) // window_size[1],
         )
+
     patches_tensor = patches.view(-1, window_size[0], window_size[1], *patches.shape[-3:])
     restored_tensor = torch.cat(torch.chunk(patches_tensor, window_size[0], dim=1), -2).squeeze(1)
     restored_tensor = torch.cat(torch.chunk(restored_tensor, window_size[1], dim=1), -1).squeeze(1)
 
     if unpadding is not None:
-        restored_tensor = torch.nn.functional.pad(restored_tensor, [-i for i in unpadding])
+        restored_tensor = F.pad(restored_tensor, [-i for i in unpadding])
     return restored_tensor
 
 
@@ -225,7 +250,7 @@ def extract_tensor_patches(
     if not torch.is_tensor(input):
         raise TypeError(f"Input input type is not a torch.Tensor. Got {type(input)}")
 
-    if not len(input.shape) == 4:
+    if len(input.shape) != 4:
         raise ValueError(f"Invalid input shape, we expect BxCxHxW. Got: {input.shape}")
 
     if padding:


Test patch:
diff --git a/test/test_contrib.py b/test/test_contrib.py
--- a/test/test_contrib.py
+++ b/test/test_contrib.py
@@ -266,7 +266,7 @@ def test_gradcheck(self, device):
 class TestCombineTensorPatches:
     def test_smoke(self, device, dtype):
         img = torch.arange(16, device=device, dtype=dtype).view(1, 1, 4, 4)
-        m = kornia.contrib.CombineTensorPatches((2, 2))
+        m = kornia.contrib.CombineTensorPatches((4, 4), (2, 2))
         patches = kornia.contrib.extract_tensor_patches(img, window_size=(2, 2), stride=(2, 2))
         assert m(patches).shape == (1, 1, 4, 4)
         assert (img == m(patches)).all()
@@ -276,21 +276,51 @@ def test_error(self, device, dtype):
             torch.arange(16, device=device, dtype=dtype).view(1, 1, 4, 4), window_size=(2, 2), stride=(2, 2), padding=1
         )
         with pytest.raises(NotImplementedError):
-            kornia.contrib.combine_tensor_patches(patches, window_size=(2, 2), stride=(3, 2))
+            kornia.contrib.combine_tensor_patches(patches, original_size=(4, 4), window_size=(2, 2), stride=(3, 2))
+
+    def test_rect_odd_dim(self, device, dtype):
+        patches = kornia.contrib.extract_tensor_patches(
+            torch.arange(12, device=device, dtype=dtype).view(1, 1, 4, 3), window_size=(2, 2), stride=(2, 2), padding=1
+        )
+        with pytest.raises(NotImplementedError):
+            kornia.contrib.combine_tensor_patches(patches, original_size=(4, 3), window_size=(2, 2), stride=(2, 2))
+
+    def test_pad_error(self, device, dtype):
+        patches = kornia.contrib.extract_tensor_patches(
+            torch.arange(64, device=device, dtype=dtype).view(1, 1, 8, 8), window_size=(4, 4), stride=(4, 4), padding=1
+        )
+        with pytest.raises(NotImplementedError):
+            kornia.contrib.combine_tensor_patches(
+                patches, original_size=(8, 8), window_size=(4, 4), stride=(4, 4), unpadding=(1, 1, 1, 1)
+            )
+
+    def test_rectangle_array(self, device, dtype):
+        img = torch.arange(24, device=device, dtype=dtype).view(1, 1, 4, 6)
+        patches = kornia.contrib.extract_tensor_patches(img, window_size=(2, 2), stride=(2, 2), padding=1)
+        m = kornia.contrib.CombineTensorPatches((4, 6), (2, 2), unpadding=1)
+        assert m(patches).shape == (1, 1, 4, 6)
+        assert (img == m(patches)).all()
 
     def test_padding1(self, device, dtype):
         img = torch.arange(16, device=device, dtype=dtype).view(1, 1, 4, 4)
         patches = kornia.contrib.extract_tensor_patches(img, window_size=(2, 2), stride=(2, 2), padding=1)
-        m = kornia.contrib.CombineTensorPatches((2, 2), unpadding=1)
+        m = kornia.contrib.CombineTensorPatches((4, 4), (2, 2), unpadding=1)
         assert m(patches).shape == (1, 1, 4, 4)
         assert (img == m(patches)).all()
 
+    def test_padding2(self, device, dtype):
+        img = torch.arange(64, device=device, dtype=dtype).view(1, 1, 8, 8)
+        patches = kornia.contrib.extract_tensor_patches(img, window_size=(2, 2), stride=(2, 2), padding=1)
+        m = kornia.contrib.CombineTensorPatches((8, 8), (2, 2), unpadding=1)
+        assert m(patches).shape == (1, 1, 8, 8)
+        assert (img == m(patches)).all()
+
     def test_gradcheck(self, device, dtype):
         patches = kornia.contrib.extract_tensor_patches(
             torch.arange(16.0, device=device, dtype=dtype).view(1, 1, 4, 4), window_size=(2, 2), stride=(2, 2)
         )
         img = utils.tensor_to_gradcheck_var(patches)  # to var
-        assert gradcheck(kornia.contrib.combine_tensor_patches, (img, (2, 2), (2, 2)), raise_exception=True)
+        assert gradcheck(kornia.contrib.combine_tensor_patches, (img, (4, 4), (2, 2), (2, 2)), raise_exception=True)
 
 
 class TestLambdaModule:
@@ -422,12 +452,20 @@ def test_value(self, device, dtype):
         h = 0.35
         input1 = torch.zeros(B, C, H, W, device=device, dtype=dtype)
         input1[:, :, 1, 1] = 1.0
-        expected_output1 = torch.tensor([[[
-            [1.4142135382, 1.0000000000, 1.4142135382, 2.2360680103],
-            [1.0000000000, 0.0000000000, 1.0000000000, 2.0000000000],
-            [1.4142135382, 1.0000000000, 1.4142135382, 2.2360680103],
-            [2.2360680103, 2.0000000000, 2.2360680103, 2.8284270763]
-        ]]], device=device, dtype=dtype)
+        expected_output1 = torch.tensor(
+            [
+                [
+                    [
+                        [1.4142135382, 1.0000000000, 1.4142135382, 2.2360680103],
+                        [1.0000000000, 0.0000000000, 1.0000000000, 2.0000000000],
+                        [1.4142135382, 1.0000000000, 1.4142135382, 2.2360680103],
+                        [2.2360680103, 2.0000000000, 2.2360680103, 2.8284270763],
+                    ]
+                ]
+            ],
+            device=device,
+            dtype=dtype,
+        )
         output1 = kornia.contrib.distance_transform(input1, kernel_size, h)
         assert_close(expected_output1, output1)