labhamlet/gramt-binaural-frame

1import torch2from torch import nn3 4 5def generate_patches(input, fstride, tstride, fshape, tshape):6    r"""Function that extract patches from tensors and stacks them.7 8    See :class:`~kornia.contrib.ExtractTensorPatches` for details.9 10    Args:11        input: tensor image where to extract the patches with shape :math:`(B, C, H, W)`.12 13    Returns:14        the tensor with the extracted patches with shape :math:`(B, N, C, H_{out}, W_{out})`.15 16    Examples:17        >>> input = torch.arange(9.).view(1, 1, 3, 3)18        >>> patches = extract_tensor_patches(input, (2, 3))19        >>> input20        tensor([[[[0., 1., 2.],21                  [3., 4., 5.],22                  [6., 7., 8.]]]])23        >>> patches[:, -1]24        tensor([[[[3., 4., 5.],25                  [6., 7., 8.]]]])26 27    """28    batch_size, num_channels = input.size()[:2]29    dims = range(2, input.dim())30    for dim, patch_size, stride in zip(dims, (fshape, tshape), (fstride, tstride)):31        input = input.unfold(dim, patch_size, stride)32    input = input.permute(0, *dims, 1, *(dim + len(dims) for dim in dims)).contiguous()33    return input.view(batch_size, -1, num_channels, fshape, tshape)34 35 36def combine_patches(37    patches,38    original_size,39    fstride,40    tstride,41    fshape,42    tshape,43    eps: float = 1e-8,44):45    r"""Restore input from patches.46 47    See :class:`~kornia.contrib.CombineTensorPatches` for details.48 49    Args:50        patches: patched tensor with shape :math:`(B, N, C, H_{out}, W_{out})`.51 52    Return:53        The combined patches in an image tensor with shape :math:`(B, C, H, W)`.54 55    Example:56        >>> out = extract_tensor_patches(torch.arange(16).view(1, 1, 4, 4), window_size=(2, 2), stride=(2, 2))57        >>> combine_tensor_patches(out, original_size=(4, 4), window_size=(2, 2), stride=(2, 2))58        tensor([[[[ 0,  1,  2,  3],59                  [ 4,  5,  6,  7],60                  [ 8,  9, 10, 11],61                  [12, 13, 14, 15]]]])62 63    .. note::64        This function is supposed to be used in conjunction with :func:`extract_tensor_patches`.65 66    """67    if patches.ndim != 5:68        raise ValueError(69            f"Invalid input shape, we expect BxNxCxHxW. Got: {patches.shape}"70        )71    ones = torch.ones(72        patches.shape[0],73        patches.shape[2],74        original_size[0],75        original_size[1],76        device=patches.device,77        dtype=patches.dtype,78    )79    restored_size = ones.shape[2:]80 81    patches = patches.permute(0, 2, 3, 4, 1)82    patches = patches.reshape(patches.shape[0], -1, patches.shape[-1])83    int_flag = 084    if not torch.is_floating_point(patches):85        int_flag = 186        dtype = patches.dtype87        patches = patches.float()88        ones = ones.float()89 90    # Calculate normalization map91    unfold_ones = torch.nn.functional.unfold(92        ones, kernel_size=(fshape, tshape), stride=(fstride, tstride)93    )94    norm_map = torch.nn.functional.fold(95        input=unfold_ones,96        output_size=restored_size,97        kernel_size=(fshape, tshape),98        stride=(fstride, tstride),99    )100    # Restored tensor101    saturated_restored_tensor = torch.nn.functional.fold(102        input=patches,103        output_size=restored_size,104        kernel_size=(fshape, tshape),105        stride=(fstride, tstride),106    )107    # Remove satuation effect due to multiple summations108    restored_tensor = saturated_restored_tensor / (norm_map + eps)109    if int_flag:110        restored_tensor = restored_tensor.to(dtype)111    return restored_tensor112 113 114# get the shape of intermediate representation.115def get_shape(fstride, tstride, input_fdim, input_tdim, fshape, tshape):116    test_input = torch.randn(1, 2, input_fdim, input_tdim)117    test_proj = nn.Conv2d(118        2,119        2,120        kernel_size=(fshape, tshape),121        stride=(fstride, tstride),122    )123    test_out = test_proj(test_input)124    f_dim = test_out.shape[2]125    t_dim = test_out.shape[3]126    return f_dim, t_dim127