def main():
device = torch.device('cuda' if FLAGS.use_cuda else 'cpu')
model = build_model()
model.to(device)
if FLAGS.log_dir is None:
writer = SummaryWriter(FLAGS.model_dir + '/tensorboard')
else:
writer = SummaryWriter(FLAGS.log_dir)
# Training
if not FLAGS.decode:
train(model, device, writer)
# Evaluating
reload_for_eval(model, FLAGS.model_dir, FLAGS.use_cuda)
evaluate(model, device)
# SI-SDR
eval_si_sdr(FLAGS.test_wav_dir, FLAGS.model_dir)
# SDR.sources
eval_sdr_sources(FLAGS.test_wav_dir, FLAGS.model_dir)
# SDR.v4
eval_sdr(FLAGS.test_wav_dir, FLAGS.model_dir)
key_s2, feat_s2 = compute_spectrum(s2_wav[i])
assert key_mix == key_s1 and key_s1 == key_s2
mask_s1 = ibm(feat_s1, feat_s2, use_log, use_power)
mask_s2 = 1 - mask_s1
key_wav, wav = readwav(mix_wav[i])
enhance_s1 = apply_mask(feat_mix, mask_s1, use_log, use_power)
enhance_s2 = apply_mask(feat_mix, mask_s2, use_log, use_power)
# Reconstruction
wav_s1 = overlap_and_add(enhance_s1, wav, sample_rate, frame_length,
frame_shift, window_type, preemphasis,
use_log, use_power, square_root_window)
wav_s2 = overlap_and_add(enhance_s2, wav, sample_rate, frame_length,
frame_shift, window_type, preemphasis,
use_log, use_power, square_root_window)
wavwrite(wav_s1, sample_rate, recons_dir + "/" + key_wav + "_1.wav")
wavwrite(wav_s2, sample_rate, recons_dir + "/" + key_wav + "_2.wav")
f_mix_wav.close()
f_s1_wav.close()
f_s2_wav.close()
# SI-SDR
eval_si_sdr(ori_dir, os.path.dirname(recons_dir))
# SDR.sources
eval_sdr_sources(ori_dir, os.path.dirname(recons_dir))
# SDR.v4
eval_sdr(ori_dir, os.path.dirname(recons_dir))