# -*- coding : utf-8 -*-

from asteroid.models import BaseModel, ConvTasNet
import soundfile as sf

def main(wav):
    # 'from_pretrained' automatically uses the right model class (asteroid.models.DPRNNTasNet).
    #model = ConvTasNet.from_pretrained("JorisCos/ConvTasNet_Libri2Mix_sepclean_8k") # sample=8k
    model = ConvTasNet.from_pretrained("JorisCos/ConvTasNet_Libri2Mix_sepclean_16k") # sample=16k
    # You can pass a NumPy array:
    mixture, sample = sf.read(wav, dtype="float32", always_2d=True)
    # Soundfile returns the mixture as shape (time, channels), and Asteroid expect
    # s (batch, channels, time)
    mixture = mixture.transpose()
    mixture = mixture.reshape(1, mixture.shape[0], mixture.shape[1])
    out_wavs = model.separate(mixture)
    a = wav[:-4]
    sf.write('./{}_dec0.wav'.format(a), out_wavs[0][0], sample)
    sf.write('./{}_dec1.wav'.format(a), out_wavs[0][1], sample)

    # Or simply a file name:
    #model.separate(wav)

if __name__ == "__main__":
    # wav = "./BAC009S0144W0121_BAC009S0047W0135_snr0.wav"
    wav = "./CallRecord_83_2022_7_22_abacf371-63d5-4aad-995b-fb225091df33_16000.wav"
    main(wav)