chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,35 @@
|
||||
import pathlib
|
||||
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
from .nets import CascadedNet
|
||||
|
||||
|
||||
class DotDict(dict):
|
||||
def __getattr__(*args):
|
||||
val = dict.get(*args)
|
||||
return DotDict(val) if type(val) is dict else val
|
||||
|
||||
__setattr__ = dict.__setitem__
|
||||
__delattr__ = dict.__delitem__
|
||||
|
||||
|
||||
def load_sep_model(model_path, device='cpu'):
|
||||
model_path = pathlib.Path(model_path)
|
||||
config_file = model_path.with_name('config.yaml')
|
||||
with open(config_file, "r") as config:
|
||||
args = yaml.safe_load(config)
|
||||
args = DotDict(args)
|
||||
model = CascadedNet(
|
||||
args.n_fft,
|
||||
args.hop_length,
|
||||
args.n_out,
|
||||
args.n_out_lstm,
|
||||
True,
|
||||
is_mono=args.is_mono
|
||||
)
|
||||
model.to(device)
|
||||
model.load_state_dict(torch.load(model_path, map_location='cpu'))
|
||||
model.eval()
|
||||
return model
|
||||
@@ -0,0 +1,166 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def crop_center(h1, h2):
|
||||
h1_shape = h1.size()
|
||||
h2_shape = h2.size()
|
||||
|
||||
if h1_shape[3] == h2_shape[3]:
|
||||
return h1
|
||||
elif h1_shape[3] < h2_shape[3]:
|
||||
raise ValueError('h1_shape[3] must be greater than h2_shape[3]')
|
||||
|
||||
# s_freq = (h2_shape[2] - h1_shape[2]) // 2
|
||||
# e_freq = s_freq + h1_shape[2]
|
||||
s_time = (h1_shape[3] - h2_shape[3]) // 2
|
||||
e_time = s_time + h2_shape[3]
|
||||
h1 = h1[:, :, :, s_time:e_time]
|
||||
|
||||
return h1
|
||||
|
||||
|
||||
class Conv2DBNActiv(nn.Module):
|
||||
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
|
||||
super(Conv2DBNActiv, self).__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
nin, nout,
|
||||
kernel_size=ksize,
|
||||
stride=stride,
|
||||
padding=pad,
|
||||
dilation=dilation,
|
||||
bias=False
|
||||
),
|
||||
nn.BatchNorm2d(nout),
|
||||
activ()
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.LeakyReLU):
|
||||
super(Encoder, self).__init__()
|
||||
self.conv1 = Conv2DBNActiv(nin, nout, ksize, stride, pad, activ=activ)
|
||||
self.conv2 = Conv2DBNActiv(nout, nout, ksize, 1, pad, activ=activ)
|
||||
|
||||
def forward(self, x):
|
||||
h = self.conv1(x)
|
||||
h = self.conv2(h)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.ReLU, dropout=False):
|
||||
super(Decoder, self).__init__()
|
||||
self.conv1 = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
|
||||
# self.conv2 = Conv2DBNActiv(nout, nout, ksize, 1, pad, activ=activ)
|
||||
self.dropout = nn.Dropout2d(0.1) if dropout else None
|
||||
|
||||
def forward(self, x, skip=None, fixed_length=True):
|
||||
if fixed_length:
|
||||
x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True)
|
||||
else:
|
||||
_, _, h, w = x.size()
|
||||
x = F.pad(x, (0, 1, 0, 1), mode='replicate')
|
||||
x = F.interpolate(x, size=(2*h+1,2*w+1), mode='bilinear', align_corners=True)
|
||||
x = x[:, :, :-1, :-1]
|
||||
|
||||
if skip is not None:
|
||||
skip = crop_center(skip, x)
|
||||
x = torch.cat([x, skip], dim=1)
|
||||
|
||||
h = self.conv1(x)
|
||||
# h = self.conv2(h)
|
||||
|
||||
if self.dropout is not None:
|
||||
h = self.dropout(h)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class Mean(nn.Module):
|
||||
def __init__(self, dim, keepdims=False):
|
||||
super(Mean, self).__init__()
|
||||
self.dim = dim
|
||||
self.keepdims = keepdims
|
||||
|
||||
def forward(self, x):
|
||||
return x.mean(self.dim, keepdims=self.keepdims)
|
||||
|
||||
|
||||
class ASPPModule(nn.Module):
|
||||
|
||||
def __init__(self, nin, nout, dilations=(4, 8, 12), activ=nn.ReLU, dropout=False):
|
||||
super(ASPPModule, self).__init__()
|
||||
self.conv1 = nn.Sequential(
|
||||
Mean(dim=-2, keepdims=True), # nn.AdaptiveAvgPool2d((1, None)),
|
||||
Conv2DBNActiv(nin, nout, 1, 1, 0, activ=activ)
|
||||
)
|
||||
self.conv2 = Conv2DBNActiv(
|
||||
nin, nout, 1, 1, 0, activ=activ
|
||||
)
|
||||
self.conv3 = Conv2DBNActiv(
|
||||
nin, nout, 3, 1, dilations[0], dilations[0], activ=activ
|
||||
)
|
||||
self.conv4 = Conv2DBNActiv(
|
||||
nin, nout, 3, 1, dilations[1], dilations[1], activ=activ
|
||||
)
|
||||
self.conv5 = Conv2DBNActiv(
|
||||
nin, nout, 3, 1, dilations[2], dilations[2], activ=activ
|
||||
)
|
||||
self.bottleneck = Conv2DBNActiv(
|
||||
nout * 5, nout, 1, 1, 0, activ=activ
|
||||
)
|
||||
self.dropout = nn.Dropout2d(0.1) if dropout else None
|
||||
|
||||
def forward(self, x):
|
||||
_, _, h, w = x.size()
|
||||
# feat1 = F.interpolate(self.conv1(x), size=(h, w), mode='bilinear', align_corners=True)
|
||||
feat1 = self.conv1(x).repeat(1, 1, h, 1)
|
||||
feat2 = self.conv2(x)
|
||||
feat3 = self.conv3(x)
|
||||
feat4 = self.conv4(x)
|
||||
feat5 = self.conv5(x)
|
||||
out = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1)
|
||||
out = self.bottleneck(out)
|
||||
|
||||
if self.dropout is not None:
|
||||
out = self.dropout(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class LSTMModule(nn.Module):
|
||||
|
||||
def __init__(self, nin_conv, nin_lstm, nout_lstm):
|
||||
super(LSTMModule, self).__init__()
|
||||
self.conv = Conv2DBNActiv(nin_conv, 1, 1, 1, 0)
|
||||
self.lstm = nn.LSTM(
|
||||
input_size=nin_lstm,
|
||||
hidden_size=nout_lstm // 2,
|
||||
bidirectional=True
|
||||
)
|
||||
self.dense = nn.Sequential(
|
||||
nn.Linear(nout_lstm, nin_lstm),
|
||||
nn.BatchNorm1d(nin_lstm),
|
||||
nn.ReLU()
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
N, _, nbins, nframes = x.size()
|
||||
h = self.conv(x)[:, 0] # N, nbins, nframes
|
||||
h = h.permute(2, 0, 1) # nframes, N, nbins
|
||||
h, _ = self.lstm(h)
|
||||
h = self.dense(h.reshape(-1, h.size()[-1])) # nframes * N, nbins
|
||||
h = h.reshape(nframes, N, 1, nbins)
|
||||
h = h.permute(1, 2, 3, 0)
|
||||
|
||||
return h
|
||||
@@ -0,0 +1,202 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from . import layers
|
||||
|
||||
|
||||
class BaseNet(nn.Module):
|
||||
|
||||
def __init__(self, nin, nout, nin_lstm, nout_lstm, dilations=((4, 2), (8, 4), (12, 6)), fixed_length=True):
|
||||
super(BaseNet, self).__init__()
|
||||
self.enc1 = layers.Conv2DBNActiv(nin, nout, 3, 1, 1)
|
||||
self.enc2 = layers.Encoder(nout, nout * 2, 3, 2, 1)
|
||||
self.enc3 = layers.Encoder(nout * 2, nout * 4, 3, 2, 1)
|
||||
self.enc4 = layers.Encoder(nout * 4, nout * 6, 3, 2, 1)
|
||||
self.enc5 = layers.Encoder(nout * 6, nout * 8, 3, 2, 1)
|
||||
|
||||
self.aspp = layers.ASPPModule(nout * 8, nout * 8, dilations, dropout=True)
|
||||
|
||||
self.dec4 = layers.Decoder(nout * (6 + 8), nout * 6, 3, 1, 1)
|
||||
self.dec3 = layers.Decoder(nout * (4 + 6), nout * 4, 3, 1, 1)
|
||||
self.dec2 = layers.Decoder(nout * (2 + 4), nout * 2, 3, 1, 1)
|
||||
self.lstm_dec2 = layers.LSTMModule(nout * 2, nin_lstm, nout_lstm)
|
||||
self.dec1 = layers.Decoder(nout * (1 + 2) + 1, nout * 1, 3, 1, 1)
|
||||
|
||||
self.fixed_length = fixed_length
|
||||
|
||||
def __call__(self, x):
|
||||
e1 = self.enc1(x)
|
||||
e2 = self.enc2(e1)
|
||||
e3 = self.enc3(e2)
|
||||
e4 = self.enc4(e3)
|
||||
e5 = self.enc5(e4)
|
||||
|
||||
h = self.aspp(e5)
|
||||
|
||||
h = self.dec4(h, e4, fixed_length=self.fixed_length)
|
||||
h = self.dec3(h, e3, fixed_length=self.fixed_length)
|
||||
h = self.dec2(h, e2, fixed_length=self.fixed_length)
|
||||
h = torch.cat([h, self.lstm_dec2(h)], dim=1)
|
||||
h = self.dec1(h, e1, fixed_length=self.fixed_length)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class CascadedNet(nn.Module):
|
||||
|
||||
def __init__(self, n_fft, hop_length, nout=32, nout_lstm=128, is_complex=False, is_mono=False, fixed_length=True):
|
||||
super(CascadedNet, self).__init__()
|
||||
self.n_fft = n_fft
|
||||
self.hop_length = hop_length
|
||||
self.seg_length = 32 * hop_length
|
||||
self.is_complex = is_complex
|
||||
self.is_mono = is_mono
|
||||
self.register_buffer("window", torch.hann_window(n_fft), persistent=False)
|
||||
self.max_bin = n_fft // 2
|
||||
self.output_bin = n_fft // 2 + 1
|
||||
self.nin_lstm = self.max_bin // 2
|
||||
self.offset = 64
|
||||
|
||||
nin = 4 if is_complex else 2
|
||||
if is_mono:
|
||||
nin = nin // 2
|
||||
|
||||
self.stg1_low_band_net = nn.Sequential(
|
||||
BaseNet(nin, nout // 2, self.nin_lstm // 2, nout_lstm, fixed_length=fixed_length),
|
||||
layers.Conv2DBNActiv(nout // 2, nout // 4, 1, 1, 0)
|
||||
)
|
||||
self.stg1_high_band_net = BaseNet(
|
||||
nin, nout // 4, self.nin_lstm // 2, nout_lstm // 2, fixed_length=fixed_length
|
||||
)
|
||||
|
||||
self.stg2_low_band_net = nn.Sequential(
|
||||
BaseNet(nout // 4 + nin, nout, self.nin_lstm // 2, nout_lstm, fixed_length=fixed_length),
|
||||
layers.Conv2DBNActiv(nout, nout // 2, 1, 1, 0)
|
||||
)
|
||||
self.stg2_high_band_net = BaseNet(
|
||||
nout // 4 + nin, nout // 2, self.nin_lstm // 2, nout_lstm // 2, fixed_length=fixed_length
|
||||
)
|
||||
|
||||
self.stg3_full_band_net = BaseNet(
|
||||
3 * nout // 4 + nin, nout, self.nin_lstm, nout_lstm, fixed_length=fixed_length
|
||||
)
|
||||
|
||||
self.out = nn.Conv2d(nout, nin, 1, bias=False)
|
||||
self.aux_out = nn.Conv2d(3 * nout // 4, nin, 1, bias=False)
|
||||
|
||||
def forward(self, x):
|
||||
if self.is_complex:
|
||||
x = torch.cat([x.real, x.imag], dim=1)
|
||||
|
||||
x = x[:, :, :self.max_bin]
|
||||
|
||||
bandw = x.size()[2] // 2
|
||||
l1_in = x[:, :, :bandw]
|
||||
h1_in = x[:, :, bandw:]
|
||||
l1 = self.stg1_low_band_net(l1_in)
|
||||
h1 = self.stg1_high_band_net(h1_in)
|
||||
aux1 = torch.cat([l1, h1], dim=2)
|
||||
|
||||
l2_in = torch.cat([l1_in, l1], dim=1)
|
||||
h2_in = torch.cat([h1_in, h1], dim=1)
|
||||
l2 = self.stg2_low_band_net(l2_in)
|
||||
h2 = self.stg2_high_band_net(h2_in)
|
||||
aux2 = torch.cat([l2, h2], dim=2)
|
||||
|
||||
f3_in = torch.cat([x, aux1, aux2], dim=1)
|
||||
f3 = self.stg3_full_band_net(f3_in)
|
||||
|
||||
if self.is_complex:
|
||||
mask = self.out(f3)
|
||||
if self.is_mono:
|
||||
mask = torch.complex(mask[:, :1], mask[:, 1:])
|
||||
else:
|
||||
mask = torch.complex(mask[:, :2], mask[:, 2:])
|
||||
mask = self.bounded_mask(mask)
|
||||
else:
|
||||
mask = torch.sigmoid(self.out(f3))
|
||||
|
||||
mask = F.pad(
|
||||
input=mask,
|
||||
pad=(0, 0, 0, self.output_bin - mask.size()[2]),
|
||||
mode='replicate'
|
||||
)
|
||||
|
||||
return mask
|
||||
|
||||
def bounded_mask(self, mask, eps=1e-8):
|
||||
mask_mag = torch.abs(mask)
|
||||
mask = torch.tanh(mask_mag) * mask / (mask_mag + eps)
|
||||
return mask
|
||||
|
||||
def predict_mask(self, x):
|
||||
mask = self.forward(x)
|
||||
|
||||
if self.offset > 0:
|
||||
mask = mask[:, :, :, self.offset:-self.offset]
|
||||
assert mask.size()[3] > 0
|
||||
|
||||
return mask
|
||||
|
||||
def predict(self, x):
|
||||
mask = self.forward(x)
|
||||
pred = x * mask
|
||||
|
||||
if self.offset > 0:
|
||||
pred = pred[:, :, :, self.offset:-self.offset]
|
||||
assert pred.size()[3] > 0
|
||||
|
||||
return pred
|
||||
|
||||
def audio2spec(self, x, use_pad=False):
|
||||
B, C, T = x.shape
|
||||
x = x.reshape(B * C, T)
|
||||
if use_pad:
|
||||
T1 = T + self.hop_length
|
||||
T_pad = self.seg_length * ((T1 - 1) // self.seg_length + 1) - T1
|
||||
nl_pad = T_pad // 2 // self.hop_length
|
||||
Tl_pad = nl_pad * self.hop_length
|
||||
x = F.pad(x, (Tl_pad, T_pad - Tl_pad))
|
||||
spec = torch.stft(
|
||||
x,
|
||||
n_fft=self.n_fft,
|
||||
hop_length=self.hop_length,
|
||||
return_complex=True,
|
||||
window=self.window,
|
||||
pad_mode='constant'
|
||||
)
|
||||
spec = spec.reshape(B, C, spec.shape[-2], spec.shape[-1])
|
||||
return spec
|
||||
|
||||
def spec2audio(self, x):
|
||||
B, C, N, T = x.shape
|
||||
x = x.reshape(-1, N, T)
|
||||
x = torch.istft(x, self.n_fft, self.hop_length, window=self.window)
|
||||
x = x.reshape(B, C, -1)
|
||||
return x
|
||||
|
||||
def predict_from_audio(self, x):
|
||||
B, C, T = x.shape
|
||||
x = x.reshape(B * C, T)
|
||||
T1 = T + self.hop_length
|
||||
T_pad = self.seg_length * ((T1 - 1) // self.seg_length + 1) - T1
|
||||
nl_pad = T_pad // 2 // self.hop_length
|
||||
Tl_pad = nl_pad * self.hop_length
|
||||
x = F.pad(x, (Tl_pad, T_pad - Tl_pad))
|
||||
spec = torch.stft(
|
||||
x,
|
||||
n_fft=self.n_fft,
|
||||
hop_length=self.hop_length,
|
||||
return_complex=True,
|
||||
window=self.window,
|
||||
pad_mode='constant'
|
||||
)
|
||||
spec = spec.reshape(B, C, spec.shape[-2], spec.shape[-1])
|
||||
mask = self.forward(spec)
|
||||
spec_pred = spec * mask
|
||||
spec_pred = spec_pred.reshape(B * C, spec.shape[-2], spec.shape[-1])
|
||||
x_pred = torch.istft(spec_pred, self.n_fft, self.hop_length, window=self.window)
|
||||
x_pred = x_pred[:, Tl_pad: Tl_pad + T]
|
||||
x_pred = x_pred.reshape(B, C, T)
|
||||
return x_pred
|
||||
Reference in New Issue
Block a user