def test_save_and_load_convtasnet(fb):
model1 = ConvTasNet(n_src=2,
n_repeats=2,
n_blocks=2,
bn_chan=16,
hid_chan=4,
skip_chan=8,
n_filters=32,
fb_name=fb)
test_input = torch.randn(1, 800)
model_conf = model1.serialize()
reconstructed_model = ConvTasNet.from_pretrained(model_conf)
assert_allclose(model1.separate(test_input),
reconstructed_model(test_input))
def main(conf):
model_path = os.path.join(conf["exp_dir"], "best_model.pth")
model = ConvTasNet.from_pretrained(model_path)
# Handle device placement
if conf["use_gpu"]:
model.cuda()
model_device = next(model.parameters()).device
torch.no_grad().__enter__()
alto, fs = sf.read(
'/home/pc2752/share//Darius/data/satb_dst/test_dcs/raw_audio/DCS_TPQuartetA/DCS_TPQuartetA_alto_1.wav'
)
bass, fs = sf.read(
'/home/pc2752/share//Darius/data/satb_dst/test_dcs/raw_audio/DCS_TPQuartetA/DCS_TPQuartetA_bass_1.wav'
)
soprano, fs = sf.read(
'/home/pc2752/share//Darius/data/satb_dst/test_dcs/raw_audio/DCS_TPQuartetA/DCS_TPQuartetA_soprano_1.wav'
)
tenor, fs = sf.read(
'/home/pc2752/share//Darius/data/satb_dst/test_dcs/raw_audio/DCS_TPQuartetA/DCS_TPQuartetA_tenor_1.wav'
)
mix, fs = sf.read(
'/home/pc2752/share//Darius/Wave-U-Net/test_set_mixes/dcs/DCS_TPQuartetA_mix.wav'
)
mix = torch.from_numpy(mix).type(torch.FloatTensor)
outputs = model.float()(mix)
sdr, sir, sar, perm = separation.bss_eval_sources(
np.array([soprano, alto, tenor, bass]),
outputs[:, :soprano.shape[0]].detach().numpy())
print("SDR: {}\nSIR: {}\nSAR: {}".format(sdr, sir, sar))
sf.write('./DCS_TPQuartetA_Soprano.wav', outputs[0], fs)
sf.write('./DCS_TPQuartetA_Alto.wav', outputs[1], fs)
sf.write('./DCS_TPQuartetA_Tenor.wav', outputs[2], fs)
sf.write('./DCS_TPQuartetA_Bass.wav', outputs[3], fs)
def convtasnet(conf):
sys.path.append('./asteroid')
from asteroid.models import ConvTasNet
from asteroid.utils import tensors_to_device
from asteroid.models import save_publishable
model_path = './models/convtasnet_usecase2.pth'
model = ConvTasNet.from_pretrained(model_path)
# Handle device placement
if conf["use_gpu"]:
model.cuda()
model_device = next(model.parameters()).device
torch.no_grad().__enter__()
mix, fs = sf.read(conf["input_path"])
mix = torch.from_numpy(mix).type(torch.FloatTensor)
outputs = model.float()(mix)
return outputs
def main(conf):
model_path = os.path.join(conf["exp_dir"], "best_model.pth")
model = ConvTasNet.from_pretrained(model_path)
# Handle device placement
if conf["use_gpu"]:
model.cuda()
model_device = next(model.parameters()).device
test_set = SourceFolderDataset(
os.path.join(conf["exp_dir"], "json/"),
conf["wav_dir"],
conf["n_src"],
conf["sample_rate"],
conf["batch_size"],
) # Uses all segment length
# Used to reorder sources only
loss_func = PITLossWrapper(pairwise_neg_sisdr, pit_from="pw_mtx")
# Randomly choose the indexes of sentences to save.
ex_save_dir = os.path.join(conf["exp_dir"], "examples/")
if conf["n_save_ex"] == -1:
conf["n_save_ex"] = len(test_set)
save_idx = random.sample(range(len(test_set)), conf["n_save_ex"])
series_list = []
torch.no_grad().__enter__()
for idx in tqdm(range(len(test_set))):
# Forward the network on the mixture.
mix, sources = tensors_to_device(test_set[idx], device=model_device)
mix = mix.unsqueeze(0)
sources = sources.unsqueeze(0)
est_sources = model(mix)
#print(test_set[idx])
#print(est_sources.shape, sources.shape, mix.shape, len(test_set))
loss, reordered_sources = loss_func(est_sources, sources, return_est=True)
#mix_np = mix.squeeze(0).cpu().data.numpy()
mix_np = mix.cpu().data.numpy()
sources_np = sources.squeeze(0).cpu().data.numpy()
est_sources_np = reordered_sources.squeeze(0).cpu().data.numpy()
utt_metrics = get_metrics(
mix_np,
sources_np,
est_sources_np,
sample_rate=conf["sample_rate"],
metrics_list=compute_metrics,
)
utt_metrics["mix_path"] = test_set.mix[idx][0]
series_list.append(pd.Series(utt_metrics))
# Save some examples in a folder. Wav files and metrics as text.
if idx in save_idx:
local_save_dir = os.path.join(ex_save_dir, "ex_{}/".format(idx))
os.makedirs(local_save_dir, exist_ok=True)
#print(mix_np.shape)
sf.write(local_save_dir + "mixture.wav", np.swapaxes(mix_np,0,1), conf["sample_rate"])
# Loop over the sources and estimates
for src_idx, src in enumerate(sources_np):
sf.write(local_save_dir + "s{}.wav".format(src_idx + 1), src, conf["sample_rate"])
for src_idx, est_src in enumerate(est_sources_np):
est_src *= np.max(np.abs(mix_np)) / np.max(np.abs(est_src))
sf.write(
local_save_dir + "s{}_estimate.wav".format(src_idx + 1),
est_src,
conf["sample_rate"],
)
# Write local metrics to the example folder.
with open(local_save_dir + "metrics.json", "w") as f:
json.dump(utt_metrics, f, indent=0)
# Save all metrics to the experiment folder.
all_metrics_df = pd.DataFrame(series_list)
all_metrics_df.to_csv(os.path.join(conf["exp_dir"], "all_metrics.csv"))
# Print and save summary metrics
final_results = {}
for metric_name in compute_metrics:
input_metric_name = "input_" + metric_name
ldf = all_metrics_df[metric_name] - all_metrics_df[input_metric_name]
final_results[metric_name] = all_metrics_df[metric_name].mean()
final_results[metric_name + "_imp"] = ldf.mean()
print("Overall metrics :")
pprint(final_results)
with open(os.path.join(conf["exp_dir"], "final_metrics.json"), "w") as f:
json.dump(final_results, f, indent=0)
model_dict = torch.load(model_path, map_location="cpu")
os.makedirs(os.path.join(conf["exp_dir"], "publish_dir"), exist_ok=True)
publishable = save_publishable(
os.path.join(conf["exp_dir"], "publish_dir"),
model_dict,
metrics=final_results,
train_conf=train_conf,
)
def main(conf):
model_path = os.path.join(conf['exp_dir'], 'best_model.pth')
model = ConvTasNet.from_pretrained(model_path)
# Handle device placement
if conf['use_gpu']:
model.cuda()
model_device = next(model.parameters()).device
test_set = WhamDataset(conf['test_dir'],
conf['task'],
sample_rate=conf['sample_rate'],
nondefault_nsrc=model.masker.n_src,
segment=None) # Uses all segment length
# Used to reorder sources only
loss_func = PITLossWrapper(pairwise_neg_sisdr, pit_from='pw_mtx')
# Randomly choose the indexes of sentences to save.
ex_save_dir = os.path.join(conf['exp_dir'], 'examples/')
if conf['n_save_ex'] == -1:
conf['n_save_ex'] = len(test_set)
save_idx = random.sample(range(len(test_set)), conf['n_save_ex'])
series_list = []
torch.no_grad().__enter__()
for idx in tqdm(range(len(test_set))):
# Forward the network on the mixture.
mix, sources = tensors_to_device(test_set[idx], device=model_device)
est_sources = model(mix[None, None])
loss, reordered_sources = loss_func(est_sources,
sources[None],
return_est=True)
mix_np = mix[None].cpu().data.numpy()
sources_np = sources.cpu().data.numpy()
est_sources_np = reordered_sources.squeeze(0).cpu().data.numpy()
utt_metrics = get_metrics(mix_np,
sources_np,
est_sources_np,
sample_rate=conf['sample_rate'],
metrics_list=compute_metrics)
utt_metrics['mix_path'] = test_set.mix[idx][0]
series_list.append(pd.Series(utt_metrics))
# Save some examples in a folder. Wav files and metrics as text.
if idx in save_idx:
local_save_dir = os.path.join(ex_save_dir, 'ex_{}/'.format(idx))
os.makedirs(local_save_dir, exist_ok=True)
sf.write(local_save_dir + "mixture.wav", mix_np[0],
conf['sample_rate'])
# Loop over the sources and estimates
for src_idx, src in enumerate(sources_np):
sf.write(local_save_dir + "s{}.wav".format(src_idx + 1), src,
conf['sample_rate'])
for src_idx, est_src in enumerate(est_sources_np):
est_src *= np.max(np.abs(mix_np)) / np.max(np.abs(est_src))
sf.write(
local_save_dir + "s{}_estimate.wav".format(src_idx + 1),
est_src, conf['sample_rate'])
# Write local metrics to the example folder.
with open(local_save_dir + 'metrics.json', 'w') as f:
json.dump(utt_metrics, f, indent=0)
# Save all metrics to the experiment folder.
all_metrics_df = pd.DataFrame(series_list)
all_metrics_df.to_csv(os.path.join(conf['exp_dir'], 'all_metrics.csv'))
# Print and save summary metrics
final_results = {}
for metric_name in compute_metrics:
input_metric_name = 'input_' + metric_name
ldf = all_metrics_df[metric_name] - all_metrics_df[input_metric_name]
final_results[metric_name] = all_metrics_df[metric_name].mean()
final_results[metric_name + '_imp'] = ldf.mean()
print('Overall metrics :')
pprint(final_results)
with open(os.path.join(conf['exp_dir'], 'final_metrics.json'), 'w') as f:
json.dump(final_results, f, indent=0)
model_dict = torch.load(model_path, map_location='cpu')
os.makedirs(os.path.join(conf['exp_dir'], 'publish_dir'), exist_ok=True)
publishable = save_publishable(os.path.join(conf['exp_dir'],
'publish_dir'),
model_dict,
metrics=final_results,
train_conf=train_conf)
def main(conf):
model_path = os.path.join(conf['exp_dir'], 'best_model.pth')
model = ConvTasNet.from_pretrained(model_path)
# Handle device placement
if conf['use_gpu']:
model.cuda()
model_device = next(model.parameters()).device
# get data for evaluation - this should change in the future to work on real test data the was not used for training
dataset = SeparationDataset(combination_list_path=os.path.join(
conf['exp_dir'], 'combination_list.pkl'))
n_val = int(
len(dataset) * conf['train_conf']['data']
['fraction_of_examples_to_use_for_validation'])
train_set, val_set = random_split(dataset,
[len(dataset) - n_val, n_val]) # noqa
# test_set = val_set
test_set = train_set
# Used to reorder sources only
loss_func = PITLossWrapper(pairwise_neg_sisdr, pit_from='pw_mtx')
# Randomly choose the indexes of sentences to save.
ex_save_dir = os.path.join(conf['exp_dir'], 'examples/')
if conf['n_save_ex'] == -1:
conf['n_save_ex'] = len(test_set)
save_idx = random.sample(range(len(test_set)), conf['n_save_ex'])
# series_list = []
torch.no_grad().__enter__()
for idx in tqdm(range(len(test_set))):
# Forward the network on the mixture.
mix, sources = tensors_to_device(test_set[idx], device=model_device)
est_sources = model(mix[None, None])
loss, reordered_sources = loss_func(est_sources,
sources[None],
return_est=True) # noqa
mix_np = to_complex(mix[None].cpu().data.numpy())
sources_np = to_complex(sources.cpu().data.numpy())
est_sources_np = to_complex(
reordered_sources.squeeze(0).cpu().data.numpy())
# utt_metrics = get_metrics(mix_np, sources_np, est_sources_np,
# sample_rate=conf['sample_rate'],
# metrics_list=compute_metrics)
# utt_metrics['mix_path'] = test_set.mix[idx][0]
# series_list.append(pd.Series(utt_metrics))
# Save some examples in a folder. Wav files and metrics as text.
if idx in save_idx:
local_save_dir = os.path.join(ex_save_dir, 'ex_{}/'.format(idx))
os.makedirs(local_save_dir, exist_ok=True)
iq_data = mix_np[0]
ax = plot_spectogram(iq_data, scale=False, show_plot=False)
ax.figure.savefig(local_save_dir + 'mixture.png')
# Loop over the sources and estimates
for src_idx, src in enumerate(sources_np):
iq_data = src
ax = plot_spectogram(iq_data, scale=False, show_plot=False)
ax.figure.savefig(local_save_dir +
"s{}.png".format(src_idx + 1))
for src_idx, est_src in enumerate(est_sources_np):
# est_src *= np.max(np.abs(mix_np))/np.max(np.abs(est_src))
iq_data = np.reshape(est_src, (32, 128)).T
ax = plot_spectogram(iq_data, scale=False, show_plot=False)
ax.figure.savefig(local_save_dir +
"s{}_estimate.png".format(src_idx + 1))