File size: 5,864 Bytes
d8ff165 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | import torch
import torch.nn as nn
import torch.nn.functional as F
class ConvBlock(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size=(3, 3), stride=(2, 1), padding=(1, 1)):
super().__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
self.bn = nn.BatchNorm2d(out_channels)
self.act = nn.LeakyReLU(0.2)
def forward(self, x):
return self.act(self.bn(self.conv(x)))
class DeconvBlock(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size=(3, 3), stride=(2, 1), padding=(1, 1), output_padding=(1, 0)):
super().__init__()
self.deconv = nn.ConvTranspose2d(
in_channels, out_channels, kernel_size,
stride=stride, padding=padding, output_padding=output_padding
)
self.bn = nn.BatchNorm2d(out_channels)
self.act = nn.LeakyReLU(0.2)
def forward(self, x):
return self.act(self.bn(self.deconv(x)))
class ComplexCRNDenoiser(nn.Module):
"""
Complex Convolutional Recurrent Network (CRN) for Audio Noise Reduction.
Operates on complex STFT representations (Real + Imaginary) to preserve both
magnitude and phase, completely eliminating phase distortion.
"""
def __init__(
self,
n_fft=512,
hop_length=128,
win_length=512,
channels=[16, 32, 64, 128, 256],
rnn_hidden=256,
rnn_layers=2
):
super().__init__()
self.n_fft = n_fft
self.hop_length = hop_length
self.win_length = win_length
self.register_buffer("window", torch.hann_window(win_length))
# Encoder (Frequency downsampling by 2 at each step)
# Input shape: (B, 2, 256, T) -> channels[0] -> channels[1] -> ... -> channels[4]
# F: 256 -> 128 -> 64 -> 32 -> 16 -> 8
self.enc1 = ConvBlock(2, channels[0])
self.enc2 = ConvBlock(channels[0], channels[1])
self.enc3 = ConvBlock(channels[1], channels[2])
self.enc4 = ConvBlock(channels[2], channels[3])
self.enc5 = ConvBlock(channels[3], channels[4])
bottleneck_f = 8
self.rnn_in_dim = channels[4] * bottleneck_f
self.rnn = nn.GRU(
input_size=self.rnn_in_dim,
hidden_size=rnn_hidden,
num_layers=rnn_layers,
batch_first=True,
bidirectional=True
)
self.rnn_fc = nn.Linear(rnn_hidden * 2, self.rnn_in_dim)
# Decoder with Skip-Connections
self.dec5 = DeconvBlock(channels[4] * 2, channels[3])
self.dec4 = DeconvBlock(channels[3] * 2, channels[2])
self.dec3 = DeconvBlock(channels[2] * 2, channels[1])
self.dec2 = DeconvBlock(channels[1] * 2, channels[0])
self.dec1 = nn.ConvTranspose2d(
channels[0] * 2, 2, kernel_size=(3, 3), stride=(2, 1),
padding=(1, 1), output_padding=(1, 0)
)
def forward_stft(self, x):
# x is (B, 2, 256, T)
e1 = self.enc1(x) # (B, C0, 128, T)
e2 = self.enc2(e1) # (B, C1, 64, T)
e3 = self.enc3(e2) # (B, C2, 32, T)
e4 = self.enc4(e3) # (B, C3, 16, T)
e5 = self.enc5(e4) # (B, C4, 8, T)
B, C, F_dim, T_dim = e5.shape
# Reshape for GRU: (B, T, C * F)
rnn_in = e5.permute(0, 3, 1, 2).reshape(B, T_dim, C * F_dim)
rnn_out, _ = self.rnn(rnn_in)
rnn_proj = self.rnn_fc(rnn_out)
# Reshape back to (B, C, F, T)
d_in = rnn_proj.reshape(B, T_dim, C, F_dim).permute(0, 2, 3, 1)
# Skip connection decoding
d5 = self.dec5(torch.cat([d_in, e5], dim=1))
d4 = self.dec4(torch.cat([d5, e4], dim=1))
d3 = self.dec3(torch.cat([d4, e3], dim=1))
d2 = self.dec2(torch.cat([d3, e2], dim=1))
mask = self.dec1(torch.cat([d2, e1], dim=1)) # (B, 2, 256, T)
# Complex mask bounded via tanh
mask = torch.tanh(mask)
return mask
def forward(self, noisy_wav):
"""
Takes raw 1D audio waveform: (Batch, AudioSamples)
Returns cleaned 1D audio waveform: (Batch, AudioSamples)
"""
orig_len = noisy_wav.shape[-1]
if self.window.device != noisy_wav.device:
self.window = self.window.to(noisy_wav.device)
# 1. STFT
stft_res = torch.stft(
noisy_wav,
n_fft=self.n_fft,
hop_length=self.hop_length,
win_length=self.win_length,
window=self.window,
return_complex=True
) # (B, 257, T)
# Separate Real and Imaginary components
real = stft_res.real
imag = stft_res.imag
# Drop highest frequency bin (257 -> 256) for power-of-2 conv pooling
real_256 = real[:, :256, :]
imag_256 = imag[:, :256, :]
stft_in = torch.stack([real_256, imag_256], dim=1) # (B, 2, 256, T)
# 2. Predict Complex Mask (Mr, Mi)
mask = self.forward_stft(stft_in)
mask_r = mask[:, 0, :, :]
mask_i = mask[:, 1, :, :]
# 3. Complex multiplication: (Yr + j Yi) = (Xr + j Xi) * (Mr + j Mi)
# Yr = Xr * Mr - Xi * Mi
# Yi = Xr * Mi + Xi * Mr
clean_real_256 = real_256 * mask_r - imag_256 * mask_i
clean_imag_256 = real_256 * mask_i + imag_256 * mask_r
# Pad 257th frequency back with zeros
clean_real = F.pad(clean_real_256, (0, 0, 0, 1))
clean_imag = F.pad(clean_imag_256, (0, 0, 0, 1))
clean_stft = torch.complex(clean_real, clean_imag)
# 4. Inverse STFT to time domain
clean_wav = torch.istft(
clean_stft,
n_fft=self.n_fft,
hop_length=self.hop_length,
win_length=self.win_length,
window=self.window,
length=orig_len
)
return clean_wav
|