mcw519 / PureSound

Make the sound you hear pure and clean by deep learning.
7 stars 0 forks source link

about tse inference #2

Closed zuowanbushiwo closed 1 year ago

zuowanbushiwo commented 1 year ago

Hi I want to test your tse model, but I have no success running your tse/demo/demo_app.py on windows. So I carefully read the content under the tse folder and wrote a simple test function, I get different results for streaming and non-streaming inference, is it right? Is my test code correct? Thanks!

import torch
from utils import DemoSpeakerNet, DemoTseNet
import soundfile
from model import init_model
if __name__ == '__main__':
    ckpt = torch.load("onnx_tse/skim_causal_460_wNoise_IS_tsdr.ckpt", map_location="cpu")
    enroll_wav_file = 'onnx_tse/s1/61-70968-0003_5105-28240-0000.wav'
    mix_wav_file = 'onnx_tse/mix/61-70968-0003_5105-28240-0000.wav'   
    enroll_wav  = torch.from_numpy(soundfile.read(enroll_wav_file, dtype="float32")[0].reshape(1, -1))
    mix_wav  = torch.from_numpy(soundfile.read(mix_wav_file, dtype="float32")[0].reshape(1, -1)) 
    speaker_net = DemoSpeakerNet()
    speaker_net.load_state_dict(ckpt["state_dict"], strict=False)
    speaker_net.eval()
    tse_net = DemoTseNet()
    tse_net.load_state_dict(ckpt["state_dict"], strict=False)
    tse_net.eval()
    tse_net.masker.init_status()
    speaker_embedding = speaker_net.get_speaker_embedding(enroll_wav) 
    stream_enh_wav  = tse_net.streaming_inference_chunk(mix_wav, speaker_embedding)
    soundfile.write('onnx_tse/enh_stream.wav', stream_enh_wav, 16000)

    model = init_model('tse_skim_v0_causal', verbose=False)
    model.load_state_dict(ckpt["state_dict"], strict=False)  # ignore loss's weight
    model.eval()      
    enh_wav = model.inference(mix_wav, enroll_wav)
    enh_wav = enh_wav.detach().cpu().numpy().reshape(-1)
    soundfile.write('onnx_tse/enh.wav', enh_wav, 16000)

    print('Hello world')

data.zip

mcw519 commented 1 year ago

Hi, Sorry about that I don't have too much time to check this issue now. (long time ago work..) You can check this first https://github.com/mcw519/PureSound/blob/main/test/test_streaming.py#L62 thanks

zuowanbushiwo commented 1 year ago

Thanks!