from typing import Optional, Tuple, Union, cast import torch import torch.nn as nn import torch.nn.functional as F from torch.nn.modules.utils import _pair PadType = Union[Tuple[int, int], Tuple[int, int, int, int]] def compute_padding( original_size: Union[int, Tuple[int, int]], window_size: Union[int, Tuple[int, int]] ) -> Tuple[int, int, int, int]: r"""Compute required padding to ensure chaining of :func:`extract_tensor_patches` and :func:`combine_tensor_patches` produces expected result. Args: original_size: the size of the original tensor. window_size: the size of the sliding window used while extracting patches. Return: The required padding for `(top, bottom, left, right)` as a tuple of 4 ints. Example: >>> image = torch.arange(12).view(1, 1, 4, 3) >>> padding = compute_padding((4,3), (3,3)) >>> out = extract_tensor_patches(image, window_size=(3, 3), stride=(3, 3), padding=padding) >>> combine_tensor_patches(out, original_size=(4, 3), window_size=(3, 3), stride=(3, 3), unpadding=padding) tensor([[[[ 0, 1, 2], [ 3, 4, 5], [ 6, 7, 8], [ 9, 10, 11]]]]) .. note:: This function is supposed to be used in conjunction with :func:`extract_tensor_patches` and :func:`combine_tensor_patches`. """ original_size = cast(Tuple[int, int], _pair(original_size)) window_size = cast(Tuple[int, int], _pair(window_size)) def paddim(dim1: int, dim2: int) -> Tuple[int, int]: if dim1 % dim2 == 0: p1 = 0 p2 = 0 else: tmp = dim2 - (dim1 % dim2) if tmp % 2 == 0: p1 = p2 = tmp // 2 else: p1 = tmp p2 = 0 return p1, p2 padt, padb = paddim(original_size[0], window_size[0]) padl, padr = paddim(original_size[1], window_size[1]) return (padt, padb, padl, padr) class ExtractTensorPatches(nn.Module): r"""Module that extract patches from tensors and stack them. In the simplest case, the output value of the operator with input size :math:`(B, C, H, W)` is :math:`(B, N, C, H_{out}, W_{out})`. where - :math:`B` is the batch size. - :math:`N` denotes the total number of extracted patches stacked in - :math:`C` denotes the number of input channels. - :math:`H`, :math:`W` the input height and width of the input in pixels. - :math:`H_{out}`, :math:`W_{out}` denote to denote to the patch size 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 be either: - a single ``int`` -- in which case the same value is used for the height and width dimension. - a ``tuple`` of two ints -- in which case, the first `int` is used for the height dimension, and the second `int` for the width dimension. :attr:`padding` can also be a ``tuple`` of four ints -- in which case, the first two ints are for the height dimension while the last two ints are for the width dimension. Args: window_size: the size of the sliding window and the output patch size. stride: stride of the sliding window. padding: Zero-padding added to both side of the input. Shape: - Input: :math:`(B, C, H, W)` - Output: :math:`(B, N, C, H_{out}, W_{out})` Returns: the tensor with the extracted patches. Examples: >>> input = torch.arange(9.).view(1, 1, 3, 3) >>> patches = extract_tensor_patches(input, (2, 3)) >>> input tensor([[[[0., 1., 2.], [3., 4., 5.], [6., 7., 8.]]]]) >>> patches[:, -1] tensor([[[[3., 4., 5.], [6., 7., 8.]]]]) """ def __init__( self, window_size: Union[int, Tuple[int, int]], stride: Optional[Union[int, Tuple[int, int]]] = 1, padding: Optional[Union[int, PadType]] = 0, ) -> None: super().__init__() self.window_size: Tuple[int, int] = _pair(window_size) self.stride: Tuple[int, int] = _pair(stride) self.padding: PadType = _pair(padding) def forward(self, input: torch.Tensor) -> torch.Tensor: # type: ignore return extract_tensor_patches(input, self.window_size, stride=self.stride, padding=self.padding) class CombineTensorPatches(nn.Module): r"""Module that combine patches from tensors. In the simplest case, the output value of the operator with input size :math:`(B, N, C, H_{out}, W_{out})` is :math:`(B, C, H, W)`. where - :math:`B` is the batch size. - :math:`N` denotes the total number of extracted patches stacked in - :math:`C` denotes the number of input channels. - :math:`H`, :math:`W` the input height and width of the input in pixels. - :math:`H_{out}`, :math:`W_{out}` denote to denote to the patch size defined in the function signature. left-right and top-bottom order. * :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 height and width dimension. - a ``tuple`` of two ints -- in which case, the first `int` is used for the height dimension, and the second `int` for the width dimension. :attr:`unpadding` can also be a ``tuple`` of four ints -- in which case, the first two ints are for the height dimension while the last two ints are for the width dimension. Args: patches: patched tensor. 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: - Input: :math:`(B, N, C, H_{out}, W_{out})` - Output: :math:`(B, C, H, W)` Example: >>> out = extract_tensor_patches(torch.arange(16).view(1, 1, 4, 4), 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, PadType] = 0, ) -> None: super().__init__() self.original_size: Tuple[int, int] = _pair(original_size) self.window_size: Tuple[int, int] = _pair(window_size) self.unpadding: PadType = _pair(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: Union[int, Tuple[int, int]], window_size: Union[int, Tuple[int, int]], stride: Union[int, Tuple[int, int]], unpadding: Union[int, PadType] = 0, ) -> torch.Tensor: r"""Restore input from patches. See :class:`~kornia.contrib.CombineTensorPatches` for details. Args: patches: patched tensor with shape :math:`(B, N, C, H_{out}, W_{out})`. 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. Return: The combined patches in an image tensor with shape :math:`(B, C, H, W)`. Example: >>> out = extract_tensor_patches(torch.arange(16).view(1, 1, 4, 4), 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 len(patches.shape) != 5: raise ValueError(f"Invalid input shape, we expect BxNxCxHxW. Got: {patches.shape}") original_size = cast(Tuple[int, int], _pair(original_size)) window_size = cast(Tuple[int, int], _pair(window_size)) stride = cast(Tuple[int, int], _pair(stride)) 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 unpadding: unpadding = cast(PadType, _pair(unpadding)) if len(unpadding) not in [2, 4]: raise AssertionError("Unpadding must be either an int, tuple of two ints or tuple of four ints") if len(unpadding) == 2: pad_vert = _pair(unpadding[0]) pad_horz = _pair(unpadding[1]) else: pad_vert = unpadding[:2] pad_horz = unpadding[2:] unpadding = cast(Tuple[int, int, int, int], pad_horz + pad_vert) hpad_check = (original_size[0] + unpadding[2] + unpadding[3]) % window_size[0] == 0 wpad_check = (original_size[1] + unpadding[0] + unpadding[1]) % window_size[1] == 0 if not hpad_check or not wpad_check: raise NotImplementedError("Insufficient padding") window_size = ( (original_size[0] + (unpadding[2] + unpadding[3])) // window_size[0], (original_size[1] + (unpadding[0] + unpadding[1])) // 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: unpadding = cast(Tuple[int, int, int, int], unpadding) restored_tensor = F.pad(restored_tensor, [-i for i in unpadding]) return restored_tensor def _extract_tensor_patchesnd( input: torch.Tensor, window_sizes: Tuple[int, ...], strides: Tuple[int, ...] ) -> torch.Tensor: batch_size, num_channels = input.size()[:2] dims = range(2, input.dim()) for dim, patch_size, stride in zip(dims, window_sizes, strides): input = input.unfold(dim, patch_size, stride) input = input.permute(0, *dims, 1, *(dim + len(dims) for dim in dims)).contiguous() return input.view(batch_size, -1, num_channels, *window_sizes) def extract_tensor_patches( input: torch.Tensor, window_size: Union[int, Tuple[int, int]], stride: Union[int, Tuple[int, int]] = 1, padding: Union[int, PadType] = 0, ) -> torch.Tensor: r"""Function that extract patches from tensors and stack them. See :class:`~kornia.contrib.ExtractTensorPatches` for details. Args: input: tensor image where to extract the patches with shape :math:`(B, C, H, W)`. window_size: the size of the sliding window and the output patch size. stride: stride of the sliding window. padding: Zero-padding added to both side of the input. Returns: the tensor with the extracted patches with shape :math:`(B, N, C, H_{out}, W_{out})`. Examples: >>> input = torch.arange(9.).view(1, 1, 3, 3) >>> patches = extract_tensor_patches(input, (2, 3)) >>> input tensor([[[[0., 1., 2.], [3., 4., 5.], [6., 7., 8.]]]]) >>> patches[:, -1] tensor([[[[3., 4., 5.], [6., 7., 8.]]]]) """ if not torch.is_tensor(input): raise TypeError(f"Input input type is not a torch.Tensor. Got {type(input)}") if len(input.shape) != 4: raise ValueError(f"Invalid input shape, we expect BxCxHxW. Got: {input.shape}") if padding: padding = cast(PadType, _pair(padding)) if len(padding) not in [2, 4]: raise AssertionError("Padding must be either an int, tuple of two ints or tuple of four ints") if len(padding) == 2: pad_vert = _pair(padding[0]) pad_horz = _pair(padding[1]) else: pad_vert = padding[:2] pad_horz = padding[2:] padding = cast(Tuple[int, int, int, int], pad_horz + pad_vert) input = F.pad(input, padding) return _extract_tensor_patchesnd(input, _pair(window_size), _pair(stride))