File size: 3,302 Bytes
64ceedd cf73d23 64ceedd d8be50a 64ceedd cf73d23 64ceedd cf73d23 64ceedd cf73d23 d8be50a 64ceedd cf73d23 64ceedd cf73d23 64ceedd cf73d23 64ceedd cf73d23 64ceedd b6c45cb |
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 |
import os
import torch
import numpy as np
import torchaudio
import yaml
from . import asteroid_test
def get_conf():
conf_filterbank = {
'n_filters': 64,
'kernel_size': 16,
'stride': 8
}
conf_masknet = {
'in_chan': 64,
'n_src': 2,
'out_chan': 64,
'ff_hid': 256,
'ff_activation': "relu",
'norm_type': "gLN",
'chunk_size': 100,
'hop_size': 50,
'n_repeats': 2,
'mask_act': 'sigmoid',
'bidirectional': True,
'dropout': 0
}
return conf_filterbank, conf_masknet
def load_dpt_model():
print('Load Separation Model...')
now_path = os.path.split(os.path.realpath(__file__))[0]
conf_filterbank, conf_masknet = get_conf()
model_path = os.path.join(now_path, "trained_model/train_dptnet_aishell_partOverlap_B2_300epoch_quan-int8.p")
model = getattr(asteroid_test, "DPTNet")(**conf_filterbank, **conf_masknet)
model = torch.quantization.quantize_dynamic(model, {torch.nn.LSTM, torch.nn.Linear}, dtype=torch.qint8)
state_dict = torch.load(model_path, map_location="cpu")
model.load_state_dict(state_dict)
model.eval()
return model
def dpt_sep_process(wav_path, model=None, outfilename=None):
if model is None:
model = load_model()
x, sr = torchaudio.load(wav_path)
x = x.cpu()
with torch.no_grad():
est_sources = model(x) # shape: (1, 2, T)
# 確保 est_sources 是 (1, 2, T)，再拆分
est_sources = est_sources.squeeze(0) # shape: (2, T)
sep_1, sep_2 = est_sources # 拆成兩個 (T, ) 的 tensor
# 正規化
max_abs = x[0].abs().max().item()
sep_1 = sep_1 * max_abs / sep_1.abs().max().item()
sep_2 = sep_2 * max_abs / sep_2.abs().max().item()
# 增加 channel 維度，變為 (1, T)
sep_1 = sep_1.unsqueeze(0)
sep_2 = sep_2.unsqueeze(0)
if outfilename is not None:
torchaudio.save(outfilename.replace('.wav', '_sep1.wav'), sep_1, sr)
torchaudio.save(outfilename.replace('.wav', '_sep2.wav'), sep_2, sr)
torchaudio.save(outfilename.replace('.wav', '_mix.wav'), x, sr)
else:
torchaudio.save(wav_path.replace('.wav', '_sep1.wav'), sep_1, sr)
torchaudio.save(wav_path.replace('.wav', '_sep2.wav'), sep_2, sr)
# def dpt_sep_process(wav_path, model=None, outfilename=None):
# if model == None:
# model = load_model()
# x, sr = torchaudio.load(wav_path)
# x = x.cpu()
# with torch.no_grad():
# est_sources = model(x)
# est_sources_np = est_sources.squeeze(0)
# sep_1, sep_2 = est_sources_np
# sep_1 = sep_1 * x[0].abs().max().item() / sep_1.abs().max().item()
# sep_2 = sep_2 * x[0].abs().max().item() / sep_2.abs().max().item()
# if outfilename != None:
# torchaudio.save(outfilename.replace('.wav', '_sep1.wav'), sep_1, sr)
# torchaudio.save(outfilename.replace('.wav', '_sep2.wav'), sep_2, sr)
# torchaudio.save(outfilename.replace('.wav', '_mix.wav'), x, sr)
# else:
# torchaudio.save(wav_path.replace('.wav', '_sep1.wav'), sep_1, sr)
# torchaudio.save(wav_path.replace('.wav', '_sep2.wav'), sep_2, sr)
if __name__ == '__main__':
print("This module should be used via Flask or Gradio.") |