Seich
/
mayavoz
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

stft

This commit is contained in:
shahules786 2022-10-29 11:34:51 +05:30
parent 085a85d9ae
commit c18a85b5c8
1 changed files with 16 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
enhancer/utils/transforms.py
View File

@ -14,9 +14,12 @@ class ConvFFT(nn.Module):
nfft: Optional[int] = None, nfft: Optional[int] = None,
window: str = "hamming", window: str = "hamming",
): ):
super().__init__()
self.window_len = window_len self.window_len = window_len
self.nfft = nfft if nfft else np.int(2 ** np.ceil(np.log2(window_len))) self.nfft = nfft if nfft else np.int(2 ** np.ceil(np.log2(window_len)))
self.window = get_window(window, window_len) self.window = torch.from_numpy(
get_window(window, window_len).astype("float32")
)
@property @property
def init_kernel(self): def init_kernel(self):
@ -24,8 +27,9 @@ class ConvFFT(nn.Module):
fourier_basis = np.fft.rfft(np.eye(self.nfft))[: self.window_len] fourier_basis = np.fft.rfft(np.eye(self.nfft))[: self.window_len]
real, imag = np.real(fourier_basis), np.imag(fourier_basis) real, imag = np.real(fourier_basis), np.imag(fourier_basis)
kernel = np.concatenate([real, imag], 1).T kernel = np.concatenate([real, imag], 1).T
kernel = torch.from_numpy(kernel.astype("float32")).unsqueeze(1)
kernel *= self.window kernel *= self.window
return torch.from_numpy(kernel) return kernel
class ConvSTFT(ConvFFT): class ConvSTFT(ConvFFT):
@ -36,9 +40,7 @@ class ConvSTFT(ConvFFT):
nfft: Optional[int] = None, nfft: Optional[int] = None,
window: str = "hamming", window: str = "hamming",
): ):
super(self, ConvSTFT).__init__( super().__init__(window_len=window_len, nfft=nfft, window=window)
window_len=window_len, nfft=nfft, window=window
)
self.hop_size = hop_size if hop_size else window_len // 2 self.hop_size = hop_size if hop_size else window_len // 2
self.register_buffer("weight", self.init_kernel) self.register_buffer("weight", self.init_kernel)
@ -69,12 +71,10 @@ class ConviSTFT(ConvFFT):
nfft: Optional[int] = None, nfft: Optional[int] = None,
window: str = "hamming", window: str = "hamming",
): ):
super(self, ConvSTFT).__init__( super().__init__(window_len=window_len, nfft=nfft, window=window)
window_len=window_len, nfft=nfft, window=window
)
self.hop_size = hop_size if hop_size else window_len // 2 self.hop_size = hop_size if hop_size else window_len // 2
self.register_buffer("weight", self.init_kernel) self.register_buffer("weight", self.init_kernel)
self.register_buffer("enframe", np.eye(window_len).unsqueeze(1)) self.register_buffer("enframe", torch.eye(window_len).unsqueeze(1))
def forward(self, input, phase=None): def forward(self, input, phase=None):
@ -82,3 +82,10 @@ class ConviSTFT(ConvFFT):
real = input * torch.cos(phase) real = input * torch.cos(phase)
imag = input * torch.sin(phase) imag = input * torch.sin(phase)
input = torch.cat([real, imag], 1) input = torch.cat([real, imag], 1)
out = F.conv_transpose1d(input, self.weight, stride=self.hop_size)
coeff = self.window.unsqueeze(1).repeat(1, 1, input.size(-1)) ** 2
coeff = F.conv_transpose1d(coeff, self.enframe, stride=self.hop_size)
out = out / coeff
pad = self.window_len - self.hop_size
out = out[..., pad:-pad]
return out