252 lines
9.1 KiB
Python
252 lines
9.1 KiB
Python
import torch.nn as nn
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import os
|
|
import sys
|
|
sys.path.append(os.path.dirname(__file__))
|
|
from models.frcrn_se.conv_stft import ConvSTFT, ConviSTFT
|
|
import numpy as np
|
|
from models.frcrn_se.unet import UNet
|
|
|
|
class FRCRN_Wrapper_StandAlone(nn.Module):
|
|
"""
|
|
A wrapper class for the DCCRN model used in standalone mode.
|
|
|
|
This class initializes the DCCRN model with predefined parameters and provides a forward method to process
|
|
input audio signals for speech enhancement.
|
|
|
|
Args:
|
|
args: Arguments containing model configuration (not used in this wrapper).
|
|
"""
|
|
def __init__(self, args):
|
|
super(FRCRN_Wrapper_StandAlone, self).__init__()
|
|
# Initialize the DCCRN model with specific parameters
|
|
self.model = DCCRN(
|
|
complex=True,
|
|
model_complexity=45,
|
|
model_depth=14,
|
|
log_amp=False,
|
|
padding_mode="zeros",
|
|
win_len=640,
|
|
win_inc=320,
|
|
fft_len=640,
|
|
win_type='hanning'
|
|
)
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Forward pass of the model.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor representing audio signals.
|
|
|
|
Returns:
|
|
torch.Tensor: Processed output tensor after applying the model.
|
|
"""
|
|
output = self.model(x)
|
|
return output[1][0] # Return estimated waveform
|
|
|
|
|
|
class FRCRN_SE_16K(nn.Module):
|
|
"""
|
|
A class for the FRCRN model specifically configured for 16 kHz input signals.
|
|
|
|
This class allows for customization of model parameters based on provided arguments.
|
|
|
|
Args:
|
|
args: Configuration parameters for the model.
|
|
"""
|
|
def __init__(self, args):
|
|
super(FRCRN_SE_16K, self).__init__()
|
|
# Initialize the DCCRN model with parameters from args
|
|
self.model = DCCRN(
|
|
complex=True,
|
|
model_complexity=45,
|
|
model_depth=14,
|
|
log_amp=False,
|
|
padding_mode="zeros",
|
|
win_len=args.win_len,
|
|
win_inc=args.win_inc,
|
|
fft_len=args.fft_len,
|
|
win_type=args.win_type
|
|
)
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Forward pass of the model.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor representing audio signals.
|
|
|
|
Returns:
|
|
torch.Tensor: Processed output tensor after applying the model.
|
|
"""
|
|
output = self.model(x)
|
|
return output[1][0] # Return estimated waveform
|
|
|
|
|
|
class DCCRN(nn.Module):
|
|
"""
|
|
We implemented our FRCRN model on the basis of DCCRN rep (https://github.com/huyanxin/DeepComplexCRN) for complex speech enhancement.
|
|
|
|
The DCCRN model (Paper: https://arxiv.org/abs/2008.00264) employs a convolutional short-time Fourier transform (STFT)
|
|
and a UNet architecture for estimating clean speech from noisy inputs, FRCRN uses an enhanced
|
|
Unet architecture.
|
|
|
|
Args:
|
|
complex (bool): Flag to determine whether to use complex numbers.
|
|
model_complexity (int): Complexity level for the model.
|
|
model_depth (int): Depth of the UNet model (14 or 20).
|
|
log_amp (bool): Whether to use log amplitude to estimate signals.
|
|
padding_mode (str): Padding mode for convolutions ('zeros', 'reflect').
|
|
win_len (int): Window length for STFT.
|
|
win_inc (int): Window increment for STFT.
|
|
fft_len (int): FFT length.
|
|
win_type (str): Window type for STFT (e.g., 'hanning').
|
|
"""
|
|
def __init__(self, complex, model_complexity, model_depth, log_amp, padding_mode, win_len=400, win_inc=100, fft_len=512, win_type='hanning'):
|
|
super().__init__()
|
|
self.feat_dim = fft_len // 2 + 1
|
|
|
|
self.win_len = win_len
|
|
self.win_inc = win_inc
|
|
self.fft_len = fft_len
|
|
self.win_type = win_type
|
|
|
|
# Initialize STFT and iSTFT layers
|
|
fix = True # Fixed STFT parameters
|
|
self.stft = ConvSTFT(self.win_len, self.win_inc, self.fft_len, self.win_type, feature_type='complex', fix=fix)
|
|
self.istft = ConviSTFT(self.win_len, self.win_inc, self.fft_len, self.win_type, feature_type='complex', fix=fix)
|
|
|
|
# Initialize two UNet models for estimating complex masks
|
|
self.unet = UNet(1, complex=complex, model_complexity=model_complexity, model_depth=model_depth, padding_mode=padding_mode)
|
|
self.unet2 = UNet(1, complex=complex, model_complexity=model_complexity, model_depth=model_depth, padding_mode=padding_mode)
|
|
|
|
def forward(self, inputs):
|
|
"""
|
|
Forward pass of the FRCRN model.
|
|
|
|
Args:
|
|
inputs (torch.Tensor): Input tensor representing audio signals.
|
|
|
|
Returns:
|
|
list: A list containing estimated spectral features, waveform, and masks.
|
|
"""
|
|
out_list = []
|
|
# Compute the complex spectrogram using STFT
|
|
cmp_spec = self.stft(inputs) # [B, D*2, T]
|
|
cmp_spec = torch.unsqueeze(cmp_spec, 1) # [B, 1, D*2, T]
|
|
|
|
# Split into real and imaginary parts
|
|
cmp_spec = torch.cat([
|
|
cmp_spec[:, :, :self.feat_dim, :], # Real part
|
|
cmp_spec[:, :, self.feat_dim:, :], # Imaginary part
|
|
], 1) # [B, 2, D, T]
|
|
|
|
cmp_spec = torch.unsqueeze(cmp_spec, 4) # [B, 2, D, T, 1]
|
|
cmp_spec = torch.transpose(cmp_spec, 1, 4) # [B, 1, D, T, 2]
|
|
|
|
# Pass through the UNet to estimate masks
|
|
unet1_out = self.unet(cmp_spec) # First UNet output
|
|
cmp_mask1 = torch.tanh(unet1_out) # First mask
|
|
|
|
unet2_out = self.unet2(unet1_out) # Second UNet output
|
|
cmp_mask2 = torch.tanh(unet2_out) # Second mask
|
|
cmp_mask2 = cmp_mask2 + cmp_mask1 # Combine masks
|
|
|
|
# Apply the estimated mask to the complex spectrogram
|
|
est_spec, est_wav, est_mask = self.apply_mask(cmp_spec, cmp_mask2)
|
|
out_list.append(est_spec)
|
|
out_list.append(est_wav)
|
|
out_list.append(est_mask)
|
|
return out_list
|
|
|
|
def inference(self, inputs):
|
|
"""
|
|
Inference method for the FRCRN model.
|
|
|
|
This method performs a forward pass through the model to estimate the clean waveform
|
|
from the noisy input.
|
|
|
|
Args:
|
|
inputs (torch.Tensor): Input tensor representing audio signals.
|
|
|
|
Returns:
|
|
torch.Tensor: Estimated waveform after processing.
|
|
"""
|
|
# Compute the complex spectrogram using STFT
|
|
cmp_spec = self.stft(inputs) # [B, D*2, T]
|
|
cmp_spec = torch.unsqueeze(cmp_spec, 1) # [B, 1, D*2, T]
|
|
|
|
# Split into real and imaginary parts
|
|
cmp_spec = torch.cat([
|
|
cmp_spec[:, :, :self.feat_dim, :], # Real part
|
|
cmp_spec[:, :, self.feat_dim:, :], # Imaginary part
|
|
], 1) # [B, 2, D, T]
|
|
|
|
cmp_spec = torch.unsqueeze(cmp_spec, 4) # [B, 2, D, T, 1]
|
|
cmp_spec = torch.transpose(cmp_spec, 1, 4) # [B, 1, D, T, 2]
|
|
|
|
# Pass through the UNet to estimate masks
|
|
unet1_out = self.unet(cmp_spec)
|
|
cmp_mask1 = torch.tanh(unet1_out)
|
|
|
|
unet2_out = self.unet2(unet1_out)
|
|
cmp_mask2 = torch.tanh(unet2_out)
|
|
cmp_mask2 = cmp_mask2 + cmp_mask1 # Combine masks
|
|
|
|
# Apply the estimated mask to compute the estimated waveform
|
|
_, est_wav, _ = self.apply_mask(cmp_spec, cmp_mask2)
|
|
return est_wav[0] # Return the estimated waveform
|
|
|
|
def apply_mask(self, cmp_spec, cmp_mask):
|
|
"""
|
|
Apply the estimated masks to the complex spectrogram.
|
|
|
|
Args:
|
|
cmp_spec (torch.Tensor): Complex spectrogram tensor.
|
|
cmp_mask (torch.Tensor): Estimated mask tensor.
|
|
|
|
Returns:
|
|
tuple: Estimated spectrogram, waveform, and mask.
|
|
"""
|
|
# Compute the estimated complex spectrogram using masks
|
|
est_spec = torch.cat([
|
|
cmp_spec[:, :, :, :, 0] * cmp_mask[:, :, :, :, 0] - cmp_spec[:, :, :, :, 1] * cmp_mask[:, :, :, :, 1],
|
|
cmp_spec[:, :, :, :, 0] * cmp_mask[:, :, :, :, 1] + cmp_spec[:, :, :, :, 1] * cmp_mask[:, :, :, :, 0]
|
|
], 1) # Combine real and imaginary parts
|
|
|
|
est_spec = torch.cat([est_spec[:, 0, :, :], est_spec[:, 1, :, :]], 1) # Flatten dimensions
|
|
cmp_mask = torch.squeeze(cmp_mask, 1)
|
|
cmp_mask = torch.cat([cmp_mask[:, :, :, 0], cmp_mask[:, :, :, 1]], 1) # Combine masks
|
|
|
|
est_wav = self.istft(est_spec) # Inverse STFT to obtain waveform
|
|
est_wav = torch.squeeze(est_wav, 1) # Remove unnecessary dimensions
|
|
return est_spec, est_wav, cmp_mask
|
|
|
|
def get_params(self, weight_decay=0.0):
|
|
"""
|
|
Get parameters for optimization with optional weight decay.
|
|
|
|
Args:
|
|
weight_decay (float): Weight decay for L2 regularization.
|
|
|
|
Returns:
|
|
list: List of dictionaries containing parameters and their weight decay settings.
|
|
"""
|
|
weights, biases = [], []
|
|
for name, param in self.named_parameters():
|
|
if 'bias' in name:
|
|
biases += [param]
|
|
else:
|
|
weights += [param]
|
|
params = [{
|
|
'params': weights,
|
|
'weight_decay': weight_decay,
|
|
}, {
|
|
'params': biases,
|
|
'weight_decay': 0.0,
|
|
}]
|
|
return params
|
|
|