| """The 1D discrete wavelet transform for PyTorch.""" | |
| from einops import rearrange | |
| import pywt | |
| import torch | |
| from torch import nn | |
| from torch.nn import functional as F | |
| from typing import Literal | |
| def get_filter_bank(wavelet): | |
| filt = torch.tensor(pywt.Wavelet(wavelet).filter_bank) | |
| if wavelet.startswith("bior") and torch.all(filt[:, 0] == 0): | |
| filt = filt[:, 1:] | |
| return filt | |
| class WaveletEncode1d(nn.Module): | |
| def __init__(self, | |
| channels, | |
| levels, | |
| wavelet: Literal["bior2.2", "bior2.4", "bior2.6", "bior2.8", "bior4.4", "bior6.8"] = "bior4.4"): | |
| super().__init__() | |
| self.wavelet = wavelet | |
| self.channels = channels | |
| self.levels = levels | |
| filt = get_filter_bank(wavelet) | |
| assert filt.shape[-1] % 2 == 1 | |
| kernel = filt[:2, None] | |
| kernel = torch.flip(kernel, dims=(-1,)) | |
| index_i = torch.repeat_interleave(torch.arange(2), channels) | |
| index_j = torch.tile(torch.arange(channels), (2,)) | |
| kernel_final = torch.zeros(channels * 2, channels, filt.shape[-1]) | |
| kernel_final[index_i * channels + index_j, index_j] = kernel[index_i, 0] | |
| self.register_buffer("kernel", kernel_final) | |
| def forward(self, x): | |
| for i in range(self.levels): | |
| low, rest = x[:, : self.channels], x[:, self.channels :] | |
| pad = self.kernel.shape[-1] // 2 | |
| low = F.pad(low, (pad, pad), "reflect") | |
| low = F.conv1d(low, self.kernel, stride=2) | |
| rest = rearrange( | |
| rest, "n (c c2) (l l2) -> n (c l2 c2) l", l2=2, c2=self.channels | |
| ) | |
| x = torch.cat([low, rest], dim=1) | |
| return x | |
| class WaveletDecode1d(nn.Module): | |
| def __init__(self, | |
| channels, | |
| levels, | |
| wavelet: Literal["bior2.2", "bior2.4", "bior2.6", "bior2.8", "bior4.4", "bior6.8"] = "bior4.4"): | |
| super().__init__() | |
| self.wavelet = wavelet | |
| self.channels = channels | |
| self.levels = levels | |
| filt = get_filter_bank(wavelet) | |
| assert filt.shape[-1] % 2 == 1 | |
| kernel = filt[2:, None] | |
| index_i = torch.repeat_interleave(torch.arange(2), channels) | |
| index_j = torch.tile(torch.arange(channels), (2,)) | |
| kernel_final = torch.zeros(channels * 2, channels, filt.shape[-1]) | |
| kernel_final[index_i * channels + index_j, index_j] = kernel[index_i, 0] | |
| self.register_buffer("kernel", kernel_final) | |
| def forward(self, x): | |
| for i in range(self.levels): | |
| low, rest = x[:, : self.channels * 2], x[:, self.channels * 2 :] | |
| pad = self.kernel.shape[-1] // 2 + 2 | |
| low = rearrange(low, "n (l2 c) l -> n c (l l2)", l2=2) | |
| low = F.pad(low, (pad, pad), "reflect") | |
| low = rearrange(low, "n c (l l2) -> n (l2 c) l", l2=2) | |
| low = F.conv_transpose1d( | |
| low, self.kernel, stride=2, padding=self.kernel.shape[-1] // 2 | |
| ) | |
| low = low[..., pad - 1 : -pad] | |
| rest = rearrange( | |
| rest, "n (c l2 c2) l -> n (c c2) (l l2)", l2=2, c2=self.channels | |
| ) | |
| x = torch.cat([low, rest], dim=1) | |
| return x |