276 lines
12 KiB
Python
276 lines
12 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
import torch.nn.functional as F
|
|
from scipy.signal import get_window
|
|
import scipy
|
|
|
|
def init_kernels(win_len, win_inc, fft_len, win_type=None, invers=False):
|
|
"""
|
|
Initialize the kernels for STFT and iSTFT operations.
|
|
|
|
This function generates the kernel for the convolutional layers used in the short-time Fourier transform (STFT)
|
|
and its inverse (iSTFT). The kernel is created based on the window type and length specified.
|
|
|
|
Args:
|
|
win_len (int): Length of the window.
|
|
win_inc (int): Window increment (hop length).
|
|
fft_len (int): Length of the FFT.
|
|
win_type (str, optional): Type of window to apply (e.g., 'hanning', 'hamming'). Default is None (rectangular window).
|
|
invers (bool, optional): If True, computes the pseudo-inverse of the kernel. Default is False.
|
|
|
|
Returns:
|
|
tuple: A tuple containing:
|
|
- torch.Tensor: The kernel used for convolution, with shape (2 * win_len, 1, fft_len).
|
|
- torch.Tensor: The window applied to the kernel, with shape (1, win_len, 1).
|
|
"""
|
|
if win_type == 'None' or win_type is None:
|
|
window = np.ones(win_len)
|
|
else:
|
|
# Convert 'hanning' to 'hann' if using scipy version 1.10.1 or higher
|
|
if scipy.__version__ >= '1.10.1' and win_type == 'hanning':
|
|
win_type = 'hann'
|
|
window = get_window(win_type, win_len, fftbins=True)**0.5
|
|
|
|
N = fft_len
|
|
fourier_basis = np.fft.rfft(np.eye(N))[:win_len] # Compute Fourier basis for the identity matrix
|
|
real_kernel = np.real(fourier_basis) # Extract real part
|
|
imag_kernel = np.imag(fourier_basis) # Extract imaginary part
|
|
kernel = np.concatenate([real_kernel, imag_kernel], 1).T # Combine real and imaginary parts
|
|
|
|
if invers:
|
|
kernel = np.linalg.pinv(kernel).T # Compute pseudo-inverse if required
|
|
|
|
kernel = kernel * window # Apply window to kernel
|
|
kernel = kernel[:, None, :] # Add singleton dimension for compatibility
|
|
return torch.from_numpy(kernel.astype(np.float32)), torch.from_numpy(window[None,:,None].astype(np.float32))
|
|
|
|
|
|
class ConvSTFT(nn.Module):
|
|
"""
|
|
Convolutional layer that performs Short-Time Fourier Transform (STFT).
|
|
|
|
This class applies the STFT to input signals using convolution with pre-computed kernels.
|
|
It can return either the complex STFT representation or the magnitude and phase components.
|
|
|
|
Attributes:
|
|
weight (nn.Parameter): Learnable convolution kernel for STFT.
|
|
feature_type (str): Specifies whether to return 'complex' or 'real' features.
|
|
stride (int): The stride used for convolution, typically equal to win_inc.
|
|
win_len (int): The length of the window used in STFT.
|
|
dim (int): The FFT length, determining the number of output features.
|
|
"""
|
|
|
|
def __init__(self, win_len, win_inc, fft_len=None, win_type='hamming', feature_type='real', fix=True):
|
|
"""
|
|
Initializes the ConvSTFT layer.
|
|
|
|
Args:
|
|
win_len (int): Length of the window.
|
|
win_inc (int): Window increment (hop length).
|
|
fft_len (int, optional): Length of the FFT. If None, it's computed based on win_len.
|
|
win_type (str, optional): Type of window to use (default is 'hamming').
|
|
feature_type (str, optional): Specifies the output feature type ('real' or 'complex'). Default is 'real'.
|
|
fix (bool, optional): If True, the kernel weights are fixed and not learnable. Default is True.
|
|
"""
|
|
super(ConvSTFT, self).__init__()
|
|
|
|
if fft_len is None:
|
|
self.fft_len = np.int(2**np.ceil(np.log2(win_len))) # Calculate fft_len based on win_len
|
|
else:
|
|
self.fft_len = fft_len
|
|
|
|
kernel, _ = init_kernels(win_len, win_inc, self.fft_len, win_type) # Initialize STFT kernel
|
|
self.weight = nn.Parameter(kernel, requires_grad=(not fix)) # Create a learnable parameter for the kernel
|
|
self.feature_type = feature_type
|
|
self.stride = win_inc
|
|
self.win_len = win_len
|
|
self.dim = self.fft_len
|
|
|
|
def forward(self, inputs):
|
|
"""
|
|
Forward pass through the ConvSTFT layer.
|
|
|
|
Args:
|
|
inputs (torch.Tensor): Input tensor of shape (batch_size, channels, length).
|
|
|
|
Returns:
|
|
tuple or torch.Tensor: Depending on feature_type, returns either:
|
|
- torch.Tensor: The complex STFT output if feature_type is 'complex'.
|
|
- tuple: A tuple containing the magnitude and phase tensors if feature_type is 'real'.
|
|
"""
|
|
if inputs.dim() == 2:
|
|
inputs = torch.unsqueeze(inputs, 1) # Add channel dimension if not present
|
|
|
|
outputs = F.conv1d(inputs, self.weight, stride=self.stride) # Perform convolution to compute STFT
|
|
|
|
if self.feature_type == 'complex':
|
|
return outputs
|
|
else:
|
|
dim = self.dim // 2 + 1 # Calculate the size for the real and imaginary components
|
|
real = outputs[:, :dim, :] # Extract real part
|
|
imag = outputs[:, dim:, :] # Extract imaginary part
|
|
mags = torch.sqrt(real**2 + imag**2) # Compute magnitude
|
|
phase = torch.atan2(imag, real) # Compute phase
|
|
return mags, phase # Return magnitude and phase
|
|
|
|
|
|
class ConviSTFT(nn.Module):
|
|
"""
|
|
Convolutional layer that performs Inverse Short-Time Fourier Transform (iSTFT).
|
|
|
|
This class applies the iSTFT to reconstruct the time-domain signal from the frequency-domain representation
|
|
obtained from the ConvSTFT layer.
|
|
|
|
Attributes:
|
|
weight (nn.Parameter): Learnable convolution kernel for iSTFT.
|
|
feature_type (str): Specifies whether to use 'real' or 'complex' features for reconstruction.
|
|
win_type (str): Type of window used during iSTFT.
|
|
win_len (int): The length of the window used in iSTFT.
|
|
win_inc (int): The window increment (hop length).
|
|
stride (int): The stride used for transposed convolution, typically equal to win_inc.
|
|
dim (int): The FFT length, determining the number of output features.
|
|
window (torch.Tensor): Buffer for the window used in iSTFT.
|
|
enframe (torch.Tensor): Buffer for the framing matrix.
|
|
"""
|
|
|
|
def __init__(self, win_len, win_inc, fft_len=None, win_type='hamming', feature_type='real', fix=True):
|
|
"""
|
|
Initializes the ConviSTFT layer.
|
|
|
|
Args:
|
|
win_len (int): Length of the window.
|
|
win_inc (int): Window increment (hop length).
|
|
fft_len (int, optional): Length of the FFT. If None, it's computed based on win_len.
|
|
win_type (str, optional): Type of window to use (default is 'hamming').
|
|
feature_type (str, optional): Specifies the output feature type ('real' or 'complex'). Default is 'real'.
|
|
fix (bool, optional): If True, the kernel weights are fixed and not learnable. Default is True.
|
|
"""
|
|
super(ConviSTFT, self).__init__()
|
|
|
|
if fft_len is None:
|
|
self.fft_len = np.int(2**np.ceil(np.log2(win_len))) # Calculate fft_len based on win_len
|
|
else:
|
|
self.fft_len = fft_len
|
|
|
|
kernel, window = init_kernels(win_len, win_inc, self.fft_len, win_type, invers=True) # Initialize iSTFT kernel
|
|
self.weight = nn.Parameter(kernel, requires_grad=(not fix)) # Create a learnable parameter for the kernel
|
|
self.feature_type = feature_type
|
|
self.win_type = win_type
|
|
self.win_len = win_len
|
|
self.win_inc = win_inc
|
|
self.stride = win_inc
|
|
self.dim = self.fft_len
|
|
self.register_buffer('window', window) # Register the window as a buffer
|
|
self.register_buffer('enframe', torch.eye(win_len)[:, None, :]) # Framing matrix for overlap-add method
|
|
|
|
def forward(self, inputs, phase=None):
|
|
"""
|
|
Forward pass through the ConviSTFT layer.
|
|
|
|
Args:
|
|
inputs (torch.Tensor): Input tensor of shape [B, N+2, T] for complex spectra
|
|
or [B, N//2+1, T] for magnitude spectra.
|
|
phase (torch.Tensor, optional): Phase tensor of shape [B, N//2+1, T]. If provided, used to reconstruct the complex spectra.
|
|
|
|
Returns:
|
|
torch.Tensor: Reconstructed time-domain signal.
|
|
"""
|
|
if phase is not None:
|
|
# Reconstruct real and imaginary components from magnitude and phase
|
|
real = inputs * torch.cos(phase) # Real part
|
|
imag = inputs * torch.sin(phase) # Imaginary part
|
|
inputs = torch.cat([real, imag], 1) # Concatenate to form complex input
|
|
|
|
outputs = F.conv_transpose1d(inputs, self.weight, stride=self.stride) # Perform transposed convolution for iSTFT
|
|
|
|
# Compute the overlap-add normalization
|
|
t = self.window.repeat(1, 1, inputs.size(-1))**2
|
|
coff = F.conv_transpose1d(t, self.enframe, stride=self.stride) # Apply the framing matrix for overlap-add
|
|
outputs = outputs / (coff + 1e-8) # Normalize the output to prevent division by zero
|
|
return outputs
|
|
|
|
|
|
def test_fft():
|
|
"""
|
|
Test the ConvSTFT layer against Librosa's STFT implementation.
|
|
|
|
This function generates a random input signal and computes its STFT using the ConvSTFT layer,
|
|
then compares the output with the STFT computed using Librosa to ensure correctness.
|
|
"""
|
|
torch.manual_seed(20)
|
|
win_len = 320
|
|
win_inc = 160
|
|
fft_len = 512
|
|
inputs = torch.randn([1, 1, 16000*4]) # Random input tensor
|
|
fft = ConvSTFT(win_len, win_inc, fft_len, win_type='hanning', feature_type='real') # Initialize ConvSTFT
|
|
|
|
outputs1 = fft(inputs)[0] # Compute STFT using ConvSTFT
|
|
outputs1 = outputs1.numpy()[0] # Convert to NumPy array for comparison
|
|
np_inputs = inputs.numpy().reshape([-1]) # Reshape input for Librosa
|
|
librosa_stft = librosa.stft(np_inputs, win_length=win_len, n_fft=fft_len, hop_length=win_inc, center=False) # Compute STFT using Librosa
|
|
print(np.mean((outputs1 - np.abs(librosa_stft))**2)) # Print mean squared error between the two STFT outputs
|
|
|
|
|
|
def test_ifft1():
|
|
"""
|
|
Test the ConviSTFT layer by reconstructing a waveform from the STFT output.
|
|
|
|
This function reads an audio file, applies the ConvSTFT to compute its STFT, and then
|
|
uses the ConviSTFT to reconstruct the time-domain signal. The reconstructed signal is saved to a file
|
|
and compared to the original to evaluate the reconstruction accuracy.
|
|
"""
|
|
import soundfile as sf
|
|
N = 100
|
|
inc = 75
|
|
fft_len = 512
|
|
torch.manual_seed(N)
|
|
|
|
# Read input audio file and reshape
|
|
data = sf.read('../../wavs/ori.wav')[0]
|
|
inputs = data.reshape([1, 1, -1]) # Reshape to [1, 1, length]
|
|
|
|
fft = ConvSTFT(N, inc, fft_len=fft_len, win_type='hanning', feature_type='complex') # Initialize ConvSTFT
|
|
ifft = ConviSTFT(N, inc, fft_len=fft_len, win_type='hanning', feature_type='complex') # Initialize ConviSTFT
|
|
|
|
inputs = torch.from_numpy(inputs.astype(np.float32)) # Convert to torch tensor
|
|
outputs1 = fft(inputs) # Compute STFT
|
|
outputs2 = ifft(outputs1) # Reconstruct waveform from STFT
|
|
sf.write('conv_stft.wav', outputs2.numpy()[0, 0, :], 16000) # Save reconstructed waveform to file
|
|
print('wav MSE', torch.mean(torch.abs(inputs[..., :outputs2.size(2)] - outputs2))) # Print mean squared error
|
|
|
|
|
|
def test_ifft2():
|
|
"""
|
|
Test the iSTFT reconstruction from a random input signal.
|
|
|
|
This function generates a random signal, computes its STFT, and then reconstructs it using the ConviSTFT layer.
|
|
The reconstructed waveform is saved to a file, and the mean squared error is printed to evaluate accuracy.
|
|
"""
|
|
N = 400
|
|
inc = 100
|
|
fft_len = 512
|
|
np.random.seed(20)
|
|
torch.manual_seed(20)
|
|
|
|
# Generate a random signal
|
|
t = np.random.randn(16000*4) * 0.005
|
|
t = np.clip(t, -1, 1) # Clip to [-1, 1] range
|
|
input = torch.from_numpy(t[None, None, :].astype(np.float32)) # Reshape for input
|
|
|
|
fft = ConvSTFT(N, inc, fft_len=fft_len, win_type='hanning', feature_type='complex') # Initialize ConvSTFT
|
|
ifft = ConviSTFT(N, inc, fft_len=fft_len, win_type='hanning', feature_type='complex') # Initialize ConviSTFT
|
|
|
|
out1 = fft(input) # Compute STFT
|
|
output = ifft(out1) # Reconstruct waveform from STFT
|
|
print('random MSE', torch.mean(torch.abs(input - output)**2)) # Print mean squared error
|
|
import soundfile as sf
|
|
sf.write('zero.wav', output[0, 0].numpy(), 16000) # Save reconstructed waveform to file
|
|
|
|
|
|
if __name__ == '__main__':
|
|
#test_fft()
|
|
test_ifft1() # Run the iSTFT reconstruction test
|
|
#test_ifft2()
|
|
|